diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index a88df42ee7b..eed287b3a08 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -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, ) diff --git a/backends/vulkan/patterns/pattern_registry.py b/backends/vulkan/patterns/pattern_registry.py index 9a906cd8770..de5c6e77a1d 100644 --- a/backends/vulkan/patterns/pattern_registry.py +++ b/backends/vulkan/patterns/pattern_registry.py @@ -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, ) diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp index ff62d4f258c..190a2b0d400 100644 --- a/backends/vulkan/runtime/VulkanBackend.cpp +++ b/backends/vulkan/runtime/VulkanBackend.cpp @@ -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(); diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl index 3469bb22fcc..63d611b5111 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl @@ -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]; diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index ec799f074ee..afce5ef3812 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -20,9 +20,18 @@ void resize_index_tensor_node( const std::vector& 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 self_sizes = graph->sizes_of(self); std::vector 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); } diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index ff709618259..c6915d37684 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -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."