diff --git a/backends/qualcomm/_passes/i64_to_i32.py b/backends/qualcomm/_passes/i64_to_i32.py index 42be979c576..a26956b5569 100644 --- a/backends/qualcomm/_passes/i64_to_i32.py +++ b/backends/qualcomm/_passes/i64_to_i32.py @@ -42,6 +42,8 @@ class I64toI32(ExportPass): exir_ops.edge.aten.gather.default: [2], exir_ops.edge.aten.scatter.src: [2], exir_ops.edge.aten.scatter.value: [2], + exir_ops.edge.aten.scatter_add.default: [2], + exir_ops.edge.aten.scatter_reduce.two: [2], } copy_op = exir_ops.edge.aten._to_copy.default @@ -170,7 +172,14 @@ def _cast_op_args_to_i64(self, graph_module: torch.fx.GraphModule): (input_node,), {"dtype": torch.int64}, ) - cast_i64_node.meta["val"] = node.meta["val"].to(torch.int64) + # This cast produces the *index* tensor, so its + # FakeTensor must be derived from the argument being + # cast, not from the op output: index.shape == + # output.shape for gather, but for scatter* the output + # takes the shape of 'self', which may differ. + cast_i64_node.meta["val"] = input_node.meta["val"].to( + torch.int64 + ) args_list = list(node.args) args_list[arg_index] = cast_i64_node node.args = tuple(args_list) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index ac146956611..51b968f1d85 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -122,6 +122,8 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.round.default, exir_ops.edge.aten.scatter.src, exir_ops.edge.aten.scatter.value, + exir_ops.edge.aten.scatter_add.default, + exir_ops.edge.aten.scatter_reduce.two, exir_ops.edge.aten.sigmoid.default, exir_ops.edge.aten.sign.default, exir_ops.edge.aten.slice_copy.Tensor, diff --git a/backends/qualcomm/builders/op_scatter_elements.py b/backends/qualcomm/builders/op_scatter_elements.py index d5abc04a375..8a3ab0b3fba 100644 --- a/backends/qualcomm/builders/op_scatter_elements.py +++ b/backends/qualcomm/builders/op_scatter_elements.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 Dict +from typing import Dict, Optional import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager @@ -22,16 +22,65 @@ @register_node_visitor class ScatterElements(NodeVisitor): - target = ["aten.scatter.src", "aten.scatter.value"] + target = [ + "aten.scatter.src", + "aten.scatter.value", + "aten.scatter_add.default", + "aten.scatter_reduce.two", + ] + + # aten reduce string -> QNN reduction mode. "mean" / "amax" / "amin" + # are intentionally absent: QNN ScatterElements cannot express them. + reduce_str_to_reduction = { + "sum": OpScatterElements.Reduction.ADD, + "prod": OpScatterElements.Reduction.MUL, + } def __init__(self, *args) -> None: super().__init__(*args) + def _get_reduction( + self, node: torch.fx.Node + ) -> Optional[OpScatterElements.Reduction]: + """ + Resolve the QNN reduction mode for this node, or None if the node + cannot be represented by QNN ScatterElements (caller falls back to CPU). + """ + op_name = node.target.__name__ + + if op_name == "aten.scatter_add.default": + return OpScatterElements.Reduction.ADD + + if op_name == "aten.scatter_reduce.two": + # include_self is keyword-only in aten.scatter_reduce.two + include_self = node.kwargs.get("include_self", True) + if not include_self: + # QNN always accumulates onto the existing values of 'self' + return None + + reduce_str = ( + node.args[4] if len(node.args) > 4 else node.kwargs.get("reduce") + ) + return self.reduce_str_to_reduction.get(reduce_str) + + # aten.scatter.src: plain overwrite + return OpScatterElements.Reduction.NONE + def define_node( self, node: torch.fx.Node, nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], ) -> PyQnnManager.PyQnnOpWrapper: + reduction = self._get_reduction(node) + if reduction is None: + # unsupported reduce mode or include_self=False -> fall back to CPU + return None + + # NOTE: QNN HTP only supports reduction != NONE in quantized mode. We + # intentionally do not gate on that here: backend capability is resolved + # by IsNodeSupportedByBackend during partitioning, so this builder stays + # backend-agnostic and picks up capability changes across SDK versions + # automatically. The fp case is covered in the rework tests. input_node = self.get_node(node.args[0]) input_tensor = self.get_tensor(input_node, node) input_tensor_wrapper = self.define_tensor( @@ -128,7 +177,7 @@ def define_node( scatter_op.AddScalarParam( OpScatterElements.param_reduction, PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, - {QCOM_DATA: np.uint32(OpScatterElements.Reduction.NONE)}, + {QCOM_DATA: np.uint32(reduction)}, ) return scatter_op diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index 6c18d78fdb5..0df2e5a4731 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -609,6 +609,8 @@ class OpScatterElements: @unique class Reduction(IntEnum): NONE = 0 + ADD = 1 + MUL = 2 @dataclass(init=False, frozen=True) diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 93f00d4e994..f4ecc07f33f 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -69,6 +69,8 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: torch.ops.aten.rms_norm.default, torch.ops.aten._safe_softmax.default, torch.ops.aten.scatter.src, + torch.ops.aten.scatter_add.default, + torch.ops.aten.scatter_reduce.two, torch.ops.aten.stack.default, torch.ops.aten.upsample_bicubic2d.vec, # This request is ignored because it is in a blocklist. Refer to exir/program/_program.py diff --git a/backends/qualcomm/quantizer/annotators/htp_rules.py b/backends/qualcomm/quantizer/annotators/htp_rules.py index 80060223508..a319c332edb 100644 --- a/backends/qualcomm/quantizer/annotators/htp_rules.py +++ b/backends/qualcomm/quantizer/annotators/htp_rules.py @@ -1444,7 +1444,12 @@ class ScaledDotProductAttention(GeneralOpDef): @register_annotator( - [torch.ops.aten.scatter.src, torch.ops.aten.scatter.value], + [ + torch.ops.aten.scatter.src, + torch.ops.aten.scatter.value, + torch.ops.aten.scatter_add.default, + torch.ops.aten.scatter_reduce.two, + ], qnn_op=None, ) class ScatterElements(GeneralOpDef): diff --git a/backends/qualcomm/quantizer/annotators/lpai_rules.py b/backends/qualcomm/quantizer/annotators/lpai_rules.py index fa68a9d3d8c..178654f18fc 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -885,7 +885,12 @@ class ScaledDotProductAttention(GeneralOpDef): @register_annotator( - [torch.ops.aten.scatter.src, torch.ops.aten.scatter.value], + [ + torch.ops.aten.scatter.src, + torch.ops.aten.scatter.value, + torch.ops.aten.scatter_add.default, + torch.ops.aten.scatter_reduce.two, + ], qnn_op=None, ) class ScatterElements(GeneralOpDef): diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 5930869ecd4..564a11a3f21 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2413,6 +2413,25 @@ def forward(self, query_layer, key_layer, value_layer, attn_mask): return attn_output +class ScatterAdd(torch.nn.Module): + def __init__(self, dim=1): + super().__init__() + self.dim = dim + + def forward(self, data, index, src): + return torch.scatter_add(data, self.dim, index, src) + + +class ScatterReduce(torch.nn.Module): + def __init__(self, dim=1, reduce="sum"): + super().__init__() + self.dim = dim + self.reduce = reduce + + def forward(self, data, index, src): + return data.scatter_reduce(self.dim, index, src, reduce=self.reduce) + + class ScatterSrc(torch.nn.Module): def __init__(self, dim=1): 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..177e3f0639d 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -1271,6 +1271,46 @@ def test_sdpa(request, kwargs): ScaledDotProductAttention.test(request, kwargs) # noqa: F405 +# QNN HTP ScatterElements with reduction != NONE is only supported in quantized +# mode; the fp16 backend validator rejects it, so the fp case falls back to CPU. +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_scatter_add(request, kwargs): + ScatterAdd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_scatter_reduce_sum(request, kwargs): + ScatterReduce.test_sum(request, kwargs) # noqa: F405 + + +# "prod" multiplies up to 3 values per output element, so the relative error +# compounds multiplicatively and needs a looser bound than "sum". +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + CosineSimilarity(0.95), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_scatter_reduce_prod(request, kwargs): + ScatterReduce.test_prod(request, kwargs) # noqa: F405 + + @enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_scatter_src(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index c2ba792811c..3db4430cefe 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -4208,6 +4208,163 @@ def test(subtests, qnn_config, quantizer, compile_spec, expected): ) +class ScatterAdd(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, data, index, src): + return torch.scatter_add(data, self.dim, index, src) + + @staticmethod + @unpack_fixtures + def test(subtests, qnn_config, quantizer, compile_spec, expected): + # duplicate indices are intentional: they are what distinguishes + # scatter_add (accumulate) from scatter (overwrite) + cases = [ + ( + 1, + ( + torch.ones(3, 5), + torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ), + ( + 0, + ( + torch.ones(3, 5), + torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ), + # negative dim exercises the "dim % rank" normalization + ( + -1, + ( + torch.ones(3, 5), + torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ), + # 4D exercises the QCOM_AXIS_ORDER remap in the builder + ( + 1, + ( + torch.ones(1, 4, 2, 3), + torch.randint(0, 4, (1, 4, 2, 3), dtype=torch.int64), + torch.rand(1, 4, 2, 3), + ), + ), + ] + for dim, inputs in cases: + with subtests.test(msg=f"dim:{dim}, shape:{tuple(inputs[0].shape)}"): + with expected as metrics: + export_and_verify( + module=__class__(dim=dim), + inputs=inputs, + qnn_config=qnn_config, + quantizer=quantizer, + compile_specs=compile_spec, + metrics=metrics, + ) + + +class ScatterReduce(torch.nn.Module): + def __init__(self, dim, reduce): + super().__init__() + self.dim = dim + self.reduce = reduce + + def forward(self, data, index, src): + return data.scatter_reduce(self.dim, index, src, reduce=self.reduce) + + # index containing duplicates, so the reduction actually combines values + _INDEX_DIM1 = torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], dtype=torch.int64 + ) + _INDEX_DIM0 = torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], dtype=torch.int64 + ) + + @staticmethod + def _run(subtests, qnn_config, quantizer, compile_spec, expected, cases): + for module, inputs in cases: + with subtests.test( + msg=f"reduce:{module.reduce}, dim:{module.dim}, " + f"shape:{tuple(inputs[0].shape)}" + ): + with expected as metrics: + export_and_verify( + module=module, + inputs=inputs, + qnn_config=qnn_config, + quantizer=quantizer, + compile_specs=compile_spec, + metrics=metrics, + ) + + @staticmethod + @unpack_fixtures + def test_sum(subtests, qnn_config, quantizer, compile_spec, expected): + cls = ScatterReduce + cases = [ + ( + cls(dim=1, reduce="sum"), + (torch.ones(3, 5), cls._INDEX_DIM1, torch.rand(3, 5)), + ), + ( + cls(dim=0, reduce="sum"), + (torch.ones(3, 5), cls._INDEX_DIM0, torch.rand(3, 5)), + ), + # negative dim exercises the "dim % rank" normalization + ( + cls(dim=-1, reduce="sum"), + (torch.ones(3, 5), cls._INDEX_DIM1, torch.rand(3, 5)), + ), + # 4D exercises the QCOM_AXIS_ORDER remap in the builder + ( + cls(dim=1, reduce="sum"), + ( + torch.ones(1, 4, 2, 3), + torch.randint(0, 4, (1, 4, 2, 3), dtype=torch.int64), + torch.rand(1, 4, 2, 3), + ), + ), + ] + cls._run(subtests, qnn_config, quantizer, compile_spec, expected, cases) + + @staticmethod + @unpack_fixtures + def test_prod(subtests, qnn_config, quantizer, compile_spec, expected): + cls = ScatterReduce + # keep src bounded away from 0 so the running product does not + # collapse toward the bottom of the quantization range. + # Only dim=0/1 here: neg-dim normalization and the 4-D axis-order + # remap are builder-level paths already covered by test_sum, and + # "prod" adds no new coverage there (only more error compounding). + cases = [ + ( + cls(dim=1, reduce="prod"), + (torch.ones(3, 5), cls._INDEX_DIM1, torch.rand(3, 5) + 0.5), + ), + ( + cls(dim=0, reduce="prod"), + (torch.ones(3, 5), cls._INDEX_DIM0, torch.rand(3, 5) + 0.5), + ), + ] + cls._run(subtests, qnn_config, quantizer, compile_spec, expected, cases) + + class ScatterSrc(torch.nn.Module): def __init__(self, dim): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 1ebda343c8c..cd755ade77c 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2274,6 +2274,11 @@ def test_qnn_backend_round(self): sample_input = (torch.randn([3, 4]),) self.lower_module_and_test_output(module, sample_input) + # NOTE: only scatter.src (reduction=NONE) is delegatable in fp16. QNN HTP + # ScatterElements rejects reduction != NONE in the fp backend validator, so + # scatter_add / scatter_reduce have no fp tests here. See + # backends/qualcomm/tests/rework/htp/op/v68/test.py, which asserts the + # expected fp failure explicitly. def test_qnn_backend_scatter_src(self): test_comb = [ { @@ -5872,6 +5877,36 @@ def test_qnn_backend_scatter_src(self): qdq_module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_scatter_add(self): + index_dim1 = torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], dtype=torch.int64 + ) + module = ScatterAdd(dim=1) # noqa: F405 + sample_input = (torch.ones(3, 5), index_dim1, torch.rand(3, 5)) + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + + def test_qnn_backend_scatter_reduce_sum(self): + index_dim1 = torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], dtype=torch.int64 + ) + module = ScatterReduce(dim=1, reduce="sum") # noqa: F405 + sample_input = (torch.ones(3, 5), index_dim1, torch.rand(3, 5)) + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + + def test_qnn_backend_scatter_reduce_prod(self): + index_dim1 = torch.tensor( + [[0, 1, 2, 0, 1], [2, 0, 1, 2, 0], [1, 2, 0, 1, 2]], dtype=torch.int64 + ) + # "prod" multiplies up to 3 values per output element, so in 8a8w the + # relative error compounds multiplicatively; loosen the bound. + self.atol, self.rtol = 3e-1, 1 + module = ScatterReduce(dim=1, reduce="prod") # noqa: F405 + sample_input = (torch.ones(3, 5), index_dim1, torch.rand(3, 5) + 0.5) + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_scatter_value(self): test_comb = [ {