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

[qti-horodnic]

Summary

Added support for the aten.scatter_add.default and aten.scatter_reduce.two core ATen ops. Extended the existing aten.scatter.src builder into a unified ScatterElements implementation. Reductions map as: scatter.src → NONE, scatter_add → ADD, scatter_reduce sum → ADD / prod → MUL.

Note the following limitations:

• No FP support for scatter_add / scatter_reduce: QNN op ScatterElements with reduction != NONE is only supported in quantized mode.
• Modes mean, amax, amin fall back to CPU.
• include_self=False not supported: QNN always accumulates onto existing self values; falls back to CPU.

Note this PR also includes a minor bug fix:

I64toI32._cast_op_args_to_i64 now derives the inserted cast node's FakeTensor from the argument being cast rather than from the op output. This incidentally fixes a latent wrong-shape meta["val"] on the pre-existing aten.scatter.src path: index.shape == output.shape for gather (no behavior change there), but scatter* outputs take the shape of self, which may differ.

Test plan

pytest backends/qualcomm/tests/rework/htp/op/v68/test.py -k "test_scatter_add" -v
pytest backends/qualcomm/tests/rework/htp/op/v68/test.py -k "test_scatter_reduce_sum" -v
pytest backends/qualcomm/tests/rework/htp/op/v68/test.py -k "test_scatter_reduce_prod" -v


python backends/qualcomm/tests/test_qnn_delegate.py TestQNNQuantizedOperator.test_qnn_backend_scatter_add --soc_model SM8750 --host aisw-vm15-labsd --device 545ee4aa --build_folder build-android
python backends/qualcomm/tests/test_qnn_delegate.py TestQNNQuantizedOperator.test_qnn_backend_scatter_reduce_sum --soc_model SM8750 --host aisw-vm15-labsd --device 545ee4aa --build_folder build-android
python backends/qualcomm/tests/test_qnn_delegate.py TestQNNQuantizedOperator.test_qnn_backend_scatter_reduce_prod --soc_model SM8750 --host aisw-vm15-labsd --device 545ee4aa --build_folder build-android

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22157

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit eb2fddb with merge base a77eb9f ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[qti-horodnic]

@pytorchbot label "release notes: qualcomm"

[psiddh]

@claude review this

[claude]

Claude finished @psiddh's task in 3m 54s —— View job

---

Code Review

I reviewed the diff against the current branch. This is a clean, well-reasoned PR — the design of routing capability decisions through the builder's None return is elegant and the commit message is unusually thorough. Below are my findings.

Tasks

• Read the full diff and surrounding context
• Review builder / partition / quantizer logic
• Review the two bug fixes
• Review test coverage
• Post review

---

✅ Strengths

• Consistent partition ↔ preprocess behavior. define_node returning None is checked both in qnn_partitioner.py:130 (via IsNodeSupportedByBackend) and qnn_preprocess.py:83. Because partitioning also calls define_node, an unsupported node (include_self=False, mean/amax/amin) is rejected at partition time and never reaches preprocess — so it can't be silently dropped from the QNN graph. This is the subtle correctness property that makes the fallback safe, and it holds.
• i64_to_i32 fix is a genuine correctness improvement (i64_to_i32.py:180). Deriving the cast's FakeTensor from input_node rather than the op output is correct: for scatter* the output takes self's shape, which can differ from the index shape. The comment explains the gather-vs-scatter distinction well.
• scatter_reduce.two arg handling is correct. include_self is keyword-only in the schema (after *), so node.kwargs.get("include_self", True) is right and it can never appear positionally. reduce at args[4] with the kwargs fallback covers both normalized and unnormalized forms.
• Test cases deliberately use duplicate indices, which is exactly what exercises accumulate-vs-overwrite semantics. Good instinct.

🔸 Minor / discussion

1. Annotation happens for modes that then fall back to CPU. The htp_rules.py / lpai_rules.py ScatterElements annotator (htp_rules.py:1446) annotates all four targets unconditionally, including scatter_reduce.two with mean/amax/amin or include_self=False. Those nodes get q/dq inserted, then the partitioner rejects them, leaving the op on CPU surrounded by quantize/dequantize. This is usually harmless (later passes fold the danglers) and matches existing patterns, but worth confirming it doesn't leave an un-foldable q/dq pair for the unsupported-reduce path.

2. No test locks in the documented CPU-fallback limitations. The PR carefully documents three fallbacks (mean/amax/amin, include_self=False), and _get_reduction implements them, but I don't see a test asserting those actually fall back (i.e. remain un-delegated rather than silently miscompiling). A single subtest for e.g. reduce="amax" or include_self=False would guard against a future QNN SDK adding partial support and this branch going stale. Fix this →

3. String-based op dispatch. _get_reduction switches on node.target.__name__ string literals (op_scatter_elements.py:49-67). This is consistent with other QNN visitors, so no change requested — just flagging it's a bit more fragile than comparing against the exir_ops.edge.aten.* overloads.

4. PR body references a conftest.py fix that isn't in this diff. Bug fix Re-sync with internal repository #2 in the description ("export_and_verify … now only inspects the exported graph on the quantized path") — the current conftest.py:481 already has the if quantizer is not None: guard, but conftest.py is not among the 11 changed files. Either that fix landed separately (via Load an older PTD file in the Python reader instead of rejecting it #22117, which last touched the file) or it's missing from this PR. If your rework FP subtests depend on it, please confirm it's actually present on the merge target, otherwise the [fp] assertions in this PR may not behave as described.

Nits

• qnn_constants.py:610 Reduction enum additions are clean.
• The tolerance loosening for prod (both test_prod cosine-sim 0.95 and the delegate test's atol/rtol) is well-justified in comments.

Overall: LGTM pending confirmation on point #4 (the conftest fix) and a nice-to-have on point #2. Nothing blocking in the builder/partition/quantizer logic itself.
· branch add_reduce_scatter

[qti-horodnic]

Regarding Claude's comments:
1-3 are general conventions in our code, will keep the current version as-is for uniformity.
Regarding 4, that comment is correct, my note was outdated since the fix has already been merged. Removed the note from the description.