Vulkan partially delegated, decomposed cross product gives incorrect outputs

[GregoryComer]

🐛 Describe the bug

When torch linalg cross products ops are present in a Vulkan delegated model, the cross product gets decomposed and then partially delegated. It gives outputs that do not match eager. The model gives corrects results with no delegation.

May be related to #12222.

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower, EdgeCompileConfig, to_edge
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
from typing import Callable, List, Optional, Tuple, Union

class LinalgCrossModel(torch.nn.Module):
    def __init__(
        self,
        dim: int = -1
    ):
        super().__init__()
        self.dim = dim
        
    def forward(self, x, y):
        return torch.linalg.cross(x, y, dim=self.dim)
        
model = LinalgCrossModel()
inputs = (
    torch.randn(3),
    torch.randn(3)
)
eager_outputs = model(*inputs)

ep = torch.export.export(model.eval(), inputs)
print(ep)
lowered = to_edge_transform_and_lower(
    ep,
    partitioner=[VulkanPartitioner()],
    compile_config=EdgeCompileConfig(_check_ir_validity=False)
).to_executorch()
print(lowered.exported_program())

et_model = _load_for_executorch_from_buffer(lowered.buffer)
et_outputs = et_model([*inputs])[0]

print(f"Inputs: {inputs}")
print(f"Eager: {eager_outputs}")
print(f"ET:    {et_outputs}")

Outputs:

Inputs: (tensor([ 6.0440e-01,  1.6187e-04, -1.7365e-01]), tensor([-0.8563, -0.7608,  1.6770]))
Eager: tensor([-0.1318, -0.8649, -0.4597])
ET:    tensor([ 0.8649,  0.0000, -0.4597])

Versions

Run on Meta internal master, Jul 3, fbcode/SwiftShader

cc @SS-JIA @manuelcandales @cbilgin