diff --git a/backends/vulkan/_passes/tag_memory_meta_pass.py b/backends/vulkan/_passes/tag_memory_meta_pass.py index f97053734f9..e74ffbd60b2 100644 --- a/backends/vulkan/_passes/tag_memory_meta_pass.py +++ b/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( diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eed287b3a08..84ebdc54c7a 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1448,26 +1448,28 @@ 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. @@ -1475,6 +1477,10 @@ def check_index_tensor_node(node: torch.fx.Node) -> bool: 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: diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index ddd8e8994b1..f5a58fb1f85 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -72,14 +72,39 @@ void index_tensor(ComputeGraph& graph, const std::vector& 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(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(dim); + VK_GET_OP_FN("aten.index_select.default") + (graph, {self, dim_ref, index, out}); } REGISTER_OPERATORS { diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index c6915d37684..fc0060ec228 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/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." diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..108ff6b38a8 100644 --- a/backends/vulkan/utils.py +++ b/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.