diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 9c5fee2e099..2a059b47471 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -18,6 +18,7 @@ from .decompose_acos import DecomposeAcos from .decompose_addmm import DecomposeAddmm from .decompose_any import DecomposeAny +from .decompose_as_strided import DecomposeAsStrided from .decompose_atan2 import DecomposeAtan2 from .decompose_binary_alpha import DecomposeBinaryAlpha from .decompose_cdist import DecomposeCDist @@ -87,6 +88,7 @@ DecomposeAcos, DecomposeAddmm, DecomposeAny, + DecomposeAsStrided, DecomposeAtan2, DecomposeBinaryAlpha, DecomposeCDist, diff --git a/backends/qualcomm/_passes/decompose_as_strided.py b/backends/qualcomm/_passes/decompose_as_strided.py new file mode 100644 index 00000000000..c850e01d05e --- /dev/null +++ b/backends/qualcomm/_passes/decompose_as_strided.py @@ -0,0 +1,226 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportPass, PassResult + +from .utils import copy_meta, merge_decomposed_graph + + +class AsStridedDecomposed(torch.nn.Module): + """Decompose as_strided into flatten + index_select + reshape. + + as_strided reinterprets memory layout using size/stride/storage_offset. + We precompute the linear indices and use index_select on the flattened input. + """ + + def __init__(self, output_size, stride, storage_offset): + super().__init__() + indices = self._compute_indices(output_size, stride, storage_offset) + self.register_buffer("indices", indices) + self.output_size = output_size + + def _compute_indices(self, size, stride, storage_offset): + """Compute flat linear indices for as_strided access pattern. + + indices[i0, i1, ..., in] = storage_offset + sum(ij * stride[j]) + Use int32 so the buffer is compatible with QNN (which does not support int64). + """ + indices = torch.zeros(size, dtype=torch.int32) + for dim in range(len(size)): + shape = [1] * len(size) + shape[dim] = size[dim] + indices = ( + indices + + torch.arange(size[dim], dtype=torch.int32).reshape(shape) + * stride[dim] + ) + indices = indices + storage_offset + return indices.flatten() + + def forward(self, x): + flat = x.reshape(-1) + gathered = torch.index_select(flat, 0, self.indices) + return gathered.reshape(self.output_size) + + +class DecomposeAsStrided(ExportPass): + """ + Decompose aten.as_strided into supported QNN primitives. + + as_strided creates a view of a tensor with specified size, stride, and + storage_offset. Since QNN has no native as_strided op, we decompose it: + + Fast path (pure view): + If strides are contiguous for the output size and storage_offset == 0, + this is equivalent to a reshape/view_copy. + + General path (via module export + merge_decomposed_graph): + 1. Flatten input to 1D + 2. index_select with precomputed linear indices (int32 constant) + 3. Reshape to target output size + + The general path materializes one int32 index per output element and is + intended for modest output sizes. + """ + + def __init__(self) -> None: + super().__init__() + self._targets = { + torch.ops.aten.as_strided.default, + torch.ops.aten.as_strided_copy.default, + exir_ops.edge.aten.as_strided_copy.default, + } + + @staticmethod + def _is_contiguous_strides(size, stride): + """Check if strides correspond to a contiguous layout for the given size.""" + if not size: + return True + expected_stride = 1 + for i in range(len(size) - 1, -1, -1): + if stride[i] != expected_stride: + return False + expected_stride *= size[i] + return True + + @staticmethod + def _max_index(size, stride, storage_offset): + """Compute maximum linear index that as_strided would access.""" + if not size or math.prod(size) == 0: + return -1 + return storage_offset + sum((s - 1) * st for s, st in zip(size, stride)) + + def _decompose_as_view(self, node, graph, input_node, output_size): + """Fast path: replace as_strided with a view/reshape.""" + input_val = input_node.meta["val"] + if list(input_val.shape) != output_size: + is_edge = isinstance(node.target, EdgeOpOverload) + view_op = ( + exir_ops.edge.aten.view_copy.default + if is_edge + else torch.ops.aten.view.default + ) + with graph.inserting_before(node): + view_node = graph.create_node( + "call_function", + view_op, + (input_node, output_size), + ) + view_node.meta = copy_meta(node.meta) + for user in node.users.copy(): + user.replace_input_with(node, view_node) + else: + for user in node.users.copy(): + user.replace_input_with(node, input_node) + + def _decompose_with_gather( # noqa: C901 + self, node, graph, graph_module, input_node, output_size, stride, storage_offset + ): + """General path: decompose via flatten + index_select + reshape.""" + model = AsStridedDecomposed(output_size, stride, storage_offset) + + input_fake = input_node.meta["val"] + decomposed_module = torch.export.export( + model, + (input_fake,), + strict=True, + ).module() + + # Register the indices buffer with a unique name + buffer_name = f"_as_strided_indices_{node.name}" + graph_module.register_buffer(buffer_name, model.indices) + + # Rename get_attr nodes in decomposed graph to use the unique buffer name + for d_node in decomposed_module.graph.nodes: + if d_node.op == "get_attr" and d_node.target == "indices": + d_node.target = buffer_name + + # Set attribute on decomposed module so node_copy works + setattr(decomposed_module, buffer_name, model.indices) + + with graph.inserting_before(node): + remap = {"x": input_node} + merge_decomposed_graph( + remap=remap, + target_node=node, + target_graph=graph, + decomposed_graph_module=decomposed_module, + ) + graph.erase_node(node) + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + + for node in list(graph.nodes): + if node.op != "call_function" or node.target not in self._targets: + continue + + input_node = node.args[0] + size_arg = node.args[1] + stride_arg = node.args[2] + storage_offset = node.args[3] if len(node.args) > 3 else 0 + if storage_offset is None: + storage_offset = 0 + + # Guard: skip SymInt (dynamic shapes) + if not all(isinstance(s, int) for s in size_arg): + continue + if not all(isinstance(s, int) for s in stride_arg): + continue + + output_size = list(size_arg) + stride = list(stride_arg) + + # Guard: skip empty output + if math.prod(output_size) == 0: + continue + + input_val = input_node.meta["val"] + + # Both decomposition paths interpret flattened contiguous storage. + # Leave non-contiguous inputs unchanged rather than alter their + # storage semantics. + if not input_val.is_contiguous(): + continue + + # Guard: validate index bounds + max_idx = self._max_index(output_size, stride, storage_offset) + assert max_idx < input_val.numel(), ( + f"DecomposeAsStrided: indices out of bounds. " + f"max_index={max_idx} >= input_numel={input_val.numel()} " + f"(size={output_size}, stride={stride}, offset={storage_offset})" + ) + + if ( + storage_offset == 0 + and self._is_contiguous_strides(output_size, stride) + and math.prod(output_size) == input_val.numel() + ): + self._decompose_as_view(node, graph, input_node, output_size) + else: + self._decompose_with_gather( + node, + graph, + graph_module, + input_node, + output_size, + stride, + storage_offset, + ) + modified = True + + if not modified: + return PassResult(graph_module, False) + + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index cd42c024147..055567c2be5 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -23,6 +23,7 @@ DecomposeAcos, DecomposeAddmm, DecomposeAny, + DecomposeAsStrided, DecomposeAtan2, DecomposeBinaryAlpha, DecomposeCDist, @@ -170,6 +171,7 @@ def get_annotation_passes(cls): ReplaceArangeArgs, DecomposeAcos, DecomposeAddmm, + DecomposeAsStrided, DecomposeAtan2, DecomposeBinaryAlpha, DecomposeCDist, @@ -204,6 +206,7 @@ def get_annotation_passes(cls): def get_export_passes(cls): """Return export pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ + DecomposeAsStrided, DecomposeBinaryAlpha, DecomposeCDist, DecomposePDist, diff --git a/backends/qualcomm/builders/README.md b/backends/qualcomm/builders/README.md index a90ea4be7a5..308f60f1b8e 100644 --- a/backends/qualcomm/builders/README.md +++ b/backends/qualcomm/builders/README.md @@ -503,6 +503,7 @@ The following PyTorch operators are supported through decomposition or annotatio | `aten.addmm` | `DecomposeAddmm` | | `aten.adaptive_avg_pool1d`, `aten.avg_pool1d` | `AnnotateAvgPool1D` | | `aten.any` | `DecomposeAny` | +| `aten.as_strided` | `DecomposeAsStrided` | | `aten.asinh` | `DecomposeHyperbolicVariants` | | `aten.atan2.default`, `aten.atan2.out` | `DecomposeAtan2` | | `aten.atanh` | `DecomposeHyperbolicVariants` | diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 5930869ecd4..cc4f08c8e5d 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -268,6 +268,17 @@ def forward(self, x, y): return squeeze_out, conv_out +class AsStrided(torch.nn.Module): + def __init__(self, size, stride, storage_offset=0): + super().__init__() + self.size = size + self.stride = stride + self.storage_offset = storage_offset + + def forward(self, x): + return torch.as_strided(x, self.size, self.stride, self.storage_offset) + + class Asinh(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/rework/htp/op/v68/test.py b/backends/qualcomm/tests/rework/htp/op/v68/test.py index 603d7d7646f..26067447ad0 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -193,6 +193,12 @@ def test_argmin(request, kwargs): ArgMin.test(request, kwargs) # noqa: F405 +@enumerate_activation_dtype([Tolerance(), Tolerance(), pytest.raises(AssertionError)]) +@with_htp_context +def test_as_strided(request, kwargs): + AsStrided.test(request, kwargs) # noqa: F405 + + @enumerate_activation_dtype( [ Tolerance(), diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index d2ac2d87fd2..2df33666f7a 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -460,6 +460,72 @@ def test_dtype_float(qnn_config, quantizer, compile_spec, expected): ) +class AsStrided(torch.nn.Module): + def __init__(self, size, stride, storage_offset): + super().__init__() + self.size = size + self.stride = stride + self.storage_offset = storage_offset + + def forward(self, x): + return torch.as_strided(x, self.size, self.stride, self.storage_offset) + + @staticmethod + @unpack_fixtures + def test(subtests, qnn_config, quantizer, compile_spec, expected): + cases = [ + # Contiguous sub-block from 4x4 + { + "size": [2, 2], + "stride": [4, 1], + "storage_offset": 0, + "input": torch.randn(4, 4), + }, + # Non-contiguous stride + non-zero offset + { + "size": [2, 3], + "stride": [6, 2], + "storage_offset": 1, + "input": torch.randn(4, 4), + }, + # 1D strided slice + {"size": [4], "stride": [2], "storage_offset": 0, "input": torch.randn(8)}, + # 3D output + { + "size": [2, 2, 2], + "stride": [8, 4, 1], + "storage_offset": 0, + "input": torch.randn(16), + }, + # Transpose-like (column-major access) + { + "size": [3, 4], + "stride": [1, 3], + "storage_offset": 0, + "input": torch.randn(3, 4), + }, + ] + for case in cases: + size, stride, offset, inp = ( + case["size"], + case["stride"], + case["storage_offset"], + case["input"], + ) + with subtests.test(msg=f"size:{size}, stride:{stride}, offset:{offset}"): + with expected as metrics: + export_and_verify( + module=__class__( + size=size, stride=stride, storage_offset=offset + ), + inputs=(inp,), + qnn_config=qnn_config, + quantizer=quantizer, + compile_specs=compile_spec, + metrics=metrics, + ) + + class ArgMax(torch.nn.Module): def __init__(self, dim, keepdim): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 1ebda343c8c..133867eaab5 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -331,6 +331,55 @@ def test_qnn_backend_argmin(self): case[QCOM_MODULE], case[QCOM_SAMPLE_INPUTS] ) + def test_qnn_backend_as_strided(self): + test_comb = [ + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 2], stride=[4, 1], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(4, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 3], stride=[6, 2], storage_offset=1 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(4, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[3, 4], stride=[1, 3], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(3, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided(size=[4], stride=[2], storage_offset=0), # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(8),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 2, 2], stride=[8, 4, 1], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(16),)], + }, + ] + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + index += 1 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_asinh(self): module = Asinh() # noqa: F405 sample_input = (torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0, 3.0]).reshape(2, 3),) @@ -3402,6 +3451,56 @@ def test_qnn_backend_argmin(self): ) self.lower_module_and_test_output(module, case[QCOM_SAMPLE_INPUTS]) + def test_qnn_backend_as_strided(self): + test_comb = [ + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 2], stride=[4, 1], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(4, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 3], stride=[6, 2], storage_offset=1 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(4, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[3, 4], stride=[1, 3], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(3, 4),)], + }, + { + QCOM_MODULE: [ + AsStrided(size=[4], stride=[2], storage_offset=0), # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(8),)], + }, + { + QCOM_MODULE: [ + AsStrided( # noqa: F405 + size=[2, 2, 2], stride=[8, 4, 1], storage_offset=0 + ), + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(16),)], + }, + ] + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + index += 1 + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_asin(self): module = Asin() # noqa: F405 sample_input = (torch.rand([3, 4]) * 2 - 1,)