Qualcomm AI Engine Direct - Adding QNN backend support for scatter_reduce.two and scatter_add core ATen ops

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)

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,

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

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)

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

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):

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):

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__()

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):