From 10f8be4efc7855c87dd83cdde1ea1c8498889f85 Mon Sep 17 00:00:00 2001 From: Arik Horodniceanu Date: Mon, 20 Jul 2026 16:28:39 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - Adding QNN backend support for reflection_pad3d core ATen op --- backends/qualcomm/_passes/decompose_pad.py | 68 ++++++++++++++++++- backends/qualcomm/builders/README.md | 2 +- backends/qualcomm/tests/models.py | 18 +++++ .../qualcomm/tests/rework/htp/op/v68/test.py | 6 ++ backends/qualcomm/tests/rework/src/op.py | 15 ++++ backends/qualcomm/tests/test_qnn_delegate.py | 41 +++++++++++ 6 files changed, 148 insertions(+), 2 deletions(-) diff --git a/backends/qualcomm/_passes/decompose_pad.py b/backends/qualcomm/_passes/decompose_pad.py index ebfa8fd1a8f..10051a3a6d0 100644 --- a/backends/qualcomm/_passes/decompose_pad.py +++ b/backends/qualcomm/_passes/decompose_pad.py @@ -8,6 +8,50 @@ from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass + +from .utils import merge_decomposed_graph + + +class ReflectionPad3d(torch.nn.Module): + """Implement reflection_pad3d using index_select + cat. + + reflection_pad3d operates on 5D tensors [N, C, D, H, W] with padding + [left, right, top, bottom, front, back]. QNN HTP's Pad op with + MIRROR_REFLECT scheme only supports max rank 4 tensors, so we implement + it using index_select + cat which are fully supported at 5D rank. + """ + + def __init__(self, padding): + super().__init__() + self.padding = padding + + @staticmethod + def _pad_dim(x, dim, pad_before, pad_after): + """Apply reflection padding along a single dimension.""" + size = x.shape[dim] + parts = [] + if pad_before > 0: + indices = torch.arange( + pad_before, 0, -1, device=x.device, dtype=torch.int32 + ) + parts.append(torch.index_select(x, dim, indices)) + parts.append(x) + if pad_after > 0: + indices = torch.arange( + size - 2, size - 2 - pad_after, -1, device=x.device, dtype=torch.int32 + ) + parts.append(torch.index_select(x, dim, indices)) + if len(parts) > 1: + return torch.cat(parts, dim=dim) + return x + + def forward(self, x): + left, right, top, bottom, front, back = self.padding + x = self._pad_dim(x, 4, left, right) + x = self._pad_dim(x, 3, top, bottom) + x = self._pad_dim(x, 2, front, back) + return x class DecomposePad(ExportPass): @@ -20,7 +64,7 @@ class DecomposePad(ExportPass): - mode='reflect', 4 padding values -> reflection_pad2d (QNN MIRROR_REFLECT, max rank 4). Not supported by QNN (max rank 4 for non-constant schemes): - - mode='reflect', 6 padding values (3d) -> reflection_pad3d (QNN MIRROR_REFLECT max rank is 4) + - mode='reflect', 6 padding values (3d) -> decomposed into index_select + cat - mode='replicate' -> QNN EDGE scheme produces incorrect results for FP32 inputs for replication_pad2d Note: reflection_pad1d is handled by PyTorch's built-in decomposition of aten.pad.default (mode='reflect', 2 padding values) @@ -46,6 +90,27 @@ def call(self, graph_module: torch.fx.GraphModule): padding = node.args[1] is_edge = isinstance(node.target, EdgeOpOverload) + + # Handle 5D reflect padding (reflection_pad3d) via decomposition into index_select + cat. + # QNN HTP's Pad op with MIRROR_REFLECT only supports max rank 4. + 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() + with graph.inserting_before(node): + remap = {"x": node.args[0]} + merge_decomposed_graph( + remap=remap, + target_node=node, + target_graph=graph, + decomposed_graph_module=decomposed_module, + ) + graph.erase_node(node) + continue + target_op = self._PAD_OPS.get((mode, len(padding), is_edge)) if target_op is None: continue @@ -54,4 +119,5 @@ def call(self, graph_module: torch.fx.GraphModule): node.args = (node.args[0], list(padding)) graph_module.recompile() + dead_code_elimination_pass(graph_module) return PassResult(graph_module, True) diff --git a/backends/qualcomm/builders/README.md b/backends/qualcomm/builders/README.md index cff3d41cd12..78f477b67ee 100644 --- a/backends/qualcomm/builders/README.md +++ b/backends/qualcomm/builders/README.md @@ -524,7 +524,7 @@ The following PyTorch operators are supported through decomposition or annotatio | `aten.pdist`, `aten._pdist_forward` | `DecomposePDist` | | `aten.reciprocal` | `DecomposeReciprocal` | | `aten.reflection_pad1d` | PyTorch built-in decomposition | -| `aten.reflection_pad2d` | `DecomposePad` | +| `aten.reflection_pad2d`, `aten.reflection_pad3d` | `DecomposePad` | | `aten.remainder.Scalar`, `aten.remainder.Tensor` | `DecomposeRemainder` | | `aten.roll` | `DecomposeRoll` | | `aten.select_scatter` | `DecomposeSelectScatter` | diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 0b35141fe41..f9581ccb7fc 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2119,6 +2119,24 @@ def forward(self, x): return self.pad(x) +class ReflectionPad3d(torch.nn.Module): + def __init__(self): + super().__init__() + self.pad = torch.nn.ReflectionPad3d((1, 1, 1, 1, 1, 1)) + + def forward(self, x): + return self.pad(x) + + +class ReflectionPad3dAsymmetric(torch.nn.Module): + def __init__(self): + super().__init__() + self.pad = torch.nn.ReflectionPad3d((2, 1, 2, 1, 2, 1)) + + def forward(self, x): + return self.pad(x) + + class Relu(torch.nn.Module): def __init__(self): 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 b227da6be93..77de76bcba2 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -1135,6 +1135,12 @@ def test_reflection_pad_2d(request, kwargs): ReflectionPad.test_4d(request, kwargs) # noqa: F405 +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + @enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_relu(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index b773b55dfad..efc67c84356 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -3754,6 +3754,21 @@ def test_4d(subtests, qnn_config, quantizer, compile_spec, expected): metrics=metrics, ) + @staticmethod + @unpack_fixtures + def test_5d(qnn_config, quantizer, compile_spec, expected): + inputs = (torch.randn(1, 3, 6, 8, 8),) + padding = (1, 1, 1, 1, 1, 1) + with expected as metrics: + export_and_verify( + module=__class__(3, padding), + inputs=inputs, + qnn_config=qnn_config, + quantizer=quantizer, + compile_specs=compile_spec, + metrics=metrics, + ) + class Relu(torch.nn.Module): def __init__(self): diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index f4c45fe209e..77cac91171f 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2074,6 +2074,27 @@ def test_qnn_backend_reflection_pad2d(self): index += 1 self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_reflection_pad3d(self): + test_comb = [ + { + QCOM_MODULE: [ + ReflectionPad3d(), # noqa: F405 + ReflectionPad3dAsymmetric(), # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [ + (torch.randn(1, 3, 6, 8, 8),), + ], + }, + ] + + 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_relu(self): module = Relu() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),) @@ -5250,6 +5271,26 @@ def test_qnn_backend_reflection_pad2d(self): qdq_module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_reflection_pad3d(self): + test_comb = [ + { + QCOM_MODULE: [ + ReflectionPad3d(), # noqa: F405 + ReflectionPad3dAsymmetric(), # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [(torch.randn(1, 3, 6, 8, 8),)], + }, + ] + + 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_relu(self): module = Relu() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),)