Vulkan embeddings give incorrect outputs

🐛 Describe the bug

The following repro gives incorrect outputs on Vulkan. Outputs match eager when not delegated.

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 Model(torch.nn.Module):
    def __init__(
        self,
    ):
        super().__init__()
        self.embedding = torch.nn.Embedding(
            num_embeddings=10,
            embedding_dim=5,
        )
        
    def forward(self, x):
        return self.embedding(x)
        
model = Model()
inputs = (
    torch.randint(0, 10, (1, 4), dtype=torch.long),
)
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([[0, 0, 5, 9]]),)
Eager: tensor([[[ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057],
         [ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057],
         [-0.3060,  0.7337,  2.3316, -1.8212,  0.7317],
         [-1.2270,  1.6362,  0.8507, -1.4619,  0.6319]]],
       grad_fn=<EmbeddingBackward0>)
ET:    tensor([[[ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057],
         [ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057],
         [ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057],
         [ 0.2909, -1.0046, -0.4795, -2.5402,  0.1057]]])

Versions

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

cc @SS-JIA @manuelcandales @cbilgin

Start by running the provided embedding reproduction and comparing eager outputs with the Executorch outputs. Inspect executorch.backends.vulkan.partitioner.vulkan_partitioner.VulkanPartitioner and the to_edge_transform_and_lower path to trace how the embedding is delegated. Done means the Vulkan-delegated embedding returns the same values as eager for the repro inputs.