Vulkan: fix embedding_texture for non width-packed outputs

embedding_texture.glsl assumed the output texel always packs along the
embedding dim: it looked up one index per texel and then read 4 consecutive
weight elements. When the output is height packed, which the memory layout
pass picks once the weight has more than 16384 rows, each texel holds 4
different rows instead, so every group of 4 output rows got the first row's
data shifted by one element.

Branch on the output's packed dim and load each component from its own weight
row when the packed dim is not the embedding dim.

Fixes pytorch#22333.

(cherry picked from commit 2e5b071)

Author: msluszniak

backends/vulkan/runtime/graph/ops/glsl/embedding_texture.glsl

@@ -50,6 +50,12 @@ int load_embedding_idx(const TensorIndex4D out_tidx) {
   return in_texel[elem_pos.comp];
 }
 
+T load_weight_element(const int embedding_idx, const int dim_idx) {
+  return T(t_weight[embedding_idx * int(width(weight)) + dim_idx]);
+}
+
+// Loads 4 consecutive elements of the embedding dim from a single weight row.
+// Only valid when the output texel packs along the embedding (width) dim.
 VEC4_T load_weight_texel(const int embedding_idx, const int dim_idx) {
   int buf_i = embedding_idx * int(width(weight)) + dim_idx;
   VEC4_T weight_texel;
@@ -66,11 +72,34 @@ void main() {
     return;
   }
 
-  TensorIndex4D out_tidx =
+  const TensorIndex4D out_tidx =
       texture_pos_to_tensor4d_idx_simple(outp, out_pos, out_layout);
-  const int embedding_idx = load_embedding_idx(out_tidx);
 
-  const VEC4_T weight_texel = load_weight_texel(embedding_idx, out_tidx.data.x);
+  const int packed_dim = get_packed_dim(out_layout);
+
+  VEC4_T weight_texel = VEC4_T(0);
+  if (packed_dim == 0) {
+    // The texel packs 4 consecutive elements of the embedding dim, so every
+    // component reads from the same weight row.
+    const int embedding_idx = load_embedding_idx(out_tidx);
+    weight_texel = load_weight_texel(embedding_idx, out_tidx.data.x);
+  } else {
+    // The texel packs 4 consecutive elements of an index dim, so every
+    // component reads a different weight row at the same embedding dim.
+    const int packed_dim_limit = safe_idx(outp.sizes, packed_dim);
+    const int packed_dim_start = safe_idx(out_tidx.data, packed_dim);
+    const int dim_idx = out_tidx.data.x;
+
+    TensorIndex4D lane_tidx = out_tidx;
+    [[unroll]] for (int i = 0; i < 4; ++i) {
+      const int packed_idx = packed_dim_start + i;
+      if (packed_idx < packed_dim_limit) {
+        safe_set(lane_tidx.data, packed_dim, packed_idx);
+        weight_texel[i] =
+            load_weight_element(load_embedding_idx(lane_tidx), dim_idx);
+      }
+    }
+  }
 
   imageStore(t_out, out_pos, weight_texel);
 }

backends/vulkan/test/test_vulkan_delegate.py

@@ -2131,6 +2131,23 @@ def forward(self, x):
             (torch.tensor([[0, 1, 0], [4, 2, 0]]),),
         )
 
+    def test_vulkan_backend_embedding_large_vocab(self):
+        # Past 16384 entries the output of the embedding is laid out as a height
+        # packed texture, so each texel holds 4 different weight rows instead of
+        # 4 elements of the same row. See #22333.
+        class EmbeddingModule(torch.nn.Module):
+            def __init__(self, embedding):
+                super().__init__()
+                self.embedding = embedding
+
+            def forward(self, x):
+                return self.embedding(x)
+
+        self.lower_module_and_test_output(
+            EmbeddingModule(torch.nn.Embedding(16385, 4)),
+            (torch.tensor([[0, 1, 16384], [7, 16000, 12345]]),),
+        )
+
     def test_vulkan_backend_embedding_3d(self):
         class EmbeddingModule(torch.nn.Module):
             def __init__(self, embedding):