From da92cc59fe4c61081f73e3620490e60e96f0c45a Mon Sep 17 00:00:00 2001 From: Arik Horodniceanu Date: Fri, 26 Jun 2026 15:33:54 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - Adding QNN backend support for scatter.value core ATen op --- backends/qualcomm/_passes/i64_to_i32.py | 1 + backends/qualcomm/_passes/layout_transform.py | 1 + .../qualcomm/builders/op_scatter_elements.py | 63 +++++++++++++--- .../quantizer/annotators/htp_rules.py | 2 +- .../quantizer/annotators/lpai_rules.py | 2 +- backends/qualcomm/tests/models.py | 10 +++ backends/qualcomm/tests/test_qnn_delegate.py | 73 +++++++++++++++++++ 7 files changed, 139 insertions(+), 13 deletions(-) diff --git a/backends/qualcomm/_passes/i64_to_i32.py b/backends/qualcomm/_passes/i64_to_i32.py index 986dd60543f..c034c41f962 100644 --- a/backends/qualcomm/_passes/i64_to_i32.py +++ b/backends/qualcomm/_passes/i64_to_i32.py @@ -41,6 +41,7 @@ class I64toI32(ExportPass): I64_IN_OPS = { exir_ops.edge.aten.gather.default: [2], exir_ops.edge.aten.scatter.src: [2], + exir_ops.edge.aten.scatter.value: [2], } copy_op = exir_ops.edge.aten._to_copy.default diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 5b9c13e6ef4..ecca8cdc23b 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -121,6 +121,7 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.relu.default, exir_ops.edge.aten.round.default, exir_ops.edge.aten.scatter.src, + exir_ops.edge.aten.scatter.value, 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 4bcf4572803..f4ab932b1b7 100644 --- a/backends/qualcomm/builders/op_scatter_elements.py +++ b/backends/qualcomm/builders/op_scatter_elements.py @@ -9,7 +9,13 @@ import numpy as np import torch -from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA +from executorch.backends.qualcomm.utils.constants import ( + QCOM_AXIS_ORDER, + QCOM_DATA, + QCOM_QUANT_ATTRS, + QCOM_SCALE, + QCOM_ZERO_POINT, +) from .node_visitor import NodeVisitor from .node_visitor_manager import register_node_visitor @@ -18,7 +24,7 @@ @register_node_visitor class ScatterElements(NodeVisitor): - target = ["aten.scatter.src"] + target = ["aten.scatter.src", "aten.scatter.value"] def __init__(self, *args) -> None: super().__init__(*args) @@ -48,15 +54,50 @@ def define_node( nodes_to_wrappers, ) - updates_node = self.get_node(node.args[3]) - updates_tensor = self.get_tensor(updates_node, node) - updates_tensor_wrapper = self.define_tensor( - updates_node, - node, - updates_tensor, - PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - nodes_to_wrappers, - ) + # Handle both scatter.src (args[3] is a Node/tensor) and + # scatter.value (args[3] is a scalar) + if isinstance(node.args[3], torch.fx.Node): + # scatter.src path: args[3] is a tensor node + updates_node = self.get_node(node.args[3]) + updates_tensor = self.get_tensor(updates_node, node) + updates_tensor_wrapper = self.define_tensor( + updates_node, + node, + updates_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + else: + # scatter.value: expand scalar to a static tensor matching index shape. + scalar_val = node.args[3] + quant_attrs = input_node.meta.get(QCOM_QUANT_ATTRS, None) + + if quant_attrs is not None: + # Quantized: pre-quantize the scalar to match the data tensor's dtype. + # QNN accepts STATIC tensors with raw quantized data. + scale = quant_attrs[QCOM_SCALE] + zero_point = quant_attrs[QCOM_ZERO_POINT] + quant_dtype = quant_attrs.get("dtype", torch.uint8) + quantized_val = int(round(scalar_val / scale) + zero_point) + info = torch.iinfo(quant_dtype) + quantized_val = max(info.min, min(info.max, quantized_val)) + updates_tensor = torch.full( + index_tensor.shape, quantized_val, dtype=quant_dtype + ) + else: + updates_tensor = torch.full( + index_tensor.shape, scalar_val, dtype=input_tensor.dtype + ) + + updates_tensor_wrapper = self.define_tensor( + index_node, + node, + updates_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + node_name=f"{node.name}_updates_value", + wrapper_idx=2, + ) output_tensor = self.get_tensor(node, node) output_tensor_wrapper = self.define_tensor( diff --git a/backends/qualcomm/quantizer/annotators/htp_rules.py b/backends/qualcomm/quantizer/annotators/htp_rules.py index ca8abb246bf..61a92501b3d 100644 --- a/backends/qualcomm/quantizer/annotators/htp_rules.py +++ b/backends/qualcomm/quantizer/annotators/htp_rules.py @@ -1427,7 +1427,7 @@ class ScaledDotProductAttention(GeneralOpDef): @register_annotator( - [torch.ops.aten.scatter.src], + [torch.ops.aten.scatter.src, torch.ops.aten.scatter.value], 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 6e5b343c5c7..9119ce33c79 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -874,7 +874,7 @@ class ScaledDotProductAttention(GeneralOpDef): @register_annotator( - [torch.ops.aten.scatter.src], + [torch.ops.aten.scatter.src, torch.ops.aten.scatter.value], qnn_op=None, ) class ScatterElements(GeneralOpDef): diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 6dbf1183659..2c439205fe4 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2304,6 +2304,16 @@ def forward(self, data, index, src): return torch.scatter(data, self.dim, index, src) +class ScatterValue(torch.nn.Module): + def __init__(self, dim=1, value=0.5): + super().__init__() + self.dim = dim + self.value = value + + def forward(self, data, index): + return torch.scatter(data, self.dim, index, self.value) + + class SelectCopy(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 98c5839c515..06af17ec3bd 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2194,6 +2194,42 @@ def test_qnn_backend_scatter_src(self): index += 1 self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_scatter_value(self): + test_comb = [ + { + QCOM_MODULE: [ScatterValue(dim=1, value=0.5)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.rand(3, 5), + torch.tensor( + [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [1, 0, 3, 4, 2]], + dtype=torch.int64, + ), + ), + ], + }, + { + QCOM_MODULE: [ScatterValue(dim=0, value=1.0)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.rand(3, 5), + torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], + dtype=torch.int64, + ), + ), + ], + }, + ] + + 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_rsqrt(self): module = Rsqrt() # noqa: F405 sample_input = (torch.abs(torch.randn([3, 4])),) @@ -5330,6 +5366,43 @@ 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_value(self): + test_comb = [ + { + QCOM_MODULE: [ScatterValue(dim=1, value=0.5)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.rand(3, 5), + torch.tensor( + [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [1, 0, 3, 4, 2]], + dtype=torch.int64, + ), + ), + ], + }, + { + QCOM_MODULE: [ScatterValue(dim=0, value=1.0)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.rand(3, 5), + torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], + dtype=torch.int64, + ), + ), + ], + }, + ] + + 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_sdpa(self): modules = [ ScaledDotProductAttention(), # noqa: F405