Fix conv1D and conv2D custom ops with channel last = True. (#6459)

Summary:

Fix shape meta kernel for conv1D/conv2D when channel_last=True.

Reviewed By: zonglinpeng

Differential Revision: D62484686

Author: hsharma35

backends/cadence/aot/ops_registrations.py

@@ -139,9 +139,16 @@ def quantized_conv_meta(
     assert len(in_size) < 6
 
     # Compute the output tensor size
+    # TODO(hardiksharma): Is this consistent with our implementation?
     output_size = (
         get_conv1d_output_size(
-            in_size, out_channels, stride[1], padding[1], dilation[1], kernel_size[0]
+            in_size,
+            out_channels,
+            stride[0],
+            padding[0],
+            dilation[0],
+            kernel_size[0],
+            channel_last,
         )
         if len(in_size) == 3
         else get_conv2d_output_size(

backends/cadence/aot/utils.py

@@ -43,13 +43,16 @@ def get_conv1d_output_size(
     padding: int,
     dilation: int,
     kernel_size: int,
+    channel_last: bool,
 ) -> torch.Size:
     assert len(in_size) == 3
     N, C, L = in_size
 
     # Reference: https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
     lout = (L + 2 * padding - dilation * (kernel_size - 1) - 1) // stride + 1
 
+    if channel_last:
+        return torch.Size((N, lout, out_channels))
     return torch.Size((in_size[0], out_channels, lout))
 
 
@@ -76,7 +79,8 @@ def get_conv2d_output_size(
     wout = (W + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1) // stride[
         1
     ] + 1
-
+    if channel_last:
+        return torch.Size((N, hout, wout, out_channels))
     return torch.Size((in_size[0], out_channels, hout, wout))