Add fusion rule to remove Expand node before broadcast-capable binary operators

Copilot · xadupre · Copilot · commit 161bcd08ba92 · 2026-03-20T16:29:06.000Z
Co-authored-by: xadupre &lt;22452781+xadupre@users.noreply.github.com&gt;
diff --git a/onnxscript/rewriter/rules/common/__init__.py b/onnxscript/rewriter/rules/common/__init__.py
@@ -12,6 +12,7 @@
     "div_by_1_rule",
     "dropout_inference_rule",
     "dropout_zero_rule",
+    "expand_before_binary_op_rules",
     "flatten_to_reshape_rule",
     "fuse_batchnorm_into_conv_rule",
     "fuse_batchnorm_into_conv_transpose_rule",
@@ -125,6 +126,9 @@
     no_op_dynamic_scatter_nd_rule,
     no_op_static_scatter_nd_rule,
 )
+from onnxscript.rewriter.rules.common._remove_expand_before_binary_op import (
+    expand_before_binary_op_rules,
+)
 from onnxscript.rewriter.rules.common._remove_optional_bias import (
     remove_optional_bias_from_conv_rule,
     remove_optional_bias_from_conv_transpose_rule,
diff --git a/onnxscript/rewriter/rules/common/_remove_expand_before_binary_op.py b/onnxscript/rewriter/rules/common/_remove_expand_before_binary_op.py
@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Fusion rule to remove an Expand node before a binary operator.
+
+This implements the optimization:
+
+    BinaryOp(Expand(x, shape), y) -> BinaryOp(x, y)
+    BinaryOp(x, Expand(y, shape)) -> BinaryOp(x, y)
+
+This is valid when the binary operator's broadcasting semantics would produce
+the same output shape as first expanding the input and then applying the op.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from onnxscript import ir
+from onnxscript.rewriter._basics import MatchResult
+from onnxscript.rewriter._ir_utils import get_numpy_value
+from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet
+
+# Binary operators in ONNX standard opset that support numpy-style broadcasting.
+_BROADCAST_BINARY_OPS: tuple[str, ...] = (
+    "Add",
+    "And",
+    "BitShift",
+    "BitwiseAnd",
+    "BitwiseOr",
+    "BitwiseXor",
+    "Div",
+    "Equal",
+    "Greater",
+    "GreaterOrEqual",
+    "Less",
+    "LessOrEqual",
+    "Mod",
+    "Mul",
+    "Or",
+    "Pow",
+    "PRelu",
+    "Sub",
+    "Xor",
+)
+
+
+def _check_expand_removable(
+    expand_input: ir.Value,
+    shape: ir.Value,
+    other_input: ir.Value,
+) -> MatchResult:
+    """Check if an Expand node can be safely removed before a binary op.
+
+    The Expand is removable if the binary op's broadcasting produces the same
+    output shape when using the original (pre-expand) tensor directly.
+
+    Args:
+        expand_input: The value fed into the Expand node.
+        shape: The target shape operand of the Expand node (must be a constant).
+        other_input: The other operand of the binary op.
+
+    Returns:
+        A MatchResult that is successful when the Expand can be removed.
+    """
+    check_result = MatchResult()
+
+    # Need static shape info for both inputs.
+    expand_input_shape = expand_input.shape
+    other_shape = other_input.shape
+    if expand_input_shape is None or other_shape is None:
+        return check_result.fail("Input shapes are not statically known.")
+
+    # Require fully static (integer-only) shapes to avoid symbolic dim issues.
+    if not expand_input_shape.is_static() or not other_shape.is_static():
+        return check_result.fail("Input shapes are not fully static.")
+
+    # The Expand target shape must be a compile-time constant.
+    expand_shape_val = get_numpy_value(shape)
+    if expand_shape_val is None:
+        return check_result.fail("Expand target shape is not a constant.")
+
+    expand_shape = tuple(int(v) for v in expand_shape_val.tolist())
+    x_shape = tuple(int(d) for d in expand_input_shape)
+    y_shape = tuple(int(d) for d in other_shape)
+
+    # Verify that removing the Expand does not change the binary op's output shape.
+    try:
+        result_with_expand = np.broadcast_shapes(expand_shape, y_shape)
+        result_without_expand = np.broadcast_shapes(x_shape, y_shape)
+    except ValueError:
+        return check_result.fail("Shapes are not broadcastable.")
+
+    if result_with_expand != result_without_expand:
+        return check_result.fail(
+            f"Removing Expand would change output shape from "
+            f"{result_with_expand} to {result_without_expand}."
+        )
+
+    return check_result
+
+
+class _ExpandFirstInput(RewriteRuleClassBase):
+    """Removes ``BinaryOp(Expand(x, shape), y)`` -> ``BinaryOp(x, y)``."""
+
+    def __init__(self, op_type: str) -> None:
+        super().__init__(f"ExpandFirst_{op_type}", remove_nodes=False)
+        self._op_type = op_type
+
+    def pattern(self, op, x: ir.Value, shape: ir.Value, y: ir.Value) -> ir.Value:
+        return getattr(op, self._op_type)(op.Expand(x, shape), y)
+
+    def check(self, context, x: ir.Value, shape: ir.Value, y: ir.Value) -> MatchResult:
+        del context  # Unused
+        return _check_expand_removable(x, shape, y)
+
+    def rewrite(self, op, x: ir.Value, shape: ir.Value, y: ir.Value) -> ir.Value:
+        return getattr(op, self._op_type)(x, y)
+
+
+class _ExpandSecondInput(RewriteRuleClassBase):
+    """Removes ``BinaryOp(x, Expand(y, shape))`` -> ``BinaryOp(x, y)``."""
+
+    def __init__(self, op_type: str) -> None:
+        super().__init__(f"ExpandSecond_{op_type}", remove_nodes=False)
+        self._op_type = op_type
+
+    def pattern(self, op, x: ir.Value, y: ir.Value, shape: ir.Value) -> ir.Value:
+        return getattr(op, self._op_type)(x, op.Expand(y, shape))
+
+    def check(self, context, x: ir.Value, y: ir.Value, shape: ir.Value) -> MatchResult:
+        del context  # Unused
+        return _check_expand_removable(y, shape, x)
+
+    def rewrite(self, op, x: ir.Value, y: ir.Value, shape: ir.Value) -> ir.Value:
+        return getattr(op, self._op_type)(x, y)
+
+
+def _make_expand_before_binary_op_rules() -> list:
+    """Create rewrite rules for removing Expand before each supported binary op."""
+    rules = []
+    for op_type in _BROADCAST_BINARY_OPS:
+        rules.append(_ExpandFirstInput.rule(op_type))
+        rules.append(_ExpandSecondInput.rule(op_type))
+    return rules
+
+
+expand_before_binary_op_rules = RewriteRuleSet(_make_expand_before_binary_op_rules())
diff --git a/onnxscript/rewriter/rules/common/_remove_expand_before_binary_op_test.py b/onnxscript/rewriter/rules/common/_remove_expand_before_binary_op_test.py
@@ -0,0 +1,228 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Tests for the remove-Expand-before-binary-op fusion rule."""
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import onnx.reference
+import parameterized
+
+import onnxscript.ir as ir
+from onnxscript.rewriter.rules.common import _remove_expand_before_binary_op as mod
+
+
+def _run_model(model: ir.Model, feeds: dict) -> list:
+    """Run a model using the ONNX reference evaluator."""
+    proto = ir.to_proto(model)
+    ref = onnx.reference.ReferenceEvaluator(proto)
+    return ref.run(None, feeds)
+
+
+class RemoveExpandBeforeBinaryOpTest(unittest.TestCase):
+    """Tests for _remove_expand_before_binary_op rules."""
+
+    def _apply_and_check(
+        self,
+        model_text: str,
+        expected_count: int,
+        expected_op_types: list[str],
+    ) -> ir.Model:
+        """Helper: apply the rules and verify the result."""
+        model = ir.from_onnx_text(model_text)
+        count = mod.expand_before_binary_op_rules.apply_to_model(model)
+        self.assertEqual(count, expected_count)
+        actual_op_types = [node.op_type for node in model.graph]
+        self.assertEqual(actual_op_types, expected_op_types)
+        return model
+
+    # ------------------------------------------------------------------
+    # Cases where the Expand should be removed
+    # ------------------------------------------------------------------
+
+    @parameterized.parameterized.expand(
+        [
+            ("Add",),
+            ("Sub",),
+            ("Mul",),
+            ("Div",),
+        ]
+    )
+    def test_expand_first_input_same_shape_is_removed(self, op_type: str):
+        """Expand producing same shape as input should be removed from BinaryOp."""
+        model_text = f"""
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3, 4] x, float[3, 4] y) => (float[3, 4] output)
+            <int64[2] shape = {{3, 4}}>
+            {{
+                expanded = Expand(x, shape)
+                output = {op_type}(expanded, y)
+            }}
+        """
+        model = self._apply_and_check(model_text, 1, [op_type])
+
+        # Verify numerical correctness
+        x = np.random.randn(3, 4).astype(np.float32)
+        y = np.random.randn(3, 4).astype(np.float32)
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    def test_expand_first_input_broadcast_covered_by_other_input(self):
+        """Expand from [3, 4] to [4, 3, 4] can be removed when y has shape [4, 3, 4]."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3, 4] x, float[4, 3, 4] y) => (float[4, 3, 4] output)
+            <int64[3] shape = {4, 3, 4}>
+            {
+                expanded = Expand(x, shape)
+                output = Add(expanded, y)
+            }
+        """
+        model = self._apply_and_check(model_text, 1, ["Add"])
+
+        x = np.random.randn(3, 4).astype(np.float32)
+        y = np.random.randn(4, 3, 4).astype(np.float32)
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    def test_expand_second_input_is_removed(self):
+        """Expand on the second input of a binary op should be removed."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 3, 4] x, float[3, 4] y) => (float[4, 3, 4] output)
+            <int64[3] shape = {4, 3, 4}>
+            {
+                expanded = Expand(y, shape)
+                output = Mul(x, expanded)
+            }
+        """
+        model = self._apply_and_check(model_text, 1, ["Mul"])
+
+        x = np.random.randn(4, 3, 4).astype(np.float32)
+        y = np.random.randn(3, 4).astype(np.float32)
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    def test_expand_with_broadcast_compatible_other_input(self):
+        """Expand from [3] to [4, 3] can be removed when y has shape [4, 1]."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3] x, float[4, 1] y) => (float[4, 3] output)
+            <int64[2] shape = {4, 3}>
+            {
+                expanded = Expand(x, shape)
+                output = Add(expanded, y)
+            }
+        """
+        model = self._apply_and_check(model_text, 1, ["Add"])
+
+        x = np.random.randn(3).astype(np.float32)
+        y = np.random.randn(4, 1).astype(np.float32)
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    def test_expand_sub_first_input_is_removed(self):
+        """Expand on the first input of Sub should be removed."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3, 4] x, float[3, 4] y) => (float[3, 4] output)
+            <int64[2] shape = {3, 4}>
+            {
+                expanded = Expand(x, shape)
+                output = Sub(expanded, y)
+            }
+        """
+        model = self._apply_and_check(model_text, 1, ["Sub"])
+
+        x = np.random.randn(3, 4).astype(np.float32)
+        y = np.random.randn(3, 4).astype(np.float32)
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    def test_expand_div_second_input_is_removed(self):
+        """Expand on the second input of Div should be removed."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 3, 4] x, float[3, 4] y) => (float[4, 3, 4] output)
+            <int64[3] shape = {4, 3, 4}>
+            {
+                expanded = Expand(y, shape)
+                output = Div(x, expanded)
+            }
+        """
+        model = self._apply_and_check(model_text, 1, ["Div"])
+
+        x = np.random.randn(4, 3, 4).astype(np.float32)
+        y = (np.random.randn(3, 4).astype(np.float32) + 2.0)  # avoid division by zero
+        original = ir.from_onnx_text(model_text)
+        expected = _run_model(original, {"x": x, "y": y})
+        got = _run_model(model, {"x": x, "y": y})
+        np.testing.assert_allclose(got[0], expected[0], rtol=1e-5)
+
+    # ------------------------------------------------------------------
+    # Cases where the Expand should NOT be removed
+    # ------------------------------------------------------------------
+
+    def test_expand_changes_output_shape_not_removed(self):
+        """Expand that changes the output shape compared to direct broadcast must be kept."""
+        # x has shape [3], expand to [4, 3], other is a scalar.
+        # With expand: broadcast([4, 3], []) = [4, 3]
+        # Without expand: broadcast([3], []) = [3]  <- different!
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3] x) => (float[4, 3] output)
+            <int64[2] shape = {4, 3}, float[1] one = {1.0}>
+            {
+                expanded = Expand(x, shape)
+                output = Add(expanded, one)
+            }
+        """
+        model = ir.from_onnx_text(model_text)
+        count = mod.expand_before_binary_op_rules.apply_to_model(model)
+        self.assertEqual(count, 0)
+
+    def test_expand_target_shape_not_constant_not_removed(self):
+        """Expand with a dynamic (non-constant) shape cannot be removed."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3, 4] x, float[3, 4] y, int64[2] shape) => (float[3, 4] output)
+            {
+                expanded = Expand(x, shape)
+                output = Add(expanded, y)
+            }
+        """
+        model = ir.from_onnx_text(model_text)
+        count = mod.expand_before_binary_op_rules.apply_to_model(model)
+        self.assertEqual(count, 0)
+
+    def test_expand_unknown_input_shape_not_removed(self):
+        """Expand cannot be removed when the input shape is not statically known."""
+        # No shape annotation on 'x'
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N] x, float[3, 4] y) => (float[3, 4] output)
+            <int64[2] shape = {3, 4}>
+            {
+                expanded = Expand(x, shape)
+                output = Add(expanded, y)
+            }
+        """
+        model = ir.from_onnx_text(model_text)
+        count = mod.expand_before_binary_op_rules.apply_to_model(model)
+        self.assertEqual(count, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()
