Support two modes of depthwise conv

Differential Revision: D94432987

Pull Request resolved: #17723

Author: DrJessop

backends/cadence/aot/ops_registrations.py

@@ -1179,18 +1179,25 @@ def quantized_conv2d_nhwc_per_tensor_meta(
     in_size = input.shape
     # Assert that the input tensor has at least 3 dimensions, and at most 6
     assert len(in_size) > 2
-    assert len(in_size) < 6
+    assert len(in_size) < 5
 
     # Determine weight layout based on depthwise vs regular conv.
     in_channels = in_size[-1]
-    if is_depthwise_conv(groups, in_channels):
-        *kernel_size, out_channels = weight.shape
-    elif len(in_size) == 3:
-        # 1D conv: weight is [OC, K, IC]
-        out_channels, *kernel_size, _ = weight.shape
+    depthwise = is_depthwise_conv(groups, in_channels)
+    if len(in_size) == 3:
+        if len(weight.shape) == 2:
+            assert depthwise
+            *kernel_size, out_channels = weight.shape
+        else:
+            out_channels, *kernel_size, _ = weight.shape
+    elif len(in_size) == 4:
+        if len(weight.shape) == 3:
+            assert depthwise
+            *kernel_size, out_channels = weight.shape
+        else:
+            out_channels, *kernel_size, _ = weight.shape
     else:
-        # 2D regular conv: weight is [OC, KH, KW, IC]
-        out_channels, *kernel_size, _ = weight.shape
+        raise ValueError("Unsupported input tensor size")
 
     # Compute the output tensor size
     output_size = (

backends/cadence/aot/ref_implementations.py

@@ -1110,7 +1110,8 @@ def quantized_conv2d_nhwc_per_tensor(
     if len(input_tensor.shape) == 3:
         # 1D conv: input is [N, L, C] -> [N, C, L]
         input_tensor = input_tensor.movedim(-1, 1).contiguous()
-        if depthwise:
+        if len(weight.shape) == 2:
+            assert depthwise, "1D depthwise conv requires 2D weight tensor"
             # 1D depthwise: weight is [K, OC] -> [OC, 1, K]
             weight = weight.permute(1, 0).unsqueeze(1).contiguous()
         else:
@@ -1120,7 +1121,8 @@ def quantized_conv2d_nhwc_per_tensor(
     else:
         # 2D conv: input is [N, H, W, C] -> [N, C, H, W]
         input_tensor = input_tensor.movedim(-1, -3)
-        if depthwise:
+        if len(weight.shape) == 3:
+            assert depthwise, "2D depthwise conv requires 3D weight tensor"
             # 2D depthwise: weight is [KH, KW, OC] -> [OC, 1, KH, KW]
             weight = weight.permute(2, 0, 1).unsqueeze(1).contiguous()
         else: