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

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)

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` |

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

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

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

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]),)