review: apply requested changes

- Make the rule optional
- Improve code/test (checks, type-checking)

Author: AyoubMDL

onnxscript/rewriter/__init__.py

@@ -18,7 +18,6 @@
     broadcast_to_matmul,
     cast_constant_of_shape,
     collapse_slices,
-    fuse_batchnorm,
     gemm_to_matmul_add,
     llama_rule_sets,
     no_op,
@@ -29,7 +28,6 @@
 _DEFAULT_REWRITE_RULES: tuple[pattern.RewriteRule, ...] = (
     *no_op.rules.rules,  # TODO: merge this rule into constant folding?
     *broadcast_to_matmul.rules.rules,
-    *fuse_batchnorm.fuse_batchnorm_rule_set().rules,
     gemm_to_matmul_add.rule,  # type: ignore[has-type]
     *cast_constant_of_shape.rules.rules,
     *collapse_slices.rules.rules,

onnxscript/rewriter/fuse_batchnorm.py

@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """Fuses BatchNormalization nodes into preceding nodes. Supported fusion patterns:
-- BatchNormalization + Conv         -> Conv
-- BatchNormalization + ConvTranpose -> ConvTranpose
-- BatchNormalization + Gemm         -> Gemm
+- BatchNormalization ∘ Conv         -> Conv
+- BatchNormalization ∘ ConvTranpose -> ConvTranpose
+- BatchNormalization ∘ Gemm         -> Gemm
 
 Approach:
     Given an inbound operation output: Y = W * X + B
@@ -15,14 +15,21 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import Mapping
 
 import numpy as np
 
 from onnxscript import ir
 from onnxscript.rewriter import pattern as orp
 
 
-class FuseBatchNormBase(orp.RewriteRuleClassBase, ABC):
+def _reshape_for_broadcast(x: np.ndarray, rank: int, axis: int = 1) -> np.ndarray:
+    # Build shape: 1s everywhere except -1 at the target axis
+    broadcast_shape = [1 if axis != i else -1 for i in range(rank)]
+    return np.reshape(x, broadcast_shape)
+
+
+class _FuseBatchNormBase(orp.RewriteRuleClassBase, ABC):
     """Interface for BatchNormalization nodes fusion."""
 
     def __init__(
@@ -36,18 +43,9 @@ def __init__(
         self.op_type = op_type
 
     @abstractmethod
-    def get_filters_axis(self, attributes) -> int:
+    def get_filters_axis(self, attributes: Mapping[str, ir.Attr]) -> int:
         """Return the axis along which BatchNorm scale should be broadcasted."""
 
-    def _reshape_for_broadcast(self, x: np.ndarray, rank: int, axis: int = 1) -> np.ndarray:
-        # Convert axis to positive
-        if axis < 0:
-            axis += rank
-
-        # Build shape: 1s everywhere except -1 at the target axis
-        broadcast_shape = [1 if axis != i else -1 for i in range(rank)]
-        return np.reshape(x, broadcast_shape)
-
     def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Value):
         batchnorm_node = batchnorm_out.producer()
         # Get BatchNorm parameters
@@ -70,7 +68,7 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu
         # Reshape scale factor so it is broadcastable
         axis = self.get_filters_axis(inbound_node.attributes)
         fused_weights = ir.tensor(
-            weights * self._reshape_for_broadcast(scale_factor, weights.ndim, axis=axis)
+            weights * _reshape_for_broadcast(scale_factor, weights.ndim, axis=axis)
         )
 
         # Update bias
@@ -92,32 +90,37 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu
             attributes=inbound_node.attributes,
         )
 
-    def check(self, context, x, inbound_out, batchnorm_out) -> orp.MatchResult:
+    def check(
+        self, context, x, inbound_out: ir.Value, batchnorm_out: ir.Value
+    ) -> orp.MatchResult:
         del context  # Unused
         check_result = orp.MatchResult()
 
         inbound_node = inbound_out.producer()
         batchnorm_node = batchnorm_out.producer()
 
         # Check that inbound weights + (inbound bias) + batchnorm params are initializers
+        # and that they are not graph inputs
         initializers = [inbound_node.inputs[1], *batchnorm_node.inputs[1:]]
         if len(inbound_node.inputs) > 2:
             initializers.append(inbound_node.inputs[2])
 
         for initializer in initializers:
             if not initializer.is_initializer() or initializer.const_value is None:
-                return check_result.fail(f"{initializer.name} is not a constant initializer")
+                return check_result.fail(f"{initializer.name} is not a constant initializer.")
+            if initializer.is_graph_input():
+                return check_result.fail(f"{initializer.name} is a graph input.")
 
         return check_result
 
 
-class FuseBatchNormIntoConv(FuseBatchNormBase):
+class FuseBatchNormIntoConv(_FuseBatchNormBase):
     """Replaces ``BatchNormalization(Conv(x))`` with ``Conv(x)``."""
 
     def __init__(self):
         super().__init__("Conv")
 
-    def get_filters_axis(self, attributes) -> int:
+    def get_filters_axis(self, attributes: Mapping[str, ir.Attr]) -> int:
         return 0
 
     def pattern(self, op, x):
@@ -128,13 +131,13 @@ def pattern(self, op, x):
         )
 
 
-class FuseBatchNormIntoConvTranspose(FuseBatchNormBase):
+class FuseBatchNormIntoConvTranspose(_FuseBatchNormBase):
     """Replaces ``BatchNormalization(ConvTranspose(x))`` with ``ConvTranspose(x)``."""
 
     def __init__(self):
         super().__init__("ConvTranspose")
 
-    def get_filters_axis(self, attributes) -> int:
+    def get_filters_axis(self, attributes: Mapping[str, ir.Attr]) -> int:
         return 1
 
     def pattern(self, op, x):
@@ -145,14 +148,16 @@ def pattern(self, op, x):
         )
 
 
-class FuseBatchNormIntoGemm(FuseBatchNormBase):
+class FuseBatchNormIntoGemm(_FuseBatchNormBase):
     """Replaces ``BatchNormalization(Gemm(x))`` with ``Gemm(x)``."""
 
     def __init__(self):
         super().__init__("Gemm")
 
-    def get_filters_axis(self, attributes) -> int:
-        return 0 if attributes.get("transB") is not None and attributes["transB"].value else 1
+    def get_filters_axis(self, attributes: Mapping[str, ir.Attr]) -> int:
+        return (
+            0 if attributes.get("transB") is not None and attributes["transB"].as_int() else 1
+        )
 
     def pattern(self, op, x):
         return op.BatchNormalization(

onnxscript/rewriter/fuse_batchnorm_test.py

@@ -12,6 +12,22 @@
 
 
 class FuseBatchnormTest(unittest.TestCase):
+    def _create_batchnorm_params(self, size: int):
+        return [
+            onnx.numpy_helper.from_array(
+                np.random.randn(size).astype(np.float32), name="gamma"
+            ),
+            onnx.numpy_helper.from_array(
+                np.random.randn(size).astype(np.float32), name="beta"
+            ),
+            onnx.numpy_helper.from_array(
+                np.random.randn(size).astype(np.float32), name="input_mean"
+            ),
+            onnx.numpy_helper.from_array(
+                np.abs(np.random.randn(size)).astype(np.float32), name="input_var"
+            ),
+        ]
+
     @parameterized.parameterized.expand(
         [
             ("bias_false", False),
@@ -45,14 +61,7 @@ def test_fuse_batchnorm_convtranspose(self, _: str, convtranspose_bias: bool):
             onnx.numpy_helper.from_array(
                 np.random.randn(32, 64, 3, 3).astype(np.float32), name="W"
             ),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="gamma"),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="beta"),
-            onnx.numpy_helper.from_array(
-                np.random.randn(64).astype(np.float32), name="input_mean"
-            ),
-            onnx.numpy_helper.from_array(
-                np.abs(np.random.randn(64)).astype(np.float32), name="input_var"
-            ),
+            *self._create_batchnorm_params(size=64),
         ]
         if convtranspose_bias:
             initializers.append(
@@ -111,14 +120,7 @@ def test_fuse_batchnorm_conv(self, _: str, conv_bias: bool):
             onnx.numpy_helper.from_array(
                 np.random.randn(64, 32, 3, 3).astype(np.float32), name="W"
             ),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="gamma"),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="beta"),
-            onnx.numpy_helper.from_array(
-                np.random.randn(64).astype(np.float32), name="input_mean"
-            ),
-            onnx.numpy_helper.from_array(
-                np.abs(np.random.randn(64)).astype(np.float32), name="input_var"
-            ),
+            *self._create_batchnorm_params(size=64),
         ]
         if conv_bias:
             initializers.append(
@@ -182,14 +184,7 @@ def test_fuse_batchnorm_gemm(self, _: str, gemm_bias: bool, transB: int):
         # Add initializers
         initializers = [
             onnx.numpy_helper.from_array(weights, name="W"),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="gamma"),
-            onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="beta"),
-            onnx.numpy_helper.from_array(
-                np.random.randn(64).astype(np.float32), name="input_mean"
-            ),
-            onnx.numpy_helper.from_array(
-                np.abs(np.random.randn(64)).astype(np.float32), name="input_var"
-            ),
+            *self._create_batchnorm_params(size=64),
         ]
         if gemm_bias:
             initializers.append(
@@ -233,6 +228,30 @@ def test_fuse_batchnorm_non_initializers(self):
         # No changes were applied
         self.assertEqual(count, 0)
 
+    def test_fuse_batchnorm_graph_inputs(self):
+        model_proto = onnx.parser.parse_model("""
+            < ir_version: 7, opset_import: ["" : 17] >
+            test_model (float[N, 32, 14, 16] X, float[64, 32, 3, 3] W) => (float [N, ?, ?, ?] Y)
+            {
+                X1 = Conv(X, W)
+                Y = BatchNormalization(X1, gamma, beta, input_mean, input_var)
+            }
+        """)
+        initializers = [
+            onnx.numpy_helper.from_array(
+                np.random.randn(64, 32, 3, 3).astype(np.float32), name="W"
+            ),
+            *self._create_batchnorm_params(size=64),
+        ]
+        model_proto.graph.initializer.extend(initializers)
+        onnx.checker.check_model(model_proto, True)
+
+        model = ir.serde.deserialize_model(model_proto)
+        count = fuse_batchnorm.fuse_batchnorm_rule_set().apply_to_model(model)
+
+        # No changes were applied as W is a graph input
+        self.assertEqual(count, 0)
+
 
 if __name__ == "__main__":
     unittest.main()