From 4bd9e509e63e5d809fa6f57474fe9d821dd3647e Mon Sep 17 00:00:00 2001 From: Iliyan Georgiev Date: Tue, 1 Apr 2025 15:49:31 +0100 Subject: [PATCH] Arm backend: Add support alias_copy operator - Add NodeVisitor factory for IDENTITY ops - Add alias_copy support - Add alias_copy tests - Move getitem to new factory Signed-off-by: Iliyan Georgiev Change-Id: Idaefb92d87c6ebd87fd514a0f0973b21e5b28800 --- .../tosa_supported_operators.py | 1 + backends/arm/operators/__init__.py | 2 +- backends/arm/operators/op_get_item.py | 35 -------- backends/arm/operators/ops_identity.py | 47 +++++++++++ .../arm/quantizer/quantization_annotator.py | 7 +- backends/arm/test/ops/test_alias_copy.py | 83 +++++++++++++++++++ 6 files changed, 138 insertions(+), 37 deletions(-) delete mode 100644 backends/arm/operators/op_get_item.py create mode 100644 backends/arm/operators/ops_identity.py create mode 100644 backends/arm/test/ops/test_alias_copy.py diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 3cb30cacc5d..952cfb17cf0 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -229,6 +229,7 @@ def is_node_supported( exir_ops.edge.aten.__lshift__.Scalar, torch.ops.aten.scalar_tensor.default, exir_ops.edge.aten.gelu.default, + exir_ops.edge.aten.alias_copy.default, ] return supported diff --git a/backends/arm/operators/__init__.py b/backends/arm/operators/__init__.py index 861ed9c7701..da050c5994e 100644 --- a/backends/arm/operators/__init__.py +++ b/backends/arm/operators/__init__.py @@ -22,7 +22,6 @@ op_erf, op_exp, op_ge, - op_get_item, op_gt, op_le, op_log, @@ -51,5 +50,6 @@ op_view, op_where, ops_binary, + ops_identity, ops_unary, ) diff --git a/backends/arm/operators/op_get_item.py b/backends/arm/operators/op_get_item.py deleted file mode 100644 index 0e1192b3bef..00000000000 --- a/backends/arm/operators/op_get_item.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright 2023-2025 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe -from typing import List - -import torch - -import tosa_tools.v0_80.serializer.tosa_serializer as ts # type: ignore -from executorch.backends.arm.operators.node_visitor import ( - NodeVisitor, - register_node_visitor, -) -from executorch.backends.arm.tosa_mapping import TosaArg - - -@register_node_visitor -class GetItemVisitor(NodeVisitor): - target = "getitem" - - def __init__(self, *args): - super().__init__(*args) - - def define_node( - self, - node: torch.fx.Node, - tosa_graph: ts.TosaSerializer, - inputs: List[TosaArg], - output: TosaArg, - ) -> None: - item_name = inputs[0].name - ## Simply add an identityOp - tosa_graph.addOperator(ts.TosaOp.Op().IDENTITY, [item_name], [output.name]) diff --git a/backends/arm/operators/ops_identity.py b/backends/arm/operators/ops_identity.py new file mode 100644 index 00000000000..0c6527cf336 --- /dev/null +++ b/backends/arm/operators/ops_identity.py @@ -0,0 +1,47 @@ +# Copyright 2025 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import List + +import torch +import torch.fx + +import tosa_tools.v0_80.serializer.tosa_serializer as ts + +from executorch.backends.arm.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.arm.tosa_mapping import TosaArg + + +def identity_operator_factory(identity_target: str): + """ + Creates and registers NodeVisitors for operators that map directly + to a TOSA IDENTITY op. + """ + + class IdentityOperatorVisitor(NodeVisitor): + target = identity_target + + def define_node( + self, + node: torch.fx.Node, + tosa_graph: ts.TosaSerializer, + inputs: List[TosaArg], + output: TosaArg, + ) -> None: + # Simply add an identityOp + tosa_graph.addOperator( + ts.TosaOp.Op().IDENTITY, [inputs[0].name], [output.name] + ) + + register_node_visitor(IdentityOperatorVisitor) + + +identity_operator_factory("getitem") +identity_operator_factory("aten.alias_copy.default") diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index 1575d59fd77..5398101fd9a 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -244,6 +244,11 @@ def _match_pattern( operator.getitem, ] +_one_to_one_shared_input_or_input_act_qspec = [ + torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten.alias_copy.default, +] + def get_quant_properties( # noqa: C901 node: Node, gm: torch.fx.GraphModule, quantization_config @@ -332,7 +337,7 @@ def any_or_hardtanh_min_zero(n: Node): _QuantProperty(2, shared_qspec), # type: ignore[arg-type] ] quant_properties.quant_output = _QuantProperty(0, shared_qspec) # type: ignore[arg-type] - elif node.target == torch.ops.aten.adaptive_avg_pool2d.default: + elif node.target in _one_to_one_shared_input_or_input_act_qspec: input_qspec = ( SharedQuantizationSpec(node.args[0]) # type: ignore[arg-type] if arm_quantizer_utils.is_output_annotated(node.args[0]) # type: ignore diff --git a/backends/arm/test/ops/test_alias_copy.py b/backends/arm/test/ops/test_alias_copy.py new file mode 100644 index 00000000000..66fa92bc445 --- /dev/null +++ b/backends/arm/test/ops/test_alias_copy.py @@ -0,0 +1,83 @@ +# Copyright 2025 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineBI, + EthosU85PipelineBI, + TosaPipelineBI, + TosaPipelineMI, +) + +input_t1 = Tuple[torch.Tensor] + + +class AliasCopy(torch.nn.Module): + """ + Tests proper handling of alias_copy when used directly. + + alias_copy can also appear from PyTorch/ExecuTorch optimizations + such as `x.transpose(0, 0)`. This is optimized to an alias_copy but + not before dq/q operators are added. + """ + + aten_op = "torch.ops.aten.alias_copy.default" + exir_op = "executorch_exir_dialects_edge__ops_aten_alias_copy_default" + + test_data: dict[input_t1] = { + "1d_ramp": (torch.arange(-16, 16, 0.2),), + "2d_ones": (torch.ones(5, 5),), + "3d_rand": (torch.rand(3, 5, 5),), + "4d_zeros": (torch.zeros(1, 10, 10, 10),), + } + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor): + return torch.alias_copy(x) + + +@common.parametrize("test_data", AliasCopy.test_data) +def test_alias_copy_tosa_MI(test_data: input_t1): + TosaPipelineMI[input_t1]( + AliasCopy(), + test_data, + AliasCopy.aten_op, + AliasCopy.exir_op, + ).run() + + +@common.parametrize("test_data", AliasCopy.test_data) +def test_alias_copy_tosa_BI(test_data: input_t1): + TosaPipelineBI[input_t1]( + AliasCopy(), + test_data, + AliasCopy.aten_op, + AliasCopy.exir_op, + ).run() + + +@common.parametrize("test_data", AliasCopy.test_data) +def test_alias_copy_u55_BI(test_data: input_t1): + EthosU55PipelineBI[input_t1]( + AliasCopy(), + test_data, + AliasCopy.aten_op, + AliasCopy.exir_op, + ).run() + + +@common.parametrize("test_data", AliasCopy.test_data) +def test_alias_copy_u85_BI(test_data: input_t1): + EthosU85PipelineBI[input_t1]( + AliasCopy(), + test_data, + AliasCopy.aten_op, + AliasCopy.exir_op, + ).run()