Vulkan: support index.Tensor with the index on any dimension

backends/vulkan/_passes/tag_memory_meta_pass.py

@@ -6,7 +6,6 @@
 
 import logging
 import operator
-
 from collections import deque
 from typing import Any
 
@@ -121,8 +120,7 @@ def single_node_impl(node: torch.fx.Node) -> bool:
         return single_node_impl(arg_node)
     elif isinstance(arg_node, (list, tuple)):
         ret: bool = False
-        for n in arg_node:
-            assert isinstance(n, torch.fx.Node)
+        for n in utils.tensor_nodes_in_arg(arg_node):
             assert utils.is_single_tensor_node(n)
             ret = single_node_impl(n) or ret
 
@@ -193,9 +191,12 @@ def is_non_constant_tensor_node(self, node: Any) -> bool:
             return True
 
         if isinstance(node, (tuple, list)):
-            for n in node:
-                if not isinstance(n, torch.fx.Node):
-                    return False
+            # A list argument may carry a None for every dimension the operator
+            # does not touch, as index.Tensor does; only the tensors matter.
+            entries = utils.tensor_nodes_in_arg(node)
+            if len(entries) == 0:
+                return False
+            for n in entries:
                 if not self.is_non_constant_tensor_node(n):
                     return False
 
@@ -259,7 +260,7 @@ def get_arg_tensor_source_repset(
 
         # Special case for cat - use the first tensor in the list as representative
         if isinstance(arg_node, list):
-            arg_node = arg_node[0]
+            arg_node = utils.tensor_nodes_in_arg(arg_node)[0]
 
         if utils.has_node_repr(arg_node):
             arg_node_repr = utils.get_node_repr(arg_node)
@@ -422,7 +423,7 @@ def constrain_op_arg_repset(self, arg_i: int, op_repsets: utils.OpRepSets) -> No
         # First, trace downstream users to discover what layout they prefer.
         arg_node = op_repsets.op_node.args[arg_i]
         if isinstance(arg_node, list):
-            arg_node = arg_node[0]
+            arg_node = utils.tensor_nodes_in_arg(arg_node)[0]
 
         arg_repset = op_repsets.get_arg_repset(arg_i)
         if not arg_repset.is_constrained():
@@ -526,8 +527,7 @@ def set_op_node_tensor_reprs(
                     or transitions_inserted
                 )
             elif isinstance(arg_node, (list, tuple)):
-                for n in arg_node:
-                    assert isinstance(n, torch.fx.Node)
+                for n in utils.tensor_nodes_in_arg(arg_node):
                     assert utils.is_single_tensor_node(n)
                     transitions_inserted = (
                         set_arg_node_repr_or_transition(

backends/vulkan/op_registry.py

@@ -1448,33 +1448,39 @@ def _index_tensor_shapes(node: torch.fx.Node):
         if self_val is None:
             return None
 
-        # Only support exactly one non-None index tensor, applied to dim 0.
+        # Exactly one index tensor; the position it sits at is the dimension
+        # the gather runs along.
         if not isinstance(indices, (list, tuple)):
             return None
-        non_none = [idx for idx in indices if idx is not None]
-        if len(non_none) != 1 or indices[0] is None:
+        positions = [i for i, idx in enumerate(indices) if idx is not None]
+        if len(positions) != 1:
             return None
-        index_arg = non_none[0]
+        dim = positions[0]
+        index_arg = indices[dim]
         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
+        return self_val, index_val, dim
 
     def check_index_tensor_node(node: torch.fx.Node) -> bool:
         shapes = _index_tensor_shapes(node)
         if shapes is None:
             return False
-        _, index_val = shapes
+        _, 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
 
     def pick_index_tensor_storage(node: torch.fx.Node):
         shapes = _index_tensor_shapes(node)
+        # A gather on any dimension but 0 is handed to index_select, which reads
+        # and writes channels-packed textures.
+        if shapes is not None and shapes[2] != 0:
+            return utils.CHANNELS_PACKED_TEXTURE, utils.CHANNELS_PACKED_TEXTURE
         # 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:

backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp

@@ -72,14 +72,39 @@ void index_tensor(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   ValueRef indices_list_ref = args[1];
   ValueRef out = args[2];
 
-  ValueListPtr indices_list = graph.get_value_list(indices_list_ref);
-  VK_CHECK_COND(
-      indices_list->size() == 1,
-      "index.Tensor: only one index tensor is supported");
-
-  ValueRef index = indices_list->at(0);
-
-  add_index_tensor_node(graph, self, index, out);
+  // The indices list carries a null for every leading dimension that is not
+  // indexed, so the position of the tensor entry is the dimension the gather
+  // runs along. Exactly one entry may be a tensor.
+  int64_t dim = -1;
+  ValueRef index = kDummyValueRef;
+  {
+    // Scoped: this pointer has to be released before add_scalar() below, which
+    // may reallocate the graph's value store.
+    ValueListPtr indices_list = graph.get_value_list(indices_list_ref);
+    for (size_t i = 0; i < indices_list->size(); ++i) {
+      const ValueRef entry = indices_list->at(i);
+      if (graph.val_is_none(entry)) {
+        continue;
+      }
+      VK_CHECK_COND(
+          dim == -1, "index.Tensor: only one index tensor is supported");
+      dim = static_cast<int64_t>(i);
+      index = entry;
+    }
+  }
+  VK_CHECK_COND(dim != -1, "index.Tensor: no index tensor found");
+
+  if (dim == 0) {
+    add_index_tensor_node(graph, self, index, out);
+    return;
+  }
+
+  // Gathering along any other dimension is exactly index_select, which already
+  // handles every dimension. Reusing it leaves the index_tensor shader, which
+  // gathers along the leading dim, untouched.
+  const ValueRef dim_ref = graph.add_scalar<int64_t>(dim);
+  VK_GET_OP_FN("aten.index_select.default")
+  (graph, {self, dim_ref, index, out});
 }
 
 REGISTER_OPERATORS {

backends/vulkan/test/test_vulkan_delegate.py

@@ -780,6 +780,23 @@ def forward(self, positions):
 
         self.lower_module_and_test_output(Gather(), sample_inputs)
 
+    def test_vulkan_backend_index_tensor_non_leading_dim(self):
+        # `x[:, :, positions]`: the index sits on the last dim, which is how a
+        # replicate-style gather over a sequence is written. The support check
+        # used to require the index to be on dim 0, so the whole op fell back to
+        # the CPU and cut the graph in two everywhere it appeared.
+        class Gather(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("positions", torch.tensor([0, 0, 1, 3, 3]))
+
+            def forward(self, x):
+                return x[:, :, self.positions] * 2.0
+
+        sample_inputs = (torch.rand(size=(1, 4, 4), dtype=torch.float32),)
+
+        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."

backends/vulkan/utils.py

@@ -264,15 +264,29 @@ def is_tensor_node(node: Any) -> bool:
     return False
 
 
-def is_tensor_arg_node(node: Any) -> bool:
+def tensor_nodes_in_arg(node: Any) -> List[torch.fx.Node]:
+    """The tensor nodes that an operator argument refers to.
+
+    An argument is either a single tensor node or a list of them. A few
+    operators pass a list with a None for every dimension that is not involved:
+    `index.Tensor(x, [None, None, idx])` gathers along dim 2. Those Nones carry
+    no tensor and are skipped, so the tensors that ARE there still get a
+    representation assigned.
+    """
     if isinstance(node, torch.fx.Node):
-        return is_tensor_node(node)
-    elif isinstance(node, (list, tuple)):
-        if len(node) == 0:
-            return False
-        return all(is_tensor_node(n) for n in node)
+        return [node] if is_tensor_node(node) else []
+    if isinstance(node, (list, tuple)):
+        entries = [n for n in node if n is not None]
+        if len(entries) == 0:
+            return []
+        if all(is_tensor_node(n) for n in entries):
+            return entries
 
-    return False
+    return []
+
+
+def is_tensor_arg_node(node: Any) -> bool:
+    return len(tensor_nodes_in_arg(node)) > 0
 
 
 def num_tensor_arg_nodes(node: torch.fx.Node) -> int:
@@ -1512,10 +1526,10 @@ def filter_invalid_reprs_for_arg(
                 arg_node.meta["val"], arg_repsets, texture_limits
             )
         elif isinstance(arg_node, list) and all(
-            is_single_tensor_node(n) for n in arg_node
+            is_single_tensor_node(n) for n in tensor_nodes_in_arg(arg_node)
         ):
             return filter_invalid_reprs_for_node_list(
-                arg_repsets, arg_node, texture_limits
+                arg_repsets, tensor_nodes_in_arg(arg_node), texture_limits
             )
         # Special case for getitem; return the repset of the particular val in the
         # list of tensors that is being extracted.