pytorch · msluszniak · Aug 14, 2026 · Aug 14, 2026
@@ -1437,33 +1437,56 @@ def register_where():
 
 @update_features(exir_ops.edge.aten.index.Tensor)
 def register_index_tensor():
-    def check_index_tensor_node(node: torch.fx.Node) -> bool:
+    def _index_tensor_shapes(node: torch.fx.Node):
+        """(self_val, index_val) for the supported single-index form, else None."""
         self_arg = node.args[0]
         indices = node.args[1]
 
-        # Only support 1D self tensor
         if not isinstance(self_arg, torch.fx.Node):
-            return False
+            return None
         self_val = self_arg.meta.get("val", None)
         if self_val is None:
-            return False
-        if len(self_val.size()) != 1:
-            return False
+            return None
 
-        # Only support exactly one non-None index tensor
+        # Only support exactly one non-None index tensor, applied to dim 0.
         if not isinstance(indices, (list, tuple)):
-            return False
+            return None
         non_none = [idx for idx in indices if idx is not None]
-        if len(non_none) != 1:
+        if len(non_none) != 1 or indices[0] is None:
+            return None
+        index_arg = non_none[0]
+        if not isinstance(index_arg, torch.fx.Node):
+            return None
+        index_val = index_arg.meta.get("val", None)
+        if index_val is None:
+            return None
+
+        return self_val, index_val
+
+    def check_index_tensor_node(node: torch.fx.Node) -> bool:
+        shapes = _index_tensor_shapes(node)
+        if shapes is None:
             return False
+        _, index_val = shapes
+        # The gather is expressed as "one index position per output slice", so
+        # the index must be 1-D. `self` may be any rank: the buffer shader
+        # copies self's trailing dims through unchanged.
+        return len(index_val.size()) == 1
 
-        return True
+    def pick_index_tensor_storage(node: torch.fx.Node):
+        shapes = _index_tensor_shapes(node)
+        # Only the buffer shader handles a higher-rank `self`; the texture
+        # variant still assumes the 1-D form (it reads self[idx, 0, 0, 0]).
+        if shapes is not None and len(shapes[0].size()) > 1:
+            return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER
+        return utils.ANY_STORAGE, utils.ANY_STORAGE
 
     return OpFeatures(
         inputs_storage=utils.ANY_STORAGE,
         inputs_dtypes=utils.FP_INT_T,
         supports_resize=True,
         are_node_inputs_supported_fn=check_index_tensor_node,
+        pick_io_storage_fn=pick_index_tensor_storage,
     )
 
 

@@ -42,10 +42,29 @@ def __init__(
 def create_pattern_match_from_internal_match(
     internal_match: InternalMatch,
 ) -> PatternMatch:
+    # `nodes_map` maps every pattern node to its match in the target graph,
+    # INCLUDING the pattern's placeholders. A placeholder's match is the node
+    # that merely *feeds* the matched subgraph, which the fused op does not
+    # compute and the partitioner therefore has no reason to claim.
+    #
+    # This matters because `all_nodes` becomes VulkanPartitioner's
+    # `fusable_nodes`, and `_is_node_supported` returns True for any node in
+    # that set *before* it consults the op's `are_node_inputs_supported_fn`.
+    # Including the feeder nodes therefore silently waives their support
+    # checks: an `aten.index.Tensor` feeding the `hf_rope` pattern got
+    # delegated despite `check_index_tensor_node` rejecting non-1-D `self`,
+    # and Vulkan's 1-D-only gather then aborted in `virtual_resize` with
+    # "new sizes cannot modify the dimensionality of the tensor".
+    placeholder_matches = set(internal_match.placeholder_nodes)
+    matched_nodes = [
+        node
+        for node in internal_match.nodes_map.values()
+        if node not in placeholder_matches
+    ]
     return PatternMatch(
         internal_match.placeholder_nodes,
         internal_match.returning_nodes,
-        list(internal_match.nodes_map.values()),
+        matched_nodes,
     )
 
 

@@ -719,6 +719,22 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
 
     const size_t num_inputs = compute_graph->inputs().size();
     const size_t num_outputs = compute_graph->outputs().size();
+    // `args` carries the delegate call's inputs followed by its outputs. If the
+    // serialized graph disagrees about either count, the `args.size() -
+    // num_outputs` computed below underflows and every output access reads
+    // through a wild pointer, so report the mismatch rather than segfault on
+    // it. A mutated buffer serialized as a graph output is one way to reach
+    // this: nothing passes such a buffer to the delegate call, so it has no
+    // argument slot.
+    VK_CHECK_COND(
+        args.size() == num_inputs + num_outputs,
+        "Vulkan graph declares ",
+        num_inputs,
+        " inputs and ",
+        num_outputs,
+        " outputs, but the delegate call supplied ",
+        args.size(),
+        " arguments");
     bool should_propagate_resize = false;
 #ifdef ET_EVENT_TRACER_ENABLED
     runtime::EventTracer* event_tracer = context.event_tracer();

@@ -43,15 +43,24 @@ void main() {
   // Convert output buffer index to tensor index
   TensorIndex out_tidx = linear_idx_to_tensor_idx(outp, out_bufi);
 
-  // Read the index value at the same tensor position
-  const uint index_bufi = tensor_idx_to_linear_idx(index, out_tidx);
+  // aten.index.Tensor with a single index tensor gathers along dim 0. This
+  // index space is WHCN-ordered -- axis 0 is the LAST pytorch dim -- so
+  // pytorch dim 0 is the highest axis, and self's trailing dims map 1:1 onto
+  // out's. With a 1-D index, rank(out) == rank(self).
+  const uint gather_axis = ndim(outp) - 1;
+  const uint gather_pos = idx_at(out_tidx, gather_axis);
+
+  // The index tensor is 1-D, so only its axis 0 (W) is populated. Indexing it
+  // with out's full coordinate is equivalent only when out is itself 1-D.
+  TensorIndex index_tidx;
+  initialize(index_tidx);
+  index_tidx.data[0][0] = gather_pos;
+  const uint index_bufi = tensor_idx_to_linear_idx(index, index_tidx);
   const int idx = t_index[index_bufi];
 
-  // Construct a tensor index for the 1D self tensor.
-  // In WHCN ordering, a 1D tensor has its elements along dim 0 (width).
-  TensorIndex self_tidx;
-  self_tidx.data[0] = uvec4(uint(idx), 0, 0, 0);
-  self_tidx.data[1] = uvec4(0);
+  // self shares out's trailing coordinates; only the gathered axis differs.
+  TensorIndex self_tidx = out_tidx;
+  self_tidx.data[div_4(gather_axis)][mod_4(gather_axis)] = uint(idx);
   const uint self_bufi = tensor_idx_to_linear_idx(inp, self_tidx);
 
   t_out[out_bufi] = t_self[self_bufi];

@@ -20,9 +20,18 @@ void resize_index_tensor_node(
     const std::vector<ValueRef>& resize_args) {
   (void)resize_args;
   const ValueRef out = args.at(0).refs.at(0);
+  const ValueRef self = args.at(1).refs.at(0);
   const ValueRef index = args.at(1).refs.at(1);
 
+  // aten.index.Tensor with a single index tensor gathers along dim 0, so
+  //   out.sizes = index.sizes ++ self.sizes[1:]
+  // Using the index's sizes alone is only correct when self is 1-D; for any
+  // higher-rank self it also changes the tensor's RANK, which virtual_resize
+  // rejects outright ("new sizes cannot modify the dimensionality").
+  const std::vector<int64_t> self_sizes = graph->sizes_of(self);
   std::vector<int64_t> out_sizes = graph->sizes_of(index);
+  out_sizes.insert(out_sizes.end(), self_sizes.begin() + 1, self_sizes.end());
+
   graph->virtual_resize(out, out_sizes);
 }
 

@@ -762,6 +762,24 @@ def forward(self, x):
 
         self.lower_module_and_test_output(model, sample_inputs)
 
+    def test_vulkan_backend_index_tensor_higher_rank_self(self):
+        # `table[positions]` with a 2-D table, the shape RoPE uses to look up
+        # its precomputed frequencies. The gather only ever needed `indices` to
+        # be 1-D -- `self`'s trailing dims come through unchanged -- but the
+        # support check used to require a 1-D `self` too, and the buffer shader
+        # only read `self[idx, 0, 0, 0]`.
+        class Gather(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.rand(32, 8))
+
+            def forward(self, positions):
+                return self.table[positions] * 2.0
+
+        sample_inputs = (torch.arange(5, dtype=torch.int32),)
+
+        self.lower_module_and_test_output(Gather(), sample_inputs)
+
     @disable_test(
         "Currently this test is failing due to weird partitioning because the eq scalar"
         "operator is not supported yet. Re-enable when the operator is supported."