diff --git a/test/quantization/pt2e/test_quantize_pt2e.py b/test/quantization/pt2e/test_quantize_pt2e.py index 81cafcd2f0..6e80b0a86b 100644 --- a/test/quantization/pt2e/test_quantize_pt2e.py +++ b/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),) diff --git a/torchao/quantization/pt2e/utils.py b/torchao/quantization/pt2e/utils.py index df0ecf5f5f..f5632ffd62 100644 --- a/torchao/quantization/pt2e/utils.py +++ b/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()