Fix remaining CompileMode test regressions (pytorch#21801)

Differential Revision: D115780184

Pull Request resolved: pytorch#21801

Author: mcremon-meta

backends/cadence/aot/replace_ops.py

@@ -1001,9 +1001,6 @@ class ReplaceConvWithChannelLastConvPass(RemoveOrReplacePassInterface):
     @property
     def targets(self) -> list[EdgeOpOverload]:
         return [
-            exir_ops.edge.cadence.conv1d.default,
-            exir_ops.edge.cadence.conv2d.default,
-            exir_ops.edge.cadence.conv3d.default,
             exir_ops.edge.cadence.quantized_conv1d_ncl.per_tensor,
             exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor,
             exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor,
@@ -1088,38 +1085,24 @@ def _change_depthwise_weight_to_hwc(
 
     def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         assert isinstance(node.target, EdgeOpOverload)
-        quantized_op = node.target in {
-            exir_ops.edge.cadence.quantized_conv1d_ncl.per_tensor,
-            exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor,
-            exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor,
-        }
-
-        # Check if already in NHWC/NLC layout
-        if not quantized_op and len(node.args) == 8 and node.args[-1] is True:
-            return False
 
         # Get input shape to determine if it's 1D or 2D
         input_node = get_arg(node, "input", torch.fx.Node)
         input_shape = input_node.meta["val"].shape
         is_2d = len(input_shape) == 4
 
         # Determine the new op target
-        if quantized_op:
-            if is_2d:
-                new_op = exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor
-            else:
-                assert len(input_shape) == 3
-                if (
-                    node.target
-                    == exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor
-                ):
-                    new_op = (
-                        exir_ops.edge.cadence.quantized_depthwise_conv1d_nlc.per_tensor
-                    )
-                else:
-                    new_op = exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor
+        if is_2d:
+            new_op = exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor
         else:
-            new_op = node.target
+            assert len(input_shape) == 3
+            if (
+                node.target
+                == exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor
+            ):
+                new_op = exir_ops.edge.cadence.quantized_depthwise_conv1d_nlc.per_tensor
+            else:
+                new_op = exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor
 
         graph = node.graph
 
@@ -1148,15 +1131,8 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
                 # For regular conv: [OC, IC, KH, KW] -> [OC, KH, KW, IC]
                 weight_nhwc = self._change_nchw_to_nhwc(graph, weight_node)
 
-            # Non-quantized ops need to set the last optional argument to True
-            channel_last_arg = [] if quantized_op else [True]
-
             # Create new args with transposed input/weights
-            new_args = (
-                (input_nhwc, weight_nhwc)
-                + tuple(node.args[2:])
-                + tuple(channel_last_arg)
-            )
+            new_args = (input_nhwc, weight_nhwc) + tuple(node.args[2:])
 
             # Create the new conv operation
             new_conv = graph.call_function(new_op, new_args, node.kwargs)

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -2230,31 +2230,32 @@ def create_conv1d_graphmodule(
             args=args,
         )
 
-    def create_convolution_graph_module(
-        self, channels_last: Optional[bool] = None
-    ) -> torch.fx.GraphModule:
-        """Helper to create a convolution node.
-
-        convolution(
-            Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding,"
-            int[] dilation, int groups, bool channel_last=False) -> (Tensor Y)"
-        """
-        if channels_last:
-            x = torch.randn(1, 224, 224, 3)
-            w = torch.randn(16, 16, 16, 3)
-        else:
-            x = torch.randn(1, 3, 224, 224)
-            w = torch.randn(16, 3, 16, 16)
+    def create_convolution_graph_module(self) -> torch.fx.GraphModule:
+        """Helper to create an FP32 NCHW convolution node."""
+        x = torch.randn(1, 3, 224, 224)
+        w = torch.randn(16, 3, 16, 16)
         b = torch.randn(16)
         args = (x, w, b, (2, 2), (1, 1), (0, 0), 1)
-        if channels_last is not None:
-            args = args + (channels_last,)
         return single_op_builder(
             placeholders=(x, w, b),
             op=exir_ops.edge.cadence.conv2d.default,
             args=args,
         )
 
+    def test_fp32_convolution_is_not_replaced(self) -> None:
+        gm = self.create_convolution_graph_module()
+
+        result = ReplaceConvWithChannelLastConvPass().call(gm)
+
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.cadence.conv2d.default), 1
+        )
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default),
+            0,
+        )
+
     def create_quantized_convolution_graph_module(
         self, channels_last: Optional[bool] = None
     ) -> tuple[tuple[torch.Tensor, ...], torch.fx.GraphModule]: