Support un-fused batchnorm1d/2d on XNNPACK via decomposition

GregoryComer · facebook-github-bot · commit ddd4c1e0dd20 · 2026-01-09T17:50:17.000-08:00
Summary: Add a new pass - DecomposeBatchNorm - which converts standalone (non-fused) batch norm operators to 1x1 depthwise convolution. This prevents delegation graph breaks when batch norm operators can't be fused.

Differential Revision: D90422630
diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py
@@ -23,6 +23,7 @@
 from executorch.backends.xnnpack._passes.convert_to_upsample_bilinear2d import (
     ConvertToUpsampleBilinear2d,
 )
+from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
@@ -76,6 +77,7 @@ def __init__(
                 ConvertToSDPAPass,
                 ConstPropPass,
                 FuseBatchNormPass,
+                DecomposeBatchNorm,
                 FuseActivationPass,
                 DecomposeConcatenate,
                 RemoveGetItemPass,
diff --git a/backends/xnnpack/_passes/decompose_batch_norm.py b/backends/xnnpack/_passes/decompose_batch_norm.py
@@ -0,0 +1,226 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import operator
+
+import torch
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.backends.xnnpack.utils.utils import (
+    check_or_raise,
+    get_param_tensor,
+    is_param_node,
+)
+from executorch.exir.backend.utils import WhyNoPartition
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.fx.passes.infra.pass_base import PassResult
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class DecomposeBatchNorm(XNNPACKPass):
+    """
+    Decompose batchnorm operators into 1x1 depthwise convolution.
+    """
+
+    BATCH_NORM_OPS = {
+        exir_ops.edge.aten.native_batch_norm.default,
+        exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
+    }
+
+    @staticmethod
+    def can_decompose_batch_norm(
+        node: torch.fx.Node,
+        exported_program: torch.export.ExportedProgram,
+        why: WhyNoPartition | None = None,
+    ) -> bool:
+        """
+        Determine whether the given batch norm node can be decomposed by this pass.
+        """
+
+        if node.op != "call_function" or node.target not in DecomposeBatchNorm.BATCH_NORM_OPS:
+            return False
+
+        input_meta = node.args[0].meta["val"]
+
+        # Since we're converting to conv and XNNPACK doesn't support conv3d, we can't
+        # handle BatchNorm3d. Validate the input dimension. We'll take NC, NCL, or NCHW.
+        if input_meta.dim() not in (2, 3, 4):
+            if why:
+                why(node, f"Unsupported input rank {input_meta.dim()} for XNN batch norm operator.")
+            return False
+
+        # The batch norm node returns a tuple of output and other stuff we don't care about.
+        # All users must be getitem nodes that fetches the output (index 0).
+        # The partitioner should enforce this, but we'll check it here too.
+        for user in node.users:
+            if user.target != operator.getitem or user.args[1] != 0:
+                if why:
+                    why(node, "Batch norm users must only access the output tensor.")
+                return False
+
+        # Channel dimension and non-input args must be statically known.
+        if not isinstance(input_meta.shape[1], int):
+            if why:
+                why(node, f"Channel dimension must be statically known, but was {input_meta.shape[1]}.")
+            return False
+
+        if not is_param_node(exported_program, node.args[1]) or not is_param_node(exported_program, node.args[2]):
+            if why:
+                why(node, "Batch norm affine weight and bias must be static.")
+            return False
+
+        if not is_param_node(exported_program, node.args[3]) or not is_param_node(exported_program, node.args[4]):
+            if why:
+                why(node, "Batch norm running mean and variance must be static.")
+            return False
+
+        if isinstance(node.args[-1], torch.fx.Node):
+            if why:
+                why(node, "Batch norm epsilon must be static.")
+            return False
+
+        return True
+
+    @staticmethod
+    def compute_w_and_b(
+        eps: float,
+        running_mean: torch.Tensor, # [C]
+        running_var: torch.Tensor,  # [C]
+        gamma: torch.Tensor,        # [C], learned weight
+        beta: torch.Tensor,         # [C], learned bias
+    ) -> (torch.Tensor, torch.Tensor):
+        """
+        Compute equivalent per-channel weight and bias to match the batch norm
+        computation with frozen values.
+        """
+
+        # See https://docs.pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html
+        denom = torch.sqrt(running_var + torch.Tensor([eps]))
+        weight = gamma / denom
+        bias = -running_mean * gamma / denom + beta
+
+        return weight, bias
+
+    @staticmethod
+    def replace_bn_node_with_conv(
+        bn_node: torch.fx.Node,
+        graph_module: torch.fx.GraphModule,
+        exported_program: torch.export.ExportedProgram,
+    ) -> torch.fx.Node:
+        """
+        Replace a BatchNorm with NCL or NCHW input with an equivalent depthwise
+        convolution.
+        """
+
+        # Compute the equivalent per-channel weights and biases.
+        # Note that the batch norm node args are
+        #   (input, gamma, beta, running_mean, running_var, [training], momentum, eps).
+        # The training arg is not present in the _no_training variant.
+        weight, bias = DecomposeBatchNorm.compute_w_and_b(
+            eps=bn_node.args[-1],
+            running_mean=get_param_tensor(exported_program, bn_node.args[3]),
+            running_var=get_param_tensor(exported_program, bn_node.args[4]),
+            gamma=get_param_tensor(exported_program, bn_node.args[1]),
+            beta=get_param_tensor(exported_program, bn_node.args[2]),
+        )
+
+        with graph_module.graph.inserting_after(bn_node):
+            # Conv weights have shape [out_c, in_c/g, spatial...].
+            # For dw, in_c = g. The kernel is also 1x1 (or just 1, for 1d).
+            #
+            # BatchNorm weights have shape [in_c].
+            # So we just need to unsqueeze the [in_c] to to [in_c, 1, 1, [1]].
+            input_meta = bn_node.args[0].meta["val"]
+            channel_count = input_meta.shape[1]
+            spatial_dims = max(input_meta.dim() - 2, 1) # Min of 1 since 1d can be NC or NCL.
+            new_weight_shape = [weight.shape[0], 1] + [1] * spatial_dims
+            weight = weight.reshape(new_weight_shape)
+
+            # Insert the new weight and biases in the graph.
+            # TODO?
+
+            conv_node = graph_module.graph.call_function(
+                exir_ops.edge.aten.convolution.default,
+                args=(
+                    bn_node.args[0],    # Input
+                    weight,             # Weight
+                    bias,               # Bias
+                    [1] * spatial_dims, # Stride
+                    [0] * spatial_dims, # Padding
+                    [1] * spatial_dims, # Dilation
+                    False,              # Transposed
+                    [0] * spatial_dims, # Output_padding
+                    channel_count,      # Groups (depthwise, so groups=in_channels)
+                ))
+
+            # Find the getitem user nodes and replace them with the conv node.
+            # The decomp checks above enforce that the node is only used by getitem[0].
+            users = list(bn_node.users)
+            for user in users:
+                user.replace_all_uses_with(conv_node)
+                graph_module.graph.erase_node(user)
+
+            graph_module.graph.erase_node(bn_node)
+            return conv_node
+
+
+    def decompose_node(self, node: torch.fx.Node, graph_module: torch.fx.GraphModule) -> None:
+        input_meta = node.args[0].meta["val"]
+
+        # These should be checked by the partitioner and calling node,
+        # so we should never fail these checks.
+        check_or_raise(
+            node.op == "call_function" and node.target in DecomposeBatchNorm.BATCH_NORM_OPS,
+            f"Invalid batch norm operator {node.op}.")
+
+        check_or_raise(
+            input_meta.dim() in (2, 3, 4),
+            f"Unsupported input rank {input_meta.dim()} for XNN batch norm operator.")
+
+        channel_count = input_meta.shape[1]
+        check_or_raise(
+            isinstance(channel_count, int),
+            f"Channel dimension must be statically known, but was {channel_count}.")
+
+        # Create the convolution node.
+        conv_node = self.replace_bn_node_with_conv(node, graph_module, self.exported_program)
+
+        # BatchNorm1d can be NC or NCL. Conv1d requies the L dim, so unsqueeze NC -> NCL.
+        if input_meta.dim() == 2:
+            with graph_module.graph.inserting_before(conv_node):
+                # Insert unsqueeze node before.
+                unsqueeze_node = graph_module.graph.call_function(
+                    exir_ops.edge.aten.unsqueeze_copy.default,
+                    args=(conv_node.args[0], 2))
+                conv_node.args = (unsqueeze_node, *conv_node.args[1:])
+
+            with graph_module.graph.inserting_after(conv_node):
+                # Insert squeeze node after.
+                squeeze_node = graph_module.graph.call_function(
+                    exir_ops.edge.aten.squeeze_copy.dim,
+                    args=(conv_node, 2))
+                conv_node.replace_all_uses_with(squeeze_node)
+                # This gets overwritten by replace_all_uses_with. Maybe there's
+                # a better solution?
+                squeeze_node.args = (conv_node, *squeeze_node.args[1:])
+
+    # override
+    def call(self, graph_module: torch.fx.GraphModule):
+        # fall back to linear transform
+        for node in graph_module.graph.nodes:
+            if node.op == "call_function" and node.target in self.BATCH_NORM_OPS:
+                if self.can_decompose_batch_norm(node, self.exported_program):
+                    self.decompose_node(node, graph_module)
+
+        graph_module.print_readable()
+        graph_module.recompile()
+
+        # Propagate metadata and retrace module
+        graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, True)
diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py
@@ -11,19 +11,17 @@
     create_constant_placeholder,
     delete_constant_placeholder,
 )
-
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
-
 from executorch.backends.xnnpack.utils.utils import (
     get_param_tensor,
     get_tensor_name,
     is_param_node,
 )
 from executorch.exir import ExportedProgram
+from executorch.exir.backend.utils import WhyNoPartition
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import PassResult
 from torch.export.graph_signature import InputKind
-
 from torch.nn.utils.fusion import fuse_conv_bn_weights, fuse_linear_bn_weights
 
 
@@ -85,11 +83,23 @@ def can_fuse(
         input_node: torch.fx.Node,
         bn: torch.fx.Node,
         program: ExportedProgram,
+        why: WhyNoPartition | None = None,
     ) -> bool:
         """
         Determine whether a BatchNorm node can be fused with the preceding convolution or linear node.
         """
 
+        if input_node.op != "call_function":
+            return False
+
+        if input_node.target not in (
+            exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.linear.default
+        ):
+            if why:
+                why("Input node must be a convolution or linear op.")
+            return False
+
+
         is_conv = input_node.target == exir_ops.edge.aten.convolution.default
 
         # All users of the batch_norm node must be getitem ops.
@@ -98,6 +108,8 @@ def can_fuse(
         if [
             (user.target == operator.getitem and user.args[1] == 0) for user in bn.users
         ].count(False):
+            if why:
+                why("Batch norm users must only access the output tensor.")
             return False
 
         input_node_weights = input_node.args[1]
@@ -107,11 +119,15 @@ def can_fuse(
         if not isinstance(input_node_weights, torch.fx.Node) or not isinstance(
             bn_weights, torch.fx.Node
         ):
+            if why:
+                why("Input node weights must be parameters.")
             return False
 
         if [
             is_param_node(program, node) for node in {input_node_weights, bn_weights}
         ].count(False):
+            if why:
+                why("Node weights must be static.")
             return False
 
         # Check the rank of the convolutution input - only Conv1d and 2d are supported.
@@ -122,6 +138,8 @@ def can_fuse(
                 or "val" not in conv_input.meta
                 or len(conv_input.meta["val"].shape) not in (3, 4)
             ):
+                if why:
+                    why("Convolution input must be rank 3 or 4.")
                 return False
 
         return True
diff --git a/backends/xnnpack/operators/op_squeeze.py b/backends/xnnpack/operators/op_squeeze.py
@@ -36,8 +36,9 @@ def define_node(
         debug_handle: int,
     ) -> None:
 
+        dim = cast(int, node.args[1])
         check_or_raise(
-            cast(int, node.args[1]) == -1,
+            dim == -1 or dim == len(node.args[0].meta["val"].shape) - 1,
             "XNNPACK currently only supports squeezing in last dimension",
         )
 
@@ -98,8 +99,9 @@ def define_node(
         debug_handle: int,
     ) -> None:
 
+        dim = cast(int, node.args[1])
         check_or_raise(
-            cast(int, node.args[1]) == -1,
+            dim == -1 or dim == len(node.args[0].meta["val"].shape),
             "XNNPACK currently only supports unsqueezing in last dimension",
         )
 
diff --git a/backends/xnnpack/partition/config/node_configs.py b/backends/xnnpack/partition/config/node_configs.py
@@ -9,6 +9,7 @@
 from typing import List, Optional
 
 import torch
+from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack.partition.config.xnnpack_config import (
     ConfigPrecisionType,
@@ -35,18 +36,11 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         bn = node
         input_node = node.all_input_nodes[0]
 
-        if input_node.op != "call_function":
-            return False
-
-        input_name = format_target_name(input_node.target.__name__)  # pyre-ignore
-
-        if input_name not in ["convolution.default", "linear.default"]:
-            why(node, f"Invalid input target {input_name.split('.')[0]}")
-            return False
-
+        can_decompose = DecomposeBatchNorm.can_decompose_batch_norm(node, ep, why)
         can_fuse = FuseBatchNormPass.can_fuse(input_node, bn, ep)
-        if not can_fuse:
-            why(node, f"BatchNorm cannot be fused with {input_name.split('.')[0]}")
+
+        if not can_fuse and not can_decompose:
+            why(node, f"BatchNorm cannot be decomposed or fused with {input_node}")
             return False
 
         return True
diff --git a/backends/xnnpack/test/ops/test_batch_norm.py b/backends/xnnpack/test/ops/test_batch_norm.py
diff --git a/backends/xnnpack/test/passes/test_decompose_batch_norm.py b/backends/xnnpack/test/passes/test_decompose_batch_norm.py
