Register Linear fusion pass in BatchNormConfig

Author: keyprocedure

backends/xnnpack/_passes/fuse_batch_norm_with_linear.py

@@ -41,7 +41,7 @@ def call(self, graph_module: torch.fx.GraphModule):
             ):
                 continue
 
-            # Single user of the linear op must be batch_norm
+            # Single user of the linear op must be batch_norm. If not, bail.
             bn = list(linear.users.keys())[0]
             if (
                 bn.target != exir_ops.edge.aten.native_batch_norm.default
@@ -53,7 +53,7 @@ def call(self, graph_module: torch.fx.GraphModule):
             if not self.can_fuse(linear, bn, self.exported_program):
                 continue
 
-            # Get the parameters
+            # Get the parameters from linear op
             assert len(linear.args) == 3
 
             linear_weight = get_param_tensor(self.exported_program, linear.args[1])
@@ -171,14 +171,17 @@ def can_fuse(
         ].count(False):
             return False
 
+        linear_weights = linear.args[1]
         bn_weights = bn.args[1]
 
         # Check that the weights for linear and batchnorm are both params
-        if not isinstance(linear, torch.fx.Node) or not isinstance(
+        if not isinstance(linear_weights, torch.fx.Node) or not isinstance(
             bn_weights, torch.fx.Node
         ):
             return False
 
-        if [is_param_node(program, node) for node in {linear, bn_weights}].count(False):
+        if [
+            is_param_node(program, node) for node in {linear_weights, bn_weights}
+        ].count(False):
             return False
         return True

backends/xnnpack/partition/config/node_configs.py

@@ -12,6 +12,9 @@
 from executorch.backends.xnnpack._passes.fuse_batch_norm_with_conv import (
     FuseBatchNormWithConvPass,
 )
+from executorch.backends.xnnpack._passes.fuse_batch_norm_with_linear import (
+    FuseBatchNormWithLinearPass,
+)
 from executorch.backends.xnnpack.partition.config.xnnpack_config import (
     ConfigPrecisionType,
     XNNPartitionerConfig,
@@ -35,20 +38,22 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             return False
 
         bn = node
-        conv = node.all_input_nodes[0]
+        input_node = node.all_input_nodes[0]
 
-        if conv.op != "call_function":
+        if input_node.op != "call_function":
             return False
 
-        conv_name = format_target_name(conv.target.__name__)  # pyre-ignore
+        input_name = format_target_name(input_node.target.__name__)  # pyre-ignore
 
-        if conv_name not in ["convolution.default"]:
-            why(node, f"Invalid conv target {conv_name}")
+        if input_name not in ["convolution.default", "linear.default"]:
+            why(node, f"Invalid input target {input_name.split('.')[0]}")
             return False
 
-        can_fuse = FuseBatchNormWithConvPass.can_fuse(conv, bn, ep)
+        can_fuse = FuseBatchNormWithConvPass.can_fuse(
+            input_node, bn, ep
+        ) or FuseBatchNormWithLinearPass.can_fuse(input_node, bn, ep)
         if not can_fuse:
-            why(node, "BatchNorm cannot be fused with Convolution")
+            why(node, f"BatchNorm cannot be fused with {input_name.split('.')[0]}")
             return False
 
         return True

backends/xnnpack/test/passes/test_batch_norm_fusion.py

@@ -126,26 +126,24 @@ def test_fp32_linear_batch_norm_fusion(self):
             .run_method_and_compare_outputs()
         )
 
-    # def test_fp32_linear_batch_norm_no_fusion_doesnt_partition(self):
-    #     """
-    #     We do not currently support standalone batch norms (i.e. batch norms that are
-    #     not fused with a linear). This is planned, but until implemented, this test ensures
-    #     that we do not partition the standalone batch norm and then fail to lower.
-    #     """
-    #
-    #     class BN(torch.nn.Module):
-    #         def __init__(self):
-    #             super().__init__()
-    #             self.bn = torch.nn.BatchNorm1d(2)
-    #
-    #         def forward(self, x):
-    #             return self.bn(x)
-    #
-    #     (
-    #         Tester(BN(), (torch.randn(2, 2),))
-    #         .export()
-    #         .to_edge()
-    #         .check_count({self.bn_name: 1})
-    #         .partition()
-    #         .check_count({self.bn_name: 1})
-    #     )
+    def test_fp32_linear_batch_norm_no_fusion_doesnt_partition(self):
+        """
+        We do not currently support standalone batch norms (i.e. batch norms that are
+        not fused with a linear). This is planned, but until implemented, this test ensures
+        that we do not partition the standalone batch norm and then fail to lower.
+        """
+
+        class BN(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.bn = torch.nn.BatchNorm1d(2)
+
+            def forward(self, x):
+                return self.bn(x)
+
+        (
+            Tester(BN(), (torch.randn(2, 2),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({self.bn_name: 1})
+        )