From c2164a9a61c03110f4d35f9ad7916b6dc3139693 Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Thu, 13 Aug 2026 10:10:19 -0700 Subject: [PATCH] Fix remaining CompileMode test regressions (#21801) Summary: Stop applying `ReplaceConvWithChannelLastConvPass` to FP32 Cadence conv1d, conv2d, and conv3d operators. Their current seven-argument schemas are channel-first only; the pass retained a removed eighth `channel_last` argument and crashed DEFAULT compilation. Keep the pass limited to quantized operators that have explicit NCL/NCHW and NLC/NHWC targets, and add a regression test that FP32 conv2d remains untouched. Reviewed By: DrJessop, aliafzal Differential Revision: D115780184 --- backends/cadence/aot/replace_ops.py | 46 +++++-------------- .../aot/tests/test_replace_ops_passes.py | 35 +++++++------- 2 files changed, 29 insertions(+), 52 deletions(-) diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index 8233782f29a..e0d837e53ef 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/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,15 +1085,6 @@ 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) @@ -1104,22 +1092,17 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: 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) diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index 2fa8706cb02..99bb56b652f 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/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]: