From 80f429f0759e180e90fee4d45b50cf34cbdcfe5e Mon Sep 17 00:00:00 2001 From: Adrian Lundell Date: Fri, 23 Oct 2026 14:20:37 +0200 Subject: [PATCH] Arm backend: Add CONST and use tosa dialect op - Add a CONST dialect operation to TOSA dialect + helper function - Add explicit CONST shift to tosa.MUL - Add explicit CONST zero shift to tosa.NEGATE - Add explicit CONST pad values to tosa.PAD This will be used for 1-1 tosa graph serialization from exir graph, for now ignore dialect CONST nodes in the existing serializer flow Signed-off-by: Adrian Lundell Change-Id: I6e5ded885d02f8ca6a700c89786b8a68234e13cd --- .../_passes/aten_to_tosa_tensor_operators.py | 36 +++++- backends/arm/_passes/exir_to_tosa_pass.py | 14 ++- .../arm/_passes/insert_dynamic_padding.py | 14 ++- backends/arm/_passes/rewrite_pad.py | 11 +- backends/arm/_passes/tosa_lowering_utils.py | 29 +++++ backends/arm/operators/op_tosa_mul.py | 13 +- backends/arm/operators/op_tosa_neg.py | 2 +- backends/arm/operators/op_tosa_pad.py | 23 ++-- backends/arm/process_node.py | 6 + .../test_tosa_dialect_constant_visitors.py | 108 ++++++++++++++++ .../misc/tosa_dialect/test_data_layout_ops.py | 10 +- .../test_tosa_dialect_binary_ops.py | 33 ++++- .../test_tosa_dialect_unary_ops.py | 21 ++-- backends/arm/test/ops/test_to_copy.py | 1 - .../arm/test/passes/test_exir_to_tosa_pass.py | 116 ++++++++++++++++++ .../passes/test_remove_data_layout_noops.py | 7 +- backends/arm/tosa/dialect/__init__.py | 1 + .../tosa/dialect/ops/binary_elementwise.py | 22 ++-- backends/arm/tosa/dialect/ops/constants.py | 23 ++++ .../arm/tosa/dialect/ops/data_layout_ops.py | 10 +- .../arm/tosa/dialect/ops/unary_elementwise.py | 12 +- 21 files changed, 457 insertions(+), 55 deletions(-) create mode 100644 backends/arm/_passes/tosa_lowering_utils.py create mode 100644 backends/arm/test/misc/test_tosa_dialect_constant_visitors.py create mode 100644 backends/arm/test/passes/test_exir_to_tosa_pass.py create mode 100644 backends/arm/tosa/dialect/ops/constants.py diff --git a/backends/arm/_passes/aten_to_tosa_tensor_operators.py b/backends/arm/_passes/aten_to_tosa_tensor_operators.py index 54f2743d397..cde64cd7691 100644 --- a/backends/arm/_passes/aten_to_tosa_tensor_operators.py +++ b/backends/arm/_passes/aten_to_tosa_tensor_operators.py @@ -3,6 +3,15 @@ # 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 cast + +import torch + +from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, + get_output_qparams, +) +from executorch.backends.arm._passes.tosa_lowering_utils import create_const from executorch.backends.transforms.aten_to_dialect_pass import ( AtenToDialectPass, DialectNodeSpec, @@ -25,6 +34,29 @@ def rewrite_rfft2(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec | Non ) +def rewrite_mul(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + with node.graph.inserting_before(node): + shift = create_const(node.graph, 0, torch.int8, node) + return DialectNodeSpec(exir_ops.backend.tosa.MUL.default, (*node.args, shift), {}) + + +def rewrite_negate(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + input_node = cast(Node, node.args[0]) + dtype = input_node.meta["val"].dtype + input_zp, output_zp = 0, 0 + if dtype == torch.int8: + input_zp = get_input_qparams(node)[0].get_zp_per_tensor() + output_zp = get_output_qparams(node)[0].get_zp_per_tensor() + with node.graph.inserting_before(node): + input_zero_point = create_const(node.graph, input_zp, dtype, input_node) + output_zero_point = create_const(node.graph, output_zp, dtype, input_node) + return DialectNodeSpec( + exir_ops.backend.tosa.NEGATE.default, + (input_node, input_zero_point, output_zero_point), + {}, + ) + + def rewrite_binary_operator( node: Node, pass_: AtenToDialectPass ) -> DialectNodeSpec | None: @@ -51,8 +83,6 @@ def rewrite_binary_operator( target = exir_ops.backend.tosa.MAXIMUM.default case exir_ops.edge.aten.minimum.default: target = exir_ops.backend.tosa.MINIMUM.default - case exir_ops.edge.aten.mul.Tensor: - target = exir_ops.backend.tosa.MUL.default case exir_ops.edge.aten.pow.Tensor_Tensor: target = exir_ops.backend.tosa.POW.default case exir_ops.edge.aten.sub.Tensor: @@ -83,8 +113,6 @@ def rewrite_unary_operator( target = exir_ops.backend.tosa.LOG.default case exir_ops.edge.aten.logical_not.default: target = exir_ops.backend.tosa.LOGICAL_NOT.default - case exir_ops.edge.aten.neg.default: - target = exir_ops.backend.tosa.NEGATE.default case exir_ops.edge.aten.reciprocal.default: target = exir_ops.backend.tosa.RECIPROCAL.default case exir_ops.edge.aten.rsqrt.default: diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py index 0acb4e59fba..d74b9a6d092 100644 --- a/backends/arm/_passes/exir_to_tosa_pass.py +++ b/backends/arm/_passes/exir_to_tosa_pass.py @@ -20,6 +20,8 @@ ) from executorch.backends.arm._passes.aten_to_tosa_tensor_operators import ( rewrite_binary_operator, + rewrite_mul, + rewrite_negate, rewrite_rfft2, rewrite_unary_operator, ) @@ -65,6 +67,16 @@ def _get_fft_replacement( return rewrite_rfft2(node, pass_) +@register_dialect_substitutions(exir_ops.edge.aten.mul.Tensor) +def _get_mul_replacement(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + return rewrite_mul(node, pass_) + + +@register_dialect_substitutions(exir_ops.edge.aten.neg.default) +def _get_negate_replacement(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + return rewrite_negate(node, pass_) + + @register_dialect_substitutions( exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.bitwise_and.Tensor, @@ -77,7 +89,6 @@ def _get_fft_replacement( exir_ops.edge.aten.logical_xor.default, exir_ops.edge.aten.maximum.default, exir_ops.edge.aten.minimum.default, - exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.pow.Tensor_Tensor, exir_ops.edge.aten.sub.Tensor, ) @@ -107,7 +118,6 @@ def _get_comparison_operator_replacement( exir_ops.edge.aten.floor.default, exir_ops.edge.aten.log.default, exir_ops.edge.aten.logical_not.default, - exir_ops.edge.aten.neg.default, exir_ops.edge.aten.reciprocal.default, exir_ops.edge.aten.rsqrt.default, exir_ops.edge.aten.sin.default, diff --git a/backends/arm/_passes/insert_dynamic_padding.py b/backends/arm/_passes/insert_dynamic_padding.py index 344f7cd5b61..92700255eb4 100644 --- a/backends/arm/_passes/insert_dynamic_padding.py +++ b/backends/arm/_passes/insert_dynamic_padding.py @@ -59,15 +59,23 @@ def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue: zero_spatial_padding = [0] * (spatial_rank * 2) padding_shape = [*zero_padding_pair, *padding, *zero_padding_pair] + pad_value = args[6] if len(args) == 8 else None + if pad_value is None: + pad_value = super().call_operator( + exir_ops.backend.tosa.CONST.default, + ([0], [1], input_tensor.data.dtype), + {}, + meta, + True, + ) pad_res = super().call_operator( exir_ops.backend.tosa.PAD.default, ( input_tensor, padding_shape, + pad_value, ), - { - "value": 0, - }, + {}, meta, True, ) diff --git a/backends/arm/_passes/rewrite_pad.py b/backends/arm/_passes/rewrite_pad.py index d3631685de0..c4f0bddca7c 100644 --- a/backends/arm/_passes/rewrite_pad.py +++ b/backends/arm/_passes/rewrite_pad.py @@ -70,10 +70,17 @@ def _rewrite_constant_pad(self, input_tensor, pad, value, meta): exir_ops.backend.tosa.CONST_SHAPE.default, (shape,), {}, meta, True ) + pad_value = super().call_operator( + exir_ops.backend.tosa.CONST.default, + ([value], [1], output_dtype), + {}, + meta, + True, + ) return super().call_operator( exir_ops.backend.tosa.PAD.default, - (input_tensor, pad_shape), - {"value": value}, + (input_tensor, pad_shape, pad_value), + {}, meta, True, ) diff --git a/backends/arm/_passes/tosa_lowering_utils.py b/backends/arm/_passes/tosa_lowering_utils.py new file mode 100644 index 00000000000..e08973dcc66 --- /dev/null +++ b/backends/arm/_passes/tosa_lowering_utils.py @@ -0,0 +1,29 @@ +# Copyright 2026 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. +"""Build explicit TOSA operands during lowering.""" + +import torch +from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph, Node + + +def create_const(graph: Graph, values, dtype: torch.dtype, from_node: Node) -> Node: + if isinstance(values, torch.Tensor): + values = values.flatten().tolist() + elif not isinstance(values, (tuple, list)): + values = [values] + values = list(values) + node = create_node( + graph, + exir_ops.backend.tosa.CONST.default, + (values, [len(values)], dtype), + from_node=from_node, + ) + node.meta.pop(TosaSpecialDtype.meta_key(), None) + fake = from_node.meta["val"] + node.meta["val"] = fake.new_empty((len(values),), dtype=dtype) + return node diff --git a/backends/arm/operators/op_tosa_mul.py b/backends/arm/operators/op_tosa_mul.py index 28a0ccc6226..d39d8e4c1c1 100644 --- a/backends/arm/operators/op_tosa_mul.py +++ b/backends/arm/operators/op_tosa_mul.py @@ -3,7 +3,7 @@ # 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 Any, List +from typing import Any, cast, List import torch.fx import tosa_serializer as ts @@ -31,11 +31,11 @@ def define_node( inputs: List[TosaArg], output: TosaArg, ) -> None: - validate_num_inputs(self.target, inputs, 2) - validate_same_dtype(self.target, [*inputs, output]) + validate_num_inputs(self.target, inputs, 3) + validate_same_dtype(self.target, [*inputs[:2], output]) validate_valid_dtype( self.target, - [*inputs, output], + [*inputs[:2], output], [ ts.DType.INT8, ts.DType.INT16, @@ -46,7 +46,10 @@ def define_node( ], ) - shift = tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift") + shift_node = cast(torch.fx.Node, node.args[2]) + shift = tosa_graph.addConst( + [1], ts.DType.INT8, shift_node.args[0], name=f"{output.name}_shift" + ) attr = ts.TosaSerializerAttribute() attr.MulAttribute() self._serialize_operator( diff --git a/backends/arm/operators/op_tosa_neg.py b/backends/arm/operators/op_tosa_neg.py index 3b59e4ec570..d1ef5f8c218 100644 --- a/backends/arm/operators/op_tosa_neg.py +++ b/backends/arm/operators/op_tosa_neg.py @@ -53,7 +53,7 @@ def define_node( ts.DType.FP32, ] - validate_num_inputs(self.target, inputs, 1) + validate_num_inputs(self.target, inputs, 3) validate_same_dtype(self.target, [*inputs, output]) validate_valid_dtype(self.target, [*inputs, output], supported_dtypes) diff --git a/backends/arm/operators/op_tosa_pad.py b/backends/arm/operators/op_tosa_pad.py index a2645463834..261d8dd0add 100644 --- a/backends/arm/operators/op_tosa_pad.py +++ b/backends/arm/operators/op_tosa_pad.py @@ -3,7 +3,7 @@ # 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 Any, List +from typing import Any, cast, List import torch @@ -20,6 +20,7 @@ validate_valid_dtype, ) from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.exir.dialects._ops import ops as exir_ops @register_node_visitor @@ -37,19 +38,23 @@ def define_node( ) -> None: supported_dtypes = supported_data_layout_dtypes(self.tosa_spec) - validate_num_inputs(self.target, inputs, 2) + validate_num_inputs(self.target, inputs, 3) validate_same_dtype(self.target, [inputs[0], output]) validate_valid_dtype( self.target, [inputs[0], output], supported_dtypes, ) - pad_const = tosa_graph.addConst( - [1], - output.dtype, - [node.kwargs.get("value", 0)], - name=node.name + "_padding_value", - ) + value_node = cast(torch.fx.Node, node.args[2]) + pad_value_name = inputs[2].name + if value_node.target == exir_ops.backend.tosa.CONST.default: + pad_value = tosa_graph.addConst( + [1], + output.dtype, + value_node.args[0], + name=node.name + "_padding_value", + ) + pad_value_name = pad_value.name attr = ts.TosaSerializerAttribute() attr.PadAttribute() @@ -61,7 +66,7 @@ def define_node( [ inputs[0].name, inputs[1].name, - pad_const.name, + pad_value_name, ], [output.name], attr, diff --git a/backends/arm/process_node.py b/backends/arm/process_node.py index aaa644266f3..180da24ab01 100644 --- a/backends/arm/process_node.py +++ b/backends/arm/process_node.py @@ -17,6 +17,7 @@ from executorch.backends.arm.tosa.mapping import TosaArg from executorch.backends.arm.tosa.specification import TosaSpecification from executorch.backends.arm.tosa.utils import normalize_symint +from executorch.exir.dialects._ops import ops as exir_ops from torch._export.utils import ( get_buffer, get_lifted_tensor_constant, @@ -144,6 +145,11 @@ def process_call_function( node_visitors: Dict[str, NodeVisitor], tosa_spec: TosaSpecification, ): + if node.target == exir_ops.backend.tosa.CONST.default: + # Temporarily skip dialect constants: visitors still create serializer + # constants until the node-visitor flow is replaced by direct lowering. + return + # Unpack arguments and convert try: inputs = [TosaArg(arg, tosa_spec) for arg in node.args] diff --git a/backends/arm/test/misc/test_tosa_dialect_constant_visitors.py b/backends/arm/test/misc/test_tosa_dialect_constant_visitors.py new file mode 100644 index 00000000000..770a496ba45 --- /dev/null +++ b/backends/arm/test/misc/test_tosa_dialect_constant_visitors.py @@ -0,0 +1,108 @@ +# Copyright 2026 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. + +import executorch.backends.arm.tosa.dialect # noqa: F401 +import numpy as np +import pytest +import torch +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm._passes.tosa_lowering_utils import create_const +from executorch.backends.arm.operators.node_visitor import get_node_visitors +from executorch.backends.arm.process_node import process_call_function, process_inputs +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx import Graph +from tosa.TosaGraph import TosaGraph # type: ignore[import-untyped] + + +@pytest.mark.parametrize( + "target,dtype,values", + [ + (exir_ops.backend.tosa.MUL.default, torch.float32, [0]), + (exir_ops.backend.tosa.MUL.default, torch.int32, [17]), + (exir_ops.backend.tosa.NEGATE.default, torch.float32, [0, 0]), + (exir_ops.backend.tosa.NEGATE.default, torch.int8, [3, -4]), + (exir_ops.backend.tosa.PAD.default, torch.float32, [2.5]), + (exir_ops.backend.tosa.PAD.default, torch.int8, [7]), + ], +) +def test_visitors_recreate_constants_without_serializing_dialect_const( + target, dtype, values +): + spec = TosaSpecification.create_from_string("TOSA-1.0+INT+FP") + visitors = get_node_visitors(spec) + assert "tosa.CONST.default" not in visitors + serializer = TosaSerializerWithConstantPool( + "", targetMajor=1, targetMinor=0, targetPatch=0, targetDraft=False + ) + graph = Graph() + tensor = graph.placeholder("input_tensor") + with FakeTensorMode(): + tensor.meta["val"] = torch.empty(2, 3, dtype=dtype) + constant_dtype = ( + torch.int8 if target == exir_ops.backend.tosa.MUL.default else dtype + ) + constants = [create_const(graph, value, constant_dtype, tensor) for value in values] + if target == exir_ops.backend.tosa.MUL.default: + args = (tensor, tensor, constants[0]) + elif target == exir_ops.backend.tosa.NEGATE.default: + args = (tensor, *constants) + else: + padding = graph.call_function( + exir_ops.backend.tosa.CONST_SHAPE.default, ([0, 0, 1, 1],) + ) + padding.meta["val"] = [0, 0, 1, 1] + padding.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.SHAPE + args = (tensor, padding, constants[0]) + operation = graph.call_function(target, args) + operation.meta["val"] = tensor.meta["val"] + if target == exir_ops.backend.tosa.NEGATE.default and dtype == torch.int8: + operation.meta["input_qparams"] = { + 0: QuantArgs(0.5, values[0], -128, 127, dtype) + } + operation.meta["output_qparams"] = { + 0: QuantArgs(0.5, values[1], -128, 127, dtype) + } + graph.output(operation) + + process_inputs(tensor, serializer, spec) + for node in graph.nodes: + if node.op == "call_function": + process_call_function(node, serializer, visitors, spec) + + block = TosaGraph.GetRootAs(serializer.serialize(), 0).Regions(0).Blocks(0) + tensor_names = { + block.Tensors(index).Name().decode() for index in range(block.TensorsLength()) + } + assert all(constant.name not in tensor_names for constant in constants) + declared_names = tensor_names | { + block.Shapes(index).Name().decode() for index in range(block.ShapesLength()) + } + for index in range(block.OperatorsLength()): + serialized = block.Operators(index) + assert all( + serialized.Inputs(input_index).decode() in declared_names + for input_index in range(serialized.InputsLength()) + ) + operation_inputs = block.Operators(block.OperatorsLength() - 1) + constant_inputs_start = 1 if target == exir_ops.backend.tosa.NEGATE.default else 2 + numpy_dtype = { + torch.float32: np.float32, + torch.int8: np.int8, + torch.int32: np.int32, + }[constant_dtype] + for index, expected in enumerate(values, start=constant_inputs_start): + name = operation_inputs.Inputs(index).decode() + serialized_tensor = next( + block.Tensors(tensor_index) + for tensor_index in range(block.TensorsLength()) + if block.Tensors(tensor_index).Name().decode() == name + ) + assert np.frombuffer( + serialized_tensor.DataAsNumpy(), dtype=numpy_dtype + ).tolist() == [expected] diff --git a/backends/arm/test/misc/tosa_dialect/test_data_layout_ops.py b/backends/arm/test/misc/tosa_dialect/test_data_layout_ops.py index 35074fc32b5..e82eee850c7 100644 --- a/backends/arm/test/misc/tosa_dialect/test_data_layout_ops.py +++ b/backends/arm/test/misc/tosa_dialect/test_data_layout_ops.py @@ -26,7 +26,9 @@ def _fake_tensor(dtype: torch.dtype, mode: FakeTensorMode) -> torch.Tensor: id="concat", ), pytest.param( - lambda x: exir_ops.backend.tosa.PAD.default(x, [1, 2, 3, 4], value=0), + lambda x: exir_ops.backend.tosa.PAD.default( + x, [1, 2, 3, 4], value=x.new_zeros((1,)) + ), (5, 10), id="pad", ), @@ -89,7 +91,9 @@ def test_data_layout_ops_positive(op, expected_shape, spec, dtype) -> None: id="concat", ), pytest.param( - lambda x: exir_ops.backend.tosa.PAD.default(x, [0, -1, 0, 0], value=0), + lambda x: exir_ops.backend.tosa.PAD.default( + x, [0, -1, 0, 0], value=x.new_zeros((1,)) + ), "non-negative", id="pad", ), @@ -165,7 +169,7 @@ def test_pad_rejects_wrong_padding_length() -> None: exir_ops.backend.tosa.PAD.default( mode.from_tensor(torch.randn((2, 3), dtype=torch.float32)), [1, 2], - value=0.0, + value=torch.zeros(1), ) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py index c4a811f3a2c..5cf0d8c4005 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py @@ -194,6 +194,8 @@ def test_tosa_binary_ops( with TosaLoweringContext( TosaSpecification.create_from_string(spec) ), FakeTensorMode() as mode: + if op_name == "MUL": + kwargs = {**kwargs, "shift": torch.zeros(1, dtype=torch.int8)} output = getattr(exir_ops.backend.tosa, op_name).default( *_to_fake(mode, input1, input2), **kwargs, @@ -343,7 +345,32 @@ def test_pow_accepts_bfloat16_with_bf16_extension() -> None: assert tuple(output.shape) == tuple(input1.shape) -def test_mul_rejects_non_zero_shift_for_non_int32() -> None: +@pytest.mark.parametrize("shift", [-1, 64]) +def test_mul_rejects_out_of_range_shift(shift): + from executorch.backends.arm.tosa.dialect.ops.binary_elementwise import MUL + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+INT")): + with pytest.raises(TosaValueError, match="range"): + MUL( + torch.ones(2, dtype=torch.int32), + torch.ones(2, dtype=torch.int32), + torch.tensor([shift], dtype=torch.int8), + ) + + +def test_mul_rejects_nonzero_shift_for_narrow_input(): + from executorch.backends.arm.tosa.dialect.ops.binary_elementwise import MUL + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+INT")): + with pytest.raises(TosaValueError, match="Only int32"): + MUL( + torch.ones(2, dtype=torch.int8), + torch.ones(2, dtype=torch.int8), + torch.tensor([1], dtype=torch.int8), + ) + + +def test_mul_rejects_non_int8_shift() -> None: input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) @@ -352,11 +379,11 @@ def test_mul_rejects_non_zero_shift_for_non_int32() -> None: ), FakeTensorMode() as mode: with pytest.raises( TosaValueError, - match="Only int32 MUL supports a non-zero shift value", + match="MUL shift must be a single-element int8 tensor", ): exir_ops.backend.tosa.MUL.default( *_to_fake(mode, input1, input2), - shift=3, + shift=torch.tensor([3], dtype=torch.int32), ) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py index 9bfd33d4e0c..89f7cb423a9 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py @@ -127,9 +127,11 @@ def test_tosa_unary_ops( with TosaLoweringContext( TosaSpecification.create_from_string(spec) ), FakeTensorMode() as mode: - output = getattr(exir_ops.backend.tosa, op_name).default( - mode.from_tensor(input_tensor) - ) + operands: tuple[torch.Tensor, ...] = (mode.from_tensor(input_tensor),) + if op_name == "NEGATE": + zero_point = torch.zeros(1, dtype=input_tensor.dtype) + operands += (zero_point, zero_point) + output = getattr(exir_ops.backend.tosa, op_name).default(*operands) assert output.dtype == input_tensor.dtype assert tuple(output.shape) == tuple(input_tensor.shape) @@ -346,9 +348,11 @@ def test_tosa_float_unary_ops_accept_bfloat16_with_bf16_extension( with TosaLoweringContext( TosaSpecification.create_from_string("TOSA-1.1+FP+bf16") ), FakeTensorMode() as mode: - output = getattr(exir_ops.backend.tosa, op_name).default( - mode.from_tensor(input_tensor) - ) + operands: tuple[torch.Tensor, ...] = (mode.from_tensor(input_tensor),) + if op_name == "NEGATE": + zero_point = torch.zeros(1, dtype=input_tensor.dtype) + operands += (zero_point, zero_point) + output = getattr(exir_ops.backend.tosa, op_name).default(*operands) assert output.dtype == torch.bfloat16 assert tuple(output.shape) == tuple(input_tensor.shape) @@ -360,8 +364,11 @@ def test_negate_rejects_bfloat16_without_bf16_extension() -> None: with TosaLoweringContext( TosaSpecification.create_from_string("TOSA-1.1+FP") ), FakeTensorMode() as mode: + zero_point = torch.zeros(1, dtype=sample_input.dtype) with pytest.raises(TosaValueError, match="doesn't support bfloat16"): - exir_ops.backend.tosa.NEGATE.default(mode.from_tensor(sample_input)) + exir_ops.backend.tosa.NEGATE.default( + mode.from_tensor(sample_input), zero_point, zero_point + ) def test_abs_rejects_int8() -> None: diff --git a/backends/arm/test/ops/test_to_copy.py b/backends/arm/test/ops/test_to_copy.py index fb6d43c6f13..adb27a80f6b 100644 --- a/backends/arm/test/ops/test_to_copy.py +++ b/backends/arm/test/ops/test_to_copy.py @@ -126,7 +126,6 @@ def test_to_tosa_FP(test_data: Tuple): aten_op=[], exir_op=[], ) - # int to int cast is not supported in TOSA+FP profile if not new_dtype.is_floating_point and not torch.is_floating_point(test_tensor): pipeline.change_args( "check_count.exir", diff --git a/backends/arm/test/passes/test_exir_to_tosa_pass.py b/backends/arm/test/passes/test_exir_to_tosa_pass.py new file mode 100644 index 00000000000..bad3ce5162c --- /dev/null +++ b/backends/arm/test/passes/test_exir_to_tosa_pass.py @@ -0,0 +1,116 @@ +# Copyright 2026 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. + +import pytest +import torch +from executorch.backends.arm._passes.exir_to_tosa_pass import ExirToTosaPass +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm._passes.rewrite_argmax_pass import RewriteArgmaxPass +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class ArgmaxToInt32(torch.nn.Module): + def forward(self, tensor): + return tensor.argmax(dim=1).to(torch.int32) + + +class CastToInt32(torch.nn.Module): + def forward(self, tensor): + return tensor.to(torch.int32) + + +class ArgmaxCastChain(torch.nn.Module): + def forward(self, tensor): + return tensor.argmax(dim=1).to(torch.int32).to(torch.float32) + + +class CastChain(torch.nn.Module): + def forward(self, tensor): + return tensor.to(torch.int32).to(torch.float32) + + +class PointwiseOperands(torch.nn.Module): + def forward(self, tensor, other): + return tensor * other, -tensor + + +@pytest.mark.parametrize( + "module,expected_casts,expected_dtype", + [ + (ArgmaxToInt32(), 0, torch.int32), + (CastToInt32(), 1, torch.int32), + (ArgmaxCastChain(), 1, torch.float32), + (CastChain(), 2, torch.float32), + ], +) +def test_lowering_preserves_remaining_dimension_order_copies( + module, expected_casts, expected_dtype +): + inputs = (torch.randn(2, 3),) + program = to_edge(torch.export.export(module, inputs)).exported_program() + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")): + rewritten = RewriteArgmaxPass()(program.graph_module) + assert rewritten is not None + result = ExirToTosaPass(program)(rewritten.graph_module) + assert result is not None + assert ( + sum( + node.target == exir_ops.edge.dim_order_ops._to_dim_order_copy.default + for node in result.graph_module.graph.nodes + ) + == expected_casts + ) + assert not any( + node.target == exir_ops.backend.tosa.CAST.default + for node in result.graph_module.graph.nodes + ) + output = result.graph_module.graph.output_node().args[0][0] + assert output.meta["val"].dtype == expected_dtype + + +@pytest.mark.parametrize("dtype", (torch.float32, torch.int8, torch.int16, torch.int32)) +def test_pointwise_rewrites_materialize_explicit_operands(dtype): + inputs = (torch.ones(2, 3, dtype=dtype), torch.ones(2, 3, dtype=dtype)) + program = to_edge( + torch.export.export(PointwiseOperands(), inputs) + ).exported_program() + input_zp, output_zp = 0, 0 + if dtype == torch.int8: + input_zp, output_zp = 3, -4 + negate = next( + node + for node in program.graph.nodes + if node.target == exir_ops.edge.aten.neg.default + ) + negate.meta["input_qparams"] = {0: QuantArgs(0.5, input_zp, -128, 127, dtype)} + negate.meta["output_qparams"] = {0: QuantArgs(0.5, output_zp, -128, 127, dtype)} + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT+FP")): + result = ExirToTosaPass(program)(program.graph_module) + assert result is not None + result.graph_module.graph.lint() + multiply = next( + node + for node in result.graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.MUL.default + ) + assert len(multiply.args) == 3 + shift = multiply.args[2] + assert shift.target == exir_ops.backend.tosa.CONST.default + assert shift.args == ([0], [1], torch.int8) + + negate = next( + node + for node in result.graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.NEGATE.default + ) + assert len(negate.args) == 3 + assert negate.args[1].args == ([input_zp], [1], dtype) + assert negate.args[2].args == ([output_zp], [1], dtype) diff --git a/backends/arm/test/passes/test_remove_data_layout_noops.py b/backends/arm/test/passes/test_remove_data_layout_noops.py index 081686973ff..93ed6b57418 100644 --- a/backends/arm/test/passes/test_remove_data_layout_noops.py +++ b/backends/arm/test/passes/test_remove_data_layout_noops.py @@ -18,6 +18,7 @@ SymbolicToTosaShapesPass, ) from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager +from executorch.backends.arm._passes.tosa_lowering_utils import create_const from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec from executorch.backends.arm.tosa.dialect.lib import TosaValueError from executorch.backends.arm.tosa.specification import ( @@ -72,7 +73,11 @@ def _tosa_data_layout_graph( ) for values in shape_values ] - kwargs = {"value": 0.0} if target == exir_ops.backend.tosa.PAD.default else {} + kwargs = ( + {"value": create_const(graph, 0.0, x.meta["val"].dtype, x)} + if target == exir_ops.backend.tosa.PAD.default + else {} + ) result = _call( graph, target, diff --git a/backends/arm/tosa/dialect/__init__.py b/backends/arm/tosa/dialect/__init__.py index ee6e56b39fb..ea129031e21 100644 --- a/backends/arm/tosa/dialect/__init__.py +++ b/backends/arm/tosa/dialect/__init__.py @@ -10,6 +10,7 @@ avg_pool2d_adaptive, binary_elementwise, comparison, + constants, conv2d, conv2d_block_scaled, conv3d, diff --git a/backends/arm/tosa/dialect/ops/binary_elementwise.py b/backends/arm/tosa/dialect/ops/binary_elementwise.py index 5c181ae0a80..2c932e7d26a 100644 --- a/backends/arm/tosa/dialect/ops/binary_elementwise.py +++ b/backends/arm/tosa/dialect/ops/binary_elementwise.py @@ -256,24 +256,26 @@ def MINIMUM( @register_fake_tosa_op( - "MUL(Tensor input1, Tensor input2, *, int shift=0) -> Tensor", + "MUL(Tensor input1, Tensor input2, Tensor shift) -> Tensor", TosaSpecification.all_versions_and_profiles(), ) def MUL( input1: torch.Tensor, input2: torch.Tensor, - *, - shift: int = 0, + shift: torch.Tensor, ) -> torch.Tensor: output_dtype = _validate_and_infer_mul_output_dtype(input1.dtype) - if shift < 0 or shift > 63: - raise TosaValueError("shift must be in the range [0, 63]", op="MUL") - if input1.dtype != torch.int32 and shift != 0: - raise TosaValueError( - "Only int32 MUL supports a non-zero shift value", - op="MUL", - ) + if shift.dtype != torch.int8 or shift.shape != (1,): + raise TosaValueError("MUL shift must be a single-element int8 tensor", op="MUL") + if not isinstance(shift, torch._subclasses.fake_tensor.FakeTensor): + shift_value = shift.item() + if not 0 <= shift_value <= 63: + raise TosaValueError("shift must be in the range [0, 63]", op="MUL") + if input1.dtype != torch.int32 and shift_value != 0: + raise TosaValueError( + "Only int32 MUL supports a non-zero shift value", op="MUL" + ) return binary_meta(input1, input2, "MUL", output_dtype=output_dtype) diff --git a/backends/arm/tosa/dialect/ops/constants.py b/backends/arm/tosa/dialect/ops/constants.py new file mode 100644 index 00000000000..aaab9428014 --- /dev/null +++ b/backends/arm/tosa/dialect/ops/constants.py @@ -0,0 +1,23 @@ +# Copyright 2026 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. +"""Explicit tensor constants in the TOSA dialect.""" + +import torch + +from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op +from executorch.backends.arm.tosa.specification import TosaSpecification + + +@register_tosa_op( + "CONST(Scalar[] values, int[] shape, ScalarType dtype) -> Tensor", + TosaSpecification.all_versions_and_profiles(), + include_real_impl=False, +) +def CONST( + values: list[int | float], shape: list[int], dtype: torch.dtype +) -> torch.Tensor: + if len(values) != torch.Size(shape).numel(): + raise ValueError("CONST values must match its shape") + return torch.empty(shape, dtype=dtype) diff --git a/backends/arm/tosa/dialect/ops/data_layout_ops.py b/backends/arm/tosa/dialect/ops/data_layout_ops.py index 52ca1366320..d938cdf1c46 100644 --- a/backends/arm/tosa/dialect/ops/data_layout_ops.py +++ b/backends/arm/tosa/dialect/ops/data_layout_ops.py @@ -125,11 +125,17 @@ def CONCAT(inputs: list[torch.Tensor], *, axis: int) -> torch.Tensor: @register_fake_tosa_op( - "PAD(Tensor input1, SymInt[] padding, *, Scalar value) -> Tensor", + "PAD(Tensor input1, SymInt[] padding, Tensor value) -> Tensor", TosaSpecification.all_versions_and_profiles(), ) -def PAD(x: torch.Tensor, padding: list[int | torch.SymInt], *, value) -> torch.Tensor: +def PAD( + x: torch.Tensor, padding: list[int | torch.SymInt], value: torch.Tensor +) -> torch.Tensor: _validate_data_layout_dtype(x.dtype, "PAD") + if value.dtype != x.dtype or value.numel() != 1: + raise TosaValueError( + "PAD value must be a scalar tensor with the input dtype", op="PAD" + ) if len(padding) != 2 * len(x.shape): raise TosaValueError( diff --git a/backends/arm/tosa/dialect/ops/unary_elementwise.py b/backends/arm/tosa/dialect/ops/unary_elementwise.py index 80808d66d2b..42fd4a64dce 100644 --- a/backends/arm/tosa/dialect/ops/unary_elementwise.py +++ b/backends/arm/tosa/dialect/ops/unary_elementwise.py @@ -152,11 +152,19 @@ def LOGICAL_NOT(input1: torch.Tensor) -> torch.Tensor: @register_fake_tosa_op( - "NEGATE(Tensor input1) -> Tensor", + "NEGATE(Tensor input1, Tensor input_zp, Tensor output_zp) -> Tensor", DUAL_PROFILE_SPECS, ) -def NEGATE(input1: torch.Tensor) -> torch.Tensor: +def NEGATE( + input1: torch.Tensor, input_zp: torch.Tensor, output_zp: torch.Tensor +) -> torch.Tensor: _validate_negate_dtype(input1.dtype) + for zero_point in (input_zp, output_zp): + if zero_point.numel() != 1 or zero_point.dtype != input1.dtype: + raise TosaValueError( + "NEGATE zero points must be scalar tensors matching the input dtype", + op="NEGATE", + ) return torch.empty_like(input1, dtype=input1.dtype)