Support common subexpression elimination pass (CSE)

onnxscript/ir/_core.py

@@ -2788,7 +2788,7 @@ def __init__(
         model_version: int | None = None,
         doc_string: str | None = None,
         functions: Sequence[Function] = (),
-        meta_data_props: dict[str, str] | None = None,
+        metadata_props: dict[str, str] | None = None,
     ) -> None:
         self.graph: Graph = graph
         self.ir_version = ir_version
@@ -2799,7 +2799,7 @@ def __init__(
         self.doc_string = doc_string
         self._functions = {func.identifier(): func for func in functions}
         self._metadata: _metadata.MetadataStore | None = None
-        self._metadata_props: dict[str, str] | None = meta_data_props
+        self._metadata_props: dict[str, str] | None = metadata_props
 
     @property
     def functions(self) -> dict[_protocols.OperatorIdentifier, Function]:

onnxscript/ir/passes/common/common_subexpression_elimination.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Eliminate common subexpression in ONNX graphs."""
+
+from __future__ import annotations
+
+__all__ = [
+    "CommonSubexpressionEliminationPass",
+]
+
+import logging
+
+from onnxscript import ir
+
+logger = logging.getLogger(__name__)
+
+
+class CommonSubexpressionEliminationPass(ir.passes.InPlacePass):
+    """Eliminate common subexpression in ONNX graphs."""
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Return the same ir.Model but with CSE applied to the graph."""
+        modified = False
+        graph = model.graph
+
+        modified = _common_subexpression_elimination(graph, modified)
+
+        return ir.passes.PassResult(
+            model,
+            modified=modified,
+        )
+
+
+def _common_subexpression_elimination(graph: ir.Graph, modified: bool) -> bool:
+    """Eliminate common subexpression in ONNX graphs."""
+
+    # node to node identifier, length of outputs, inputs, and attributes
+    existing_node_info_to_the_node: dict[
+        tuple[
+            tuple[str, str, str],  # op_identifier
+            int,  # len(outputs)
+            tuple[int, ...],  # input ids
+            tuple[tuple[str, object], ...],  # attributes
+        ],
+        ir.Node,
+    ] = {}
+    previous_node = None
+
+    for node in graph:
+        # Use equality to check if the node is a common subexpression.
+        attributes = {}
+        for k, v in node.attributes.items():
+            assert isinstance(v, ir.Attr)
+            if isinstance(v.value, ir.Graph):
+                modified = _common_subexpression_elimination(v.value, modified)
+            attributes[k] = v.value
+
+        node_info = (
+            node.op_identifier(),
+            len(node.outputs),
+            tuple(id(input) for input in node.inputs),
+            tuple(sorted(attributes.items())),
+        )
+        # Check if the node is a common subexpression.
+        if node_info in existing_node_info_to_the_node:
+            # If it is, this node is already in the new graph, so
+            # we don't need to create a new node.
+            modified = True
+            existing_node = existing_node_info_to_the_node[node_info]
+            ir.convenience.replace_nodes_and_values(
+                graph,
+                insertion_point=previous_node or node,
+                old_nodes=[node],
+                new_nodes=[existing_node],
+                old_values=node.outputs,
+                new_values=existing_node.outputs,
+            )
+            previous_node = existing_node
+            logger.debug("Reusing node %s", existing_node.name)
+        else:
+            # If it is not, add to the mapping.
+            existing_node_info_to_the_node[node_info] = node
+            previous_node = node
+    return modified

onnxscript/ir/passes/common/common_subexpression_elimination_test.py

@@ -0,0 +1,237 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import onnxruntime as ort
+
+from onnxscript import FLOAT, ir, script
+from onnxscript import opset18 as op
+from onnxscript.ir.passes.common import common_subexpression_elimination
+
+
+class TestCommonSubexpressionEliminationPass(unittest.TestCase):
+    def check_graph(self, model: ir.Model, inputs: list[ir.Value], delta_nodes: list[int]):
+        """Check if the model applied the CSE pass correctly."""
+        assert len(list(model.graphs())) == len(delta_nodes)
+        # Log all results from the original model.
+        # 1. model graph node counts
+        original_graphs_node_count = np.array([graph.num_nodes() for graph in model.graphs()])
+        model_proto = ir.serde.serialize_model(model)
+
+        # 2. model outputs
+        ort_inputs = {
+            k.name: np.random.rand(*v.shape).astype(np.float32)
+            for k, v in zip(model.graph.inputs, inputs)
+        }
+        original_model_session = ort.InferenceSession(model_proto.SerializeToString())
+        original_model_results = original_model_session.run(None, ort_inputs)
+
+        result = common_subexpression_elimination.CommonSubexpressionEliminationPass()(model)
+
+        result_graphs_node_count = np.array([graph.num_nodes() for graph in model.graphs()])
+
+        # Check if the number of nodes in the model is correct
+        self.assertTrue(
+            np.array_equal(
+                original_graphs_node_count, np.add(result_graphs_node_count, delta_nodes)
+            )
+        )
+        self.assertEqual(
+            result.modified, any(original_graphs_node_count > result_graphs_node_count)
+        )
+
+        result_proto = ir.serde.serialize_model(result.model)
+        result_session = ort.InferenceSession(result_proto.SerializeToString())
+        result_results = result_session.run(None, ort_inputs)
+
+        # Check if the models produce the same output
+        # with the same inputs
+        for idx, original_model_result in enumerate(original_model_results):
+            np.testing.assert_allclose(
+                original_model_result, result_results[idx], rtol=1e-5, atol=1e-5
+            )
+
+    def test_two_branches_with_the_same_operations_is_csed(self):
+        """Test if two branches with the same operations are CSEd.
+
+        def test_simple(self):
+            def f(x):
+                a = x.cos()
+                b = x.cos()
+                c = a + a
+                d = b + b
+                return c + d
+
+            x = torch.randn(2, 2)
+        """
+
+        @script()
+        def test_model(x: FLOAT[2, 2]) -> FLOAT[2, 2]:
+            a = op.Cos(x)
+            b = op.Cos(x)
+            c = a + a
+            d = b + b
+            return c + d
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+
+        self.check_graph(model, [np.random.rand(2, 2)], delta_nodes=[2])
+
+    def test_more_operations_in_two_branches_with_the_same_operations_is_csed(self):
+        """Test if two branches with the same operations are CSEd.
+
+        def test_simple(self):
+            def f(x):
+                a = x.cos().sin()
+                b = x.cos().sin()
+                c = a + a
+                d = b + b
+                return c + d
+
+            x = torch.randn(2, 2)
+        """
+
+        @script()
+        def test_model(x: FLOAT[1]) -> FLOAT[1]:
+            a = op.Sin(op.Cos(x))
+            b = op.Sin(op.Cos(x))
+            c = a + a
+            d = b + b
+            return c + d
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.check_graph(model, [np.random.rand(1)], delta_nodes=[3])
+
+    def test_multiple_same_ops_with_attributes_are_csed(self):
+        """Test if multiple same ops are CSEd.
+
+        def f(x):
+            a = x.sum()
+            b = x.sum()
+            c = x.sum()
+            d = x.sum()
+            return a + b + c + d
+
+        x = torch.randn(2, 2)
+
+        """
+
+        @script()
+        def test_model(x: FLOAT[2, 2]) -> FLOAT[2, 2]:
+            a = op.ReduceSum(x, keepdims=False)
+            b = op.ReduceSum(x, keepdims=False)
+            c = op.ReduceSum(x, keepdims=False)
+            d = op.ReduceSum(x, keepdims=False)
+            return a + b + c + d
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.check_graph(model, [np.random.rand(2, 2)], delta_nodes=[3])
+
+    def test_the_ops_with_the_same_inputs_but_different_attributes_are_not_csed(self):
+        """Test if the ops with the same inputs but different attributes are not CSEd.
+
+        def f(x):
+            a = x.sum()
+            b = x.sum(keepdims=True)
+            c = x.sum()
+            d = x.sum(keepdims=True)
+            return a + b + c + d
+
+        x = torch.randn(2, 2)
+
+        """
+
+        @script()
+        def test_model(x: FLOAT[2, 2]) -> FLOAT[2, 2]:
+            a = op.ReduceSum(x, keepdims=False)
+            b = op.ReduceSum(x, keepdims=True)
+            return a + b
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.check_graph(model, [np.random.rand(2, 2)], delta_nodes=[0])
+
+    def test_control_flow_if_ops_are_not_csed_as_graph_attr_is_not_matched(self):
+        """Test if control flow ops are not CSEd.
+
+        def f(a, b):
+            rank = a.rank()
+            if rank == 2:
+                result1 = a - b
+            else:
+                result1 = a + b
+            if rank == 2:
+                result2 = a - b
+            else:
+                result2 = a + b
+            return result1 + result2
+
+        x = torch.randn(2, 2)
+
+        """
+
+        @script()
+        def test_model(a: FLOAT[2, 2], b: FLOAT[2, 2]) -> FLOAT[2, 2]:
+            rank = op.Size(op.Shape(a))
+            if rank == 2:
+                result1 = a - b
+            else:
+                result1 = a + b
+            if rank == 2:
+                result2 = a - b
+            else:
+                result2 = a + b
+            return result1 + result2
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.check_graph(
+            model, [np.random.rand(2, 2), np.random.rand(2, 2)], delta_nodes=[0, 0, 0, 0, 0]
+        )
+
+    def test_subgraph_is_csed(self):
+        """Test if control flow ops are not CSEd.
+
+        def f(x):
+            rank = x.rank()
+            if rank == 2:
+                a = x.cos()
+                b = x.cos()
+                c = a + a
+                d = b + b
+                return c + d
+            else:
+                a = x.sin()
+                b = x.sin()
+                c = a + a
+                d = b + b
+            return c + d
+
+        x = torch.randn(2, 2)
+
+        """
+
+        @script()
+        def test_model(x: FLOAT[2, 2]) -> FLOAT[2, 2]:
+            rank = op.Size(op.Shape(x))
+            if rank == 2:
+                a = op.Cos(x)
+                b = op.Cos(x)
+                c = a + a
+                d = b + b
+            else:
+                a = op.Sin(x)
+                b = op.Sin(x)
+                c = a + a
+                d = b + b
+            return c + d
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.check_graph(model, [np.random.rand(2, 2)], delta_nodes=[0, 2, 2])

onnxscript/ir/serde.py

@@ -511,7 +511,7 @@ def deserialize_model(proto: onnx.ModelProto) -> _core.Model:
         model_version=_get_field(proto, "model_version"),
         doc_string=_get_field(proto, "doc_string"),
         functions=functions,
-        meta_data_props=deserialize_metadata_props(proto.metadata_props),
+        metadata_props=deserialize_metadata_props(proto.metadata_props),
     )
 
     # Handle experimental value info for functions created by the dynamo exporter in IR version 9