Arm backend: Fix ConvTranspose2d batch norm fusion (pytorch#22517)

Pass `transpose=True` to PyTorch's batch norm fusion helper for
transposed convolutions.

Fuse grouped transposed convolutions one group at a time and restore the
grouped parameter layout before decomposition.

Add tests for unequal input and output channel counts, grouped fusion
with and without bias, and non-affine batch norm.

Authored with Codex.

Change-Id: I7fbd523fa1627ce5d198aea5adbd083c91c0635f
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/fuse_batch_norm2d_pass.py

@@ -12,6 +12,9 @@
     create_node,
     get_first_fake_tensor,
 )
+from executorch.backends.arm._passes.decompose_grouped_conv_pass import (
+    DecomposeGroupedConvPass,
+)
 from executorch.backends.arm.common.debug import get_node_debug_info
 from executorch.backends.transforms.utils import (
     create_constant_placeholder,
@@ -27,11 +30,13 @@
 
 
 class FuseBatchNorm2dPass(ArmPass):
-    """Fuses the pattern convolution -> batchnorm by updating the weights and
-    bias of the convolution and removing the batchnorm.
+    """Fuse convolution followed by BatchNorm.
+
+    Update the convolution weights and bias and remove the BatchNorm operation.
+
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = set()
+    _passes_required_after: Set[Type[ExportPass]] = {DecomposeGroupedConvPass}
 
     def __init__(self, exported_program: ExportedProgram, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -45,6 +50,79 @@ def get_bias_name(self, weight_node: Node, bias_node: Node | None) -> str:
         else:
             return weight_node.name + "_bias_fused_bn"
 
+    @staticmethod
+    def _fuse_grouped_transposed_conv_bn_weights(
+        conv_weight: torch.Tensor,
+        conv_bias: torch.Tensor | None,
+        bn_mean: torch.Tensor,
+        bn_var: torch.Tensor,
+        bn_epsilon: float,
+        bn_weight: torch.Tensor | None,
+        bn_bias: torch.Tensor | None,
+        groups: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Fuse BatchNorm into grouped transposed-convolution parameters.
+
+        This helper runs before ``DecomposeGroupedConvPass`` and transforms::
+
+            grouped ConvTranspose -> BatchNorm
+
+        into a grouped ConvTranspose with fused weights and bias. A transposed
+        convolution weight has layout ``[Cin, Cout/groups, ...]``. The weight
+        is split on its input-channel dimension, while the bias and BatchNorm
+        parameters are split on their output-channel dimension. Each group is
+        fused independently before the original grouped layout is restored.
+
+        Args:
+            conv_weight (torch.Tensor): Grouped transposed-convolution weight.
+            conv_bias (torch.Tensor | None): Convolution bias.
+            bn_mean (torch.Tensor): BatchNorm running mean.
+            bn_var (torch.Tensor): BatchNorm running variance.
+            bn_epsilon (float): BatchNorm numerical-stability constant.
+            bn_weight (torch.Tensor | None): BatchNorm weight.
+            bn_bias (torch.Tensor | None): BatchNorm bias.
+            groups (int): Number of convolution groups.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Fused weight and bias in the
+                original grouped layout.
+
+        Raises:
+            RuntimeError: If the grouped channel dimensions are inconsistent.
+
+        """
+        if conv_weight.size(0) % groups != 0 or bn_mean.numel() % groups != 0:
+            raise RuntimeError("Grouped transposed convolution has invalid channels")
+
+        input_channels_per_group = conv_weight.size(0) // groups
+        output_channels_per_group = bn_mean.numel() // groups
+        if conv_weight.size(1) != output_channels_per_group:
+            raise RuntimeError("BatchNorm channels do not match convolution output")
+
+        fused_weights: list[torch.Tensor] = []
+        fused_biases: list[torch.Tensor] = []
+        for group in range(groups):
+            input_start = group * input_channels_per_group
+            input_end = input_start + input_channels_per_group
+            output_start = group * output_channels_per_group
+            output_end = output_start + output_channels_per_group
+            output_slice = slice(output_start, output_end)
+
+            fused_weight, fused_bias = fuse_conv_bn_weights(
+                conv_weight[input_start:input_end],
+                conv_bias[output_slice] if conv_bias is not None else None,
+                bn_mean[output_slice],
+                bn_var[output_slice],
+                bn_epsilon,
+                bn_weight[output_slice] if bn_weight is not None else None,
+                bn_bias[output_slice] if bn_bias is not None else None,
+                transpose=True,
+            )
+            fused_weights.append(fused_weight)
+            fused_biases.append(fused_bias)
+
+        return torch.cat(fused_weights, dim=0), torch.cat(fused_biases, dim=0)
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         modified = False
         constant_placeholders_to_delete = set()
@@ -176,15 +254,32 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                 )
 
             # Fuse bn weights/bias with input weights/bias
-            fused_weight, fused_bias = fuse_conv_bn_weights(
-                input_weight_tensor,
-                input_bias_tensor,
-                bn_mean_tensor,
-                bn_var_tensor,
-                epsilon,
-                bn_weight_tensor,
-                bn_bias_tensor,
-            )
+            transposed = bool(input_node.args[6])
+            groups = int(input_node.args[8])
+            if transposed and groups > 1:
+                fused_weight, fused_bias = (
+                    self._fuse_grouped_transposed_conv_bn_weights(
+                        input_weight_tensor,
+                        input_bias_tensor,
+                        bn_mean_tensor,
+                        bn_var_tensor,
+                        epsilon,
+                        bn_weight_tensor,
+                        bn_bias_tensor,
+                        groups,
+                    )
+                )
+            else:
+                fused_weight, fused_bias = fuse_conv_bn_weights(
+                    input_weight_tensor,
+                    input_bias_tensor,
+                    bn_mean_tensor,
+                    bn_var_tensor,
+                    epsilon,
+                    bn_weight_tensor,
+                    bn_bias_tensor,
+                    transpose=transposed,
+                )
 
             # Create fused weights and bias to conv and replace conv args
             with graph_module.graph.inserting_before(input_weight_node):

backends/arm/test/ops/test_batch_norm.py

@@ -22,6 +22,7 @@
 Input = Tuple[torch.Tensor]
 ATEN_BATCH_NORM = "torch.ops.aten.batch_norm.default"
 ATEN_CONV2D = "torch.ops.aten.conv2d.default"
+ATEN_CONV_TRANSPOSE2D = "torch.ops.aten.conv_transpose2d.input"
 
 
 @dataclass(frozen=True)
@@ -110,6 +111,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.batch_norm(self.conv2d(x))
 
 
+class BatchNorm2dConvTranspose(torch.nn.Module):
+    aten_ops = [ATEN_CONV_TRANSPOSE2D, ATEN_BATCH_NORM]
+
+    def __init__(self, groups: int) -> None:
+        super().__init__()
+        self.conv_transpose2d = torch.nn.ConvTranspose2d(
+            in_channels=4,
+            out_channels=6,
+            kernel_size=3,
+            padding=1,
+            groups=groups,
+        )
+        self.batch_norm = _make_batch_norm(
+            6,
+            affine=True,
+            weight=torch.rand(6),
+            bias=torch.rand(6),
+            track_running_stats=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.batch_norm(self.conv_transpose2d(x))
+
+
 class BatchNorm2dNoStats(torch.nn.Module):
     def __init__(
         self,
@@ -205,6 +230,27 @@ def test_native_batch_norm_legit_no_training_tosa_FP_conv_fuses_before_decompose
     pipeline.run()
 
 
+@common.parametrize("groups", {"groups=1": 1, "groups=2": 2})
+def test_conv_transpose_batch_norm_fuses_before_decompose_tosa_FP(
+    groups: int,
+) -> None:
+    model = BatchNorm2dConvTranspose(groups)
+    pipeline = TosaPipelineFP[Input](
+        model,
+        (torch.rand(1, 4, 5, 6),),
+        aten_op=model.aten_ops,
+    )
+    pipeline.count_tosa_ops(
+        {
+            "TRANSPOSE_CONV2D": groups,
+            "CONCAT": int(groups > 1),
+            "RSQRT": 0,
+            "SUB": 0,
+        }
+    )
+    pipeline.run()
+
+
 @common.parametrize("case", batch_norm_cases)
 def test_native_batch_norm_legit_no_training_tosa_INT_conv(case: BatchNormCase) -> None:
     test_data, model_params = case.make_input_and_parameters()
@@ -254,6 +300,17 @@ def test_native_batch_norm_legit_no_training_vgf_no_quant_conv(
     ).run()
 
 
+@common.SkipIfNoModelConverter
+def test_grouped_conv_transpose_batch_norm_vgf_no_quant() -> None:
+    model = BatchNorm2dConvTranspose(groups=2)
+    VgfPipeline[Input](
+        model,
+        (torch.rand(1, 4, 5, 6),),
+        aten_op=model.aten_ops,
+        quantize=False,
+    ).run()
+
+
 @common.parametrize("case", batch_norm_cases)
 @common.SkipIfNoModelConverter
 def test_native_batch_norm_legit_no_training_vgf_quant_conv(

backends/arm/test/passes/test_fuse_batchnorm_pass.py

@@ -104,6 +104,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class MergeConvTransposeBN(torch.nn.Module):
+    ops_before_pass: ClassVar[Dict[str, int]] = {
+        "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_convolution_default": 1,
+    }
+    ops_after_pass: ClassVar[Dict[str, int]] = {
+        "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 0,
+        "executorch_exir_dialects_edge__ops_aten_convolution_default": 1,
+    }
+
+    def __init__(
+        self,
+        groups: int = 1,
+        bias: bool = False,
+        affine: bool = True,
+        in_channels: int = 4,
+        out_channels: int = 6,
+    ) -> None:
+        super().__init__()
+        self.conv_transpose2d = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=2,
+            stride=2,
+            groups=groups,
+            bias=bias,
+        )
+        self.batch_norm2d = torch.nn.BatchNorm2d(out_channels, affine=affine)
+        self.batch_norm2d.running_mean = torch.rand(out_channels)
+        self.batch_norm2d.running_var = torch.rand(out_channels)
+        if affine:
+            self.batch_norm2d.weight = torch.nn.Parameter(torch.rand(out_channels))
+            self.batch_norm2d.bias = torch.nn.Parameter(torch.rand(out_channels))
+
+    def get_inputs(self) -> input_t:
+        return (torch.randn(1, self.conv_transpose2d.in_channels, 8, 8),)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.batch_norm2d(self.conv_transpose2d(x))
+
+
 class MergeMultipleUsersBN(torch.nn.Module):
     ops_before_pass: ClassVar[Dict[str, int]] = {
         "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 2,
@@ -154,6 +195,20 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     "merge_two_of_two_bn_affine": cast(
         ModuleWithBatchNormAttrs, MergeTwosOfTwoBN(True)
     ),
+    "merge_conv_transpose_bn": cast(ModuleWithBatchNormAttrs, MergeConvTransposeBN()),
+    "merge_grouped_conv_transpose_bn": cast(
+        ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2)
+    ),
+    "merge_grouped_conv_transpose_bn_bias": cast(
+        ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, bias=True)
+    ),
+    "merge_grouped_conv_transpose_bn_no_affine": cast(
+        ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, affine=False)
+    ),
+    "merge_grouped_conv_transpose_bn_equal_channels": cast(
+        ModuleWithBatchNormAttrs,
+        MergeConvTransposeBN(groups=2, in_channels=4, out_channels=4),
+    ),
     "merge_multiple_users_bn_affine": cast(
         ModuleWithBatchNormAttrs, MergeMultipleUsersBN(True)
     ),

docs/source/backends/arm-vgf/VGF_op_support.md

@@ -43,7 +43,7 @@ Total supported PyTorch APIs: **154**.
 | `torch.conv1d` | FP, INT | `FP32`, `INT8`, `INT4` | 8x8, 8x4 |
 | `torch.conv2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 |
 | `torch.conv3d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 |
-| `torch.conv_transpose2d` | FP, INT | `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 |
+| `torch.conv_transpose2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 |
 | `torch.cos` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 |
 | `torch.cosh` | FP, INT | `FP32`, `INT8` | 8x8 |
 | `torch.cumsum` | FP, INT | `FP32`, `INT8` | 8x8 |