From 5e0aff3351e833cf418b3cf78297985639236f56 Mon Sep 17 00:00:00 2001 From: Youngsik Yang Date: Sun, 31 May 2026 01:05:01 +0900 Subject: [PATCH] Arm backend: support depthwise Conv3D DecomposeGroupedConvPass skips depthwise convolutions (in_channels == groups) because Conv2D depthwise maps to TOSA DEPTHWISE_CONV2D. For Conv3D (rank-5 input) there is no DEPTHWISE_CONV3D op, so the skip caused RewriteConvPass to raise a RuntimeError during preprocess. Narrow the depthwise skip to Conv2D only. For rank-5 inputs the pass falls through to the existing slice->CONV3D->cat decomposition, which already handles arbitrary groups and produces the correct per-channel computations that depthwise convolution requires. Signed-off-by: Youngsik Yang --- .../_passes/decompose_grouped_conv_pass.py | 6 ++- backends/arm/_passes/rewrite_conv_pass.py | 10 ++-- backends/arm/test/ops/test_conv3d.py | 50 +++++++++++++++---- 3 files changed, 48 insertions(+), 18 deletions(-) diff --git a/backends/arm/_passes/decompose_grouped_conv_pass.py b/backends/arm/_passes/decompose_grouped_conv_pass.py index ed0adbe83d7..a0180fe79e7 100644 --- a/backends/arm/_passes/decompose_grouped_conv_pass.py +++ b/backends/arm/_passes/decompose_grouped_conv_pass.py @@ -252,8 +252,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] diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index a51f1ae0555..54c443dd04a 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/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 diff --git a/backends/arm/test/ops/test_conv3d.py b/backends/arm/test/ops/test_conv3d.py index 7348809a0de..ee24e8a7d8d 100644 --- a/backends/arm/test/ops/test_conv3d.py +++ b/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)