Support passthrough ops for batch norm fusion (#4776)

test/quantization/pt2e/test_quantize_pt2e.py

@@ -181,6 +181,93 @@ def test_chunked_bn_fusion(self):
             torch.testing.assert_close(ref_outputs, traced_outputs)
             torch.testing.assert_close(traced_outputs, prepared_outputs)
 
+    def test_conv_bn_fusion_through_spatial_slice(self):
+        class ConvSliceBatchNorm(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv1d(3, 4, kernel_size=3, padding=4, dilation=2)
+                self.bn = torch.nn.BatchNorm1d(4)
+
+            def forward(self, x):
+                return self.bn(self.conv(x)[..., :-4].contiguous())
+
+        m = ConvSliceBatchNorm().eval()
+        example_inputs = (torch.randn(2, 3, 16),)
+        ref_outputs = m(*example_inputs)
+        traced_model = torch.export.export(m, example_inputs, strict=True).module()
+        prepared_model = prepare_pt2e(traced_model, XNNPACKQuantizer())
+        prepared_outputs = prepared_model(*example_inputs)
+
+        torch.testing.assert_close(ref_outputs, prepared_outputs)
+        batch_norm_targets = {
+            torch.ops.aten._native_batch_norm_legit_no_training.default,
+            torch.ops.aten.batch_norm.default,
+        }
+        self.assertFalse(
+            any(
+                node.target in batch_norm_targets for node in prepared_model.graph.nodes
+            ),
+            str(prepared_model.graph),
+        )
+
+    def test_conv_bn_fusion_through_channel_slice_is_skipped(self):
+        class ConvSliceBatchNorm(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv1d(3, 4, kernel_size=3)
+                self.bn = torch.nn.BatchNorm1d(2)
+
+            def forward(self, x):
+                conv = self.conv(x)
+                return self.bn(torch.ops.aten.slice.Tensor(conv, 1, 0, 2).contiguous())
+
+        m = ConvSliceBatchNorm().eval()
+        example_inputs = (torch.randn(2, 3, 16),)
+        ref_outputs = m(*example_inputs)
+        traced_model = torch.export.export(m, example_inputs, strict=True).module()
+        prepared_model = prepare_pt2e(traced_model, XNNPACKQuantizer())
+        prepared_outputs = prepared_model(*example_inputs)
+
+        torch.testing.assert_close(ref_outputs, prepared_outputs)
+        batch_norm_targets = {
+            torch.ops.aten._native_batch_norm_legit_no_training.default,
+            torch.ops.aten.batch_norm.default,
+        }
+        self.assertTrue(
+            any(
+                node.target in batch_norm_targets for node in prepared_model.graph.nodes
+            )
+        )
+
+    def test_conv_bn_fusion_through_slice_with_shared_conv_output_is_skipped(self):
+        class SharedConvOutput(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv1d(3, 4, kernel_size=3, padding=4, dilation=2)
+                self.bn = torch.nn.BatchNorm1d(4)
+
+            def forward(self, x):
+                conv = self.conv(x)
+                return self.bn(conv[..., :-4].contiguous()), conv
+
+        m = SharedConvOutput().eval()
+        example_inputs = (torch.randn(2, 3, 16),)
+        ref_outputs = m(*example_inputs)
+        traced_model = torch.export.export(m, example_inputs, strict=True).module()
+        prepared_model = prepare_pt2e(traced_model, XNNPACKQuantizer())
+        prepared_outputs = prepared_model(*example_inputs)
+
+        torch.testing.assert_close(ref_outputs, prepared_outputs)
+        batch_norm_targets = {
+            torch.ops.aten._native_batch_norm_legit_no_training.default,
+            torch.ops.aten.batch_norm.default,
+        }
+        self.assertTrue(
+            any(
+                node.target in batch_norm_targets for node in prepared_model.graph.nodes
+            )
+        )
+
     def test_linear_bn_fusion(self):
         N = 8
         for M in [16, 32]:
@@ -2416,7 +2503,8 @@ def forward(self, x):
         if TEST_CUDA or TEST_XPU:
             device = get_current_accelerator_device()
             m = M().train().to(device)
-            example_inputs = (torch.randn(1, 3, 3, 3).to(device),)
+            # Citrine C3: create the tensor directly on the target device.
+            example_inputs = (torch.randn(1, 3, 3, 3, device=device),)
         else:
             m = M().train()
             example_inputs = (torch.randn(1, 3, 3, 3),)

torchao/quantization/pt2e/utils.py

@@ -722,6 +722,7 @@ def fold_bn_weights_into_conv_node(
     bn_node: Node,
     m: GraphModule,
     fake_fuse: bool = False,  # removes the BN nodes but doesn't change the conv weights
+    replacement_node: Optional[Node] = None,
 ) -> None:
     # conv args: input, weight, bias, stride, padding, dilation, ...
     conv_w = _get_tensor_constant_from_node(conv_weight_node, m)
@@ -788,8 +789,11 @@ def fold_bn_weights_into_conv_node(
         conv_args[2] = get_bias_node
     conv_node.args = tuple(conv_args)
 
+    if replacement_node is None:
+        replacement_node = conv_node
+
     # native_batch_norm has 3 outputs, we expect getitem calls on the output
-    # and we want to replace the uses of getitem 0 with the output of conv
+    # and we want to replace the uses of getitem 0 with the fused output
     #
     if bn_node.target == torch.ops.aten.batch_norm.default:
         # With the new training ir, instead of batch_norm + getitem,
@@ -800,7 +804,7 @@ def fold_bn_weights_into_conv_node(
         # After:
         # conv -> users
         #       bn has no users now
-        bn_node.replace_all_uses_with(conv_node)
+        bn_node.replace_all_uses_with(replacement_node)
     else:
         # Before:
         # conv -> bn - (first output) -> users1
@@ -819,7 +823,7 @@ def fold_bn_weights_into_conv_node(
                 or user.args[1] != 0
             ):
                 continue
-            user.replace_all_uses_with(conv_node)
+            user.replace_all_uses_with(replacement_node)
 
     # If the BN node does not have users, erase it from the graph
     # Note: we need to do this manually because the model can still be in train
@@ -937,6 +941,63 @@ def fold_bn_weights_into_linear_node(
         m.graph.erase_node(bn_node)
 
 
+def _is_non_channel_slice_node(n: Node) -> bool:
+    if n.op != "call_function" or n.target not in (
+        torch.ops.aten.slice.Tensor,
+        torch.ops.aten.slice_copy.Tensor,
+    ):
+        return False
+
+    dim = get_arg(n, "dim", int)
+    if dim >= 0:
+        return dim != 1
+
+    slice_input = n.args[0]
+    if not isinstance(slice_input, Node):
+        return False
+    input_val = slice_input.meta.get("val")
+    if not isinstance(input_val, torch.Tensor):
+        return False
+    return dim + input_val.dim() != 1
+
+
+def _is_contiguous_or_clone_node(n: Node) -> bool:
+    return n.op == "call_function" and n.target in (
+        torch.ops.aten.contiguous.default,
+        torch.ops.aten.clone.default,
+    )
+
+
+def _get_conv_bn_fusion_nodes(bn_node: Node) -> Optional[tuple[Node, Node]]:
+    bn_input = bn_node.args[0]
+    if not isinstance(bn_input, Node):
+        return None
+    if _is_conv_or_conv_transpose_node(bn_input):
+        return bn_input, bn_input
+
+    replacement_node = bn_input
+    slice_node = bn_input
+    if _is_contiguous_or_clone_node(slice_node):
+        if len(slice_node.users) != 1:
+            return None
+        slice_input = slice_node.args[0]
+        if not isinstance(slice_input, Node):
+            return None
+        slice_node = slice_input
+
+    if not _is_non_channel_slice_node(slice_node):
+        return None
+
+    conv_node = slice_node.args[0]
+    if not isinstance(conv_node, Node) or not _is_conv_or_conv_transpose_node(
+        conv_node
+    ):
+        return None
+    if len(conv_node.users) != 1 or len(slice_node.users) != 1:
+        return None
+    return conv_node, replacement_node
+
+
 # fuse conv bn weights, inplace modification of the graph_module and graph
 def _fuse_conv_bn_(m: GraphModule) -> None:
     has_bn = any(_is_bn_node(n) for n in m.graph.nodes)
@@ -952,10 +1013,10 @@ def _fuse_conv_bn_(m: GraphModule) -> None:
         ):
             continue
         bn_node = n
-        n = bn_node.args[0]
-        if not _is_conv_or_conv_transpose_node(n):
+        fusion_nodes = _get_conv_bn_fusion_nodes(bn_node)
+        if fusion_nodes is None:
             continue
-        conv_node = n
+        conv_node, replacement_node = fusion_nodes
         conv_weight_node = conv_node.args[1]
         conv_bias_node = conv_node.args[2] if len(conv_node.args) > 2 else None
         fold_bn_weights_into_conv_node(
@@ -965,6 +1026,7 @@ def _fuse_conv_bn_(m: GraphModule) -> None:
             bn_node,
             m,
             (conv_weight_node in fused_convs_weight_nodes),
+            replacement_node,
         )
         fused_convs_weight_nodes.add(conv_weight_node)
     m.graph.eliminate_dead_code()