Arm backend: support depthwise Conv3D (pytorch#19902)

## Summary

Depthwise Conv3D (`in_channels == groups`, rank-5 input) previously
crashed with a `RuntimeError` inside `RewriteConvPass` because TOSA has
no `DEPTHWISE_CONV3D` op. `DecomposeGroupedConvPass` already handles
non-depthwise grouped Conv3D by splitting it into `groups==1`
convolutions via slice→conv→cat, but it explicitly skipped the depthwise
case since Conv2D depthwise maps to the native `DEPTHWISE_CONV2D` TOSA
op.

For Conv3D there is no such native op, so the fix is to extend
`DecomposeGroupedConvPass` to stop skipping depthwise when the input is
rank 5(Conv3D).
The existing slice→`CONV3D`→cat decomposition can handle it correctly.

```mermaid
flowchart LR
    DW2D["Depthwise Conv2D\n(in_channels == groups, rank 4)"]
    DW3D["Depthwise Conv3D\n(in_channels == groups, rank 5)"]
    GRP["DecomposeGroupedConvPass"]
    RC2D["RewriteConvPass"]
    RC3D["RewriteConvPass"]
    DELEGATE_CONV2D["DEPTHWISE_CONV2D"]
    DELEGATE_CONV3D["CONV3D"]

    DW2D --> RC2D
    DW3D -->|"decomposed"| GRP
    GRP -->|"CONV3D (groups==1)"| RC3D
    RC2D -->|"delegated to native op"| DELEGATE_CONV2D
    RC3D -->|"delegated to native op"| DELEGATE_CONV3D

```

## Files changed:

| File | Change |
| --- | --- |
| `backends/arm/_passes/decompose_grouped_conv_pass.py` | In
`call_operator`, narrow the depthwise skip to Conv2D only
(`len(input.data.shape) != 5`); for rank-5 inputs(Conv3D) fall through
to the existing decomposition. |
| `backends/arm/_passes/rewrite_conv_pass.py` | Update comment in
`_is_conv3d` to reflect that both grouped and depthwise Conv3D are now
decomposed upstream; retain the `RuntimeError` as defense-in-depth. |
| `backends/arm/test/ops/test_conv3d.py` | Rewrite
`test_convolution_3d_tosa_FP_depthwise` to assert delegation |

## Test result

```bash
python -m pytest backends/arm/test/ops/test_conv3d.py::test_convolution_u55_INT_not_delegated_3d
# 2 passed, 0 failed.
```

```bash
lintrunner -a \
    backends/arm/_passes/decompose_grouped_conv_pass.py \
    backends/arm/_passes/rewrite_conv_pass.py \
    backends/arm/test/ops/test_conv3d.py
# ok No lint issues.
```

Signed-off-by: Youngsik Yang <vacu9708@gmail.com>

Author: vacu9708

backends/arm/_passes/decompose_grouped_conv_pass.py

@@ -257,8 +257,10 @@ def call_operator(self, op, args, kwargs, meta):
 
         input_node = args[0]
         if DecomposeGroupedConvPass._is_depthwise_conv(input_node, groups, transposed):
-            # This is a depthwise convolution which is handled elsewhere
-            return super().call_operator(op, args, kwargs, meta)
+            # Conv2D depthwise maps to TOSA DEPTHWISE_CONV2D — handled in RewriteConvPass.
+            # Conv3D has no DEPTHWISE_CONV3D, so fall through and decompose like grouped conv.
+            if len(input_node.data.shape) != 5:
+                return super().call_operator(op, args, kwargs, meta)
 
         weight_node = args[1]
         bias_node = args[2]

backends/arm/_passes/rewrite_conv_pass.py

@@ -129,13 +129,13 @@ def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool:
 
     def _is_conv3d(self, rank, groups) -> bool:
         if rank == 5:
-            # A Conv3D is considered depthwise if Group == InChannels and
-            # Group * N == OutChannels, where N is a possitive integer.
-            # Currently we do not support depthwise or grouped conv3d.
-            # @TODO Add grouped/depthwise conv3d support or reject in partitioner.
+            # Both grouped and depthwise Conv3D are decomposed into groups==1
+            # convolutions by DecomposeGroupedConvPass before reaching here.
+            # This guard is defense-in-depth for paths that bypass that pass.
             if groups != 1:
                 raise RuntimeError(
-                    "CONV3D with groups != 1 is not supported in the Arm backend."
+                    "CONV3D with groups != 1 reached unexpectedly; "
+                    "DecomposeGroupedConvPass should have decomposed it first."
                 )
             return True
         return False

backends/arm/test/ops/test_conv3d.py

@@ -212,6 +212,32 @@ def forward(self, x):
         return self.conv(x)
 
 
+class GroupedConv3d(torch.nn.Module):
+    """Non-depthwise grouped Conv3d (in_channels != groups).
+
+    Split into ``groups`` plain convolutions by DecomposeGroupedConvPass, so it
+    is delegated unlike the depthwise case.
+
+    """
+
+    def __init__(self, dtype=torch.float):
+        super().__init__()
+        self.dtype = dtype
+        self.conv = torch.nn.Conv3d(
+            in_channels=4,
+            out_channels=4,
+            kernel_size=(3, 3, 3),
+            padding=1,
+            groups=2,
+        ).to(dtype)
+
+    def get_inputs(self):
+        return (torch.randn(1, 4, 8, 8, 8).to(self.dtype),)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
 conv3d_2x2_3x2x14x14_nobias = Conv3d(
     in_channels=2,
     out_channels=3,
@@ -623,19 +649,21 @@ def test_convolution_3d_tosa_INT_multi_op():
 
 
 def test_convolution_3d_tosa_FP_depthwise():
-    """Depthwise or Grouped Conv3d should be rejected until grouped support
-    exists.
+    """Depthwise Conv3d should be delegated, decomposed into groups==1
+    convolutions by DecomposeGroupedConvPass.
     """
     model = DepthwiseConv3d()
-    pipeline = TosaPipelineFP[input_t](
-        model,
-        model.get_inputs(),
-        aten_op,
-        exir_op,
-        run_on_tosa_ref_model=False,
-    )
-    with pytest.raises(RuntimeError, match="CONV3D with groups != 1"):
-        pipeline.run()
+    pipeline = TosaPipelineFP[input_t](model, model.get_inputs(), aten_op, exir_op)
+    pipeline.run()
+
+
+def test_convolution_3d_tosa_FP_grouped():
+    """Non-depthwise grouped Conv3d should be delegated, decomposed into
+    groups==1 convolutions by DecomposeGroupedConvPass.
+    """
+    model = GroupedConv3d()
+    pipeline = TosaPipelineFP[input_t](model, model.get_inputs(), aten_op, exir_op)
+    pipeline.run()
 
 
 @common.parametrize("test_data", test_data_INT)