Qualcomm AI Engine Direct - Adding QNN backend support for reflection_pad3d core ATen op - #21502

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 94 New Failures, 3 Pending, 5 Unclassified Failures

As of commit 1738773 with merge base 454f979 ():

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

@pytorchbot label "release notes: qualcomm"

@claude pls review

Claude finished @psiddh's task in 8m 56s —— View job

Code Review: reflection_pad3d via index_select + cat

• Read PR diff and surrounding pass/pipeline code
• Analyze decomposition correctness
• Check dialect/pipeline consistency
• Review tests & docs

Overall this is a clean, well-scoped addition and the decomposition approach is sound. The _pad_dim index math is correct for reflection semantics (excludes the edge row/col, matching torch.nn.ReflectionPad3d), the sequential per-dim padding is a valid way to compose N-D reflection, and asymmetric padding is covered by tests. A few points worth addressing:

1. Inserted ops are raw ATen, not edge-dialect (main concern)

DecomposePad runs in the to_edge capture-program pipeline where the pad node can be an EdgeOpOverload — the existing reflection_pad2d path explicitly branches on is_edge and picks the edge op variant (_PAD_OPS[("reflect", 4, True)]). The new branch ignores is_edge and always inserts ATen index_select/cat from torch.export.export(...).module():

# backends/qualcomm/_passes/decompose_pad.py:96-112
if mode == "reflect" and len(padding) == 6:
    model = ReflectionPad3d(list(padding))
    decomposed_module = torch.export.export(model, (node.args[0].meta["val"],), strict=True).module()
    ...

Compare with DecomposeMaxPool3d, which deliberately wraps the decomposed module in to_edge(...) for the non-quantization (edge) path so the inserted nodes are edge-dialect (decompose_maxpool3d.py:112-118). DecomposeRoll inserts raw ATen slice/cat and appears to work, so the QNN partitioner likely tolerates this — but the inconsistency is worth resolving. Please confirm the ATen index_select/cat nodes are recognized by the partitioner in the edge pipeline (partitioner matches on node.target.__name__), and consider aligning with the DecomposeMaxPool3d to_edge-wrapping pattern for consistency. Fix this →

2. dead_code_elimination_pass now runs for every invocation

It was added at the end of call() unconditionally. That's fine functionally, but it now runs even when no pad node was rewritten. Not a correctness issue — just noting it's slightly broader than needed. Acceptable to keep.

3. Docstring placement is now misleading

reflection_pad3d is still listed under the "Not supported by QNN" heading in the class docstring, even though it's now supported (via decomposition):

Not supported by QNN (max rank 4 for non-constant schemes):
- mode='reflect', 6 padding values (3d) -> decomposed into index_select + cat

Consider moving it under the "Supported" section (e.g. "mode='reflect', 6 padding values (3d) -> decomposed into index_select + cat, since QNN MIRROR_REFLECT max rank is 4").

4. Minor: implicit precondition on padding vs. dim size

_pad_dim assumes pad < size (true for valid reflection padding — PyTorch enforces it at export). No action needed, but if a malformed pad ever reached here, torch.arange(size-2, size-2-pad_after, -1) could yield negative indices silently. The upstream guarantee makes this safe.

Tests & docs

• Both FP and quantized delegate tests added, plus rework test_5d and symmetric + asymmetric models.py variants — good coverage.
• README table update is correct.

Nothing here is blocking beyond confirming point (1). Nice work.
· branch reflection_pad3d

@psiddh
Addressed the minor comments from Claude:

• Comment 1: The pass only targets aten.pad.default (not edge dialect). The 5D branch can only trigger in ATen IR during the export pipeline, so the inserted ATen nodes are correct. Clarified the docstring.

• Comment 2: Acknowledged. Consistent with other passes (e.g. DecomposeCDist). No change.

• Comment 3: Fixed, moved to the Supported section.

• Comment 4: Good catch. PyTorch enforces pad < size at export time, so this is safe. No action needed.