Arm backend: Extend convolution support check to 3d

backends/arm/operator_support/convolution_support.py

@@ -55,7 +55,7 @@ def _is_node_supported_u55(self, node: fx.Node):
 
         C_in = shape_in[1]
         C_out = shape_out[1]
-        if (C_in == group) and (C_out % C_in) == 0:
+        if (C_in == group) and (C_out % C_in) == 0 and len(shape_in) <= 4:
             # Depthwise convolution
             for dim in shape_in[1:]:
                 if not 1 <= dim <= 65536:
@@ -74,13 +74,19 @@ def _is_node_supported_u55(self, node: fx.Node):
 
         kernel_w = kernel[2]
         kernel_h = kernel[3] if len(kernel) > 3 else 1
+        kernel_z = kernel[4] if len(kernel) > 4 else 1
         # Kernel condition misses constraint on sum of absolute weights
         if not 1 <= kernel_h <= 64 or not 1 <= kernel_w * kernel_h <= 4096:
             self.reporter.report_reject(
                 node,
                 f"Convolution needs to have kernel_y<=64, kernel_x*kernel_y<=4096, got kernel ({kernel_w}, {kernel_h})",
             )
             return False
+        if kernel_z != 1:
+            self.reporter.report_reject(
+                node, f"Convolution3d needs to have kernel_z==1, got {kernel_z}."
+            )
+            return False
 
         if not self._stride_condition(node):
             self.reporter.report_reject(
@@ -107,6 +113,14 @@ def _stride_condition(self, node: fx.Node) -> bool:
         if len(strides) == 1:
             strides = [strides[0]] * 2
 
+        if len(strides) > 2:
+            stride_z = strides[2]
+            if stride_z > 1:
+                self.reporter.report_reject(
+                    node, f"Convolution3d only supports stride_z<=1, got {stride_z}."
+                )
+                return False
+
         for stride, dilation in zip(strides, dilations):
             stride_condition = 1 <= stride <= 3
             dilation_condition = (not has_padding) and (dilation == 1)

backends/arm/test/ops/test_conv2d.py

@@ -8,10 +8,10 @@
 
 import torch
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.arm_tester import ArmTester
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineBI,
     EthosU85PipelineBI,
+    OpNotSupportedPipeline,
     TosaPipelineBI,
     TosaPipelineMI,
 )
@@ -34,9 +34,9 @@ def __init__(
         in_channels: Union[List, int, None] = None,
         out_channels: Union[List, int, None] = None,
         kernel_size: Union[List, Tuple, None] = None,
-        stride: Union[List, Tuple, None] = None,
-        padding: Union[List, Tuple, None] = None,
-        dilation: Union[List, Tuple, None] = None,
+        stride: Union[List, Tuple, int, None] = None,
+        padding: Union[List, Tuple, int, None] = None,
+        dilation: Union[List, Tuple, int, None] = None,
         groups: Union[List, int, None] = None,
         bias: Union[List, bool, None] = None,
         padding_mode: Union[List, str, None] = None,
@@ -446,17 +446,9 @@ def test_convolution_2d_u85_BI_on_fvp(test_module):
 def test_reject_convolution_2d_u55_BI(
     module: Conv2d,
 ):
-    (
-        ArmTester(
-            module,
-            example_inputs=module.get_inputs(),
-            compile_spec=common.get_u55_compile_spec(),
-        )
-        .quantize()
-        .export()
-        .check_count({"torch.ops.aten.conv2d.default": 1})
-        .check(["torch.ops.quantized_decomposed"])
-        .to_edge_transform_and_lower()
-        .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
-        .check_count({"torch.ops.higher_order.executorch_call_delegate": 0})
-    )
+    OpNotSupportedPipeline(
+        module,
+        module.get_inputs(),
+        "TOSA-0.80+BI+u55",
+        {"executorch_exir_dialects_edge__ops_aten_convolution_default": 1},
+    ).run()