diff --git a/tests/transforms/scalar/test_algebraic_simplify.py b/tests/transforms/scalar/test_algebraic_simplify.py index 706465c..4db78c7 100644 --- a/tests/transforms/scalar/test_algebraic_simplify.py +++ b/tests/transforms/scalar/test_algebraic_simplify.py @@ -7,8 +7,7 @@ import unittest import tensorflow.compat.v1 as tf -from graph_optimizer.core import GraphOptimizer -from graph_optimizer.transforms.scalar.algebraic_simplify import AlgebraicSimplifyPass +from graph_optimizer.runner import OptimizationPipeline from graph_optimizer.utils.graph_utils import create_node, create_const_node @@ -16,6 +15,11 @@ class AlgebraicSimplifyPassTest(unittest.TestCase): + def _run_optimization(self, graph_def, protected_nodes=None): + """Helper to run the level 1 optimization pipeline.""" + pipeline = OptimizationPipeline(graph_def=graph_def, level=1, protected_nodes=protected_nodes or []) + return pipeline.run() + def create_graph(self, nodes): """Helper to create a GraphDef from node list.""" graph_def = tf.GraphDef() @@ -28,13 +32,8 @@ def test_add_zero_left(self): add = create_node("Add", name="add", inputs=["zero", "x"]) graph = self.create_graph([x, zero, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - # add should be replaced by x - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("add", names) @@ -44,12 +43,8 @@ def test_add_zero_right(self): add = create_node("Add", name="add", inputs=["x", "zero"]) graph = self.create_graph([x, zero, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("add", names) @@ -59,12 +54,8 @@ def test_sub_zero(self): sub = create_node("Sub", name="sub", inputs=["x", "zero"]) graph = self.create_graph([x, zero, sub]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("sub", names) @@ -74,12 +65,8 @@ def test_mul_one_left(self): mul = create_node("Mul", name="mul", inputs=["one", "x"]) graph = self.create_graph([x, one, mul]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("mul", names) @@ -89,12 +76,8 @@ def test_mul_one_right(self): mul = create_node("Mul", name="mul", inputs=["x", "one"]) graph = self.create_graph([x, one, mul]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("mul", names) @@ -105,21 +88,14 @@ def test_mul_zero_left(self): mul = create_node("Mul", name="mul", inputs=["zero", "x"]) graph = self.create_graph([x, zero, mul]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["mul_zero"] - ) - - # result should be a zero const + optimized_graph = self._run_optimization(graph, protected_nodes=["mul_zero"]) zeros = [ n - for n in optimizer.graph_def.node + for n in optimized_graph.node if n.op == "Const" and n.name == "mul_zero" ] self.assertEqual(len(zeros), 1) - self.assertNotIn("mul", {n.name for n in optimizer.graph_def.node}) + self.assertNotIn("mul", {n.name for n in optimized_graph.node}) def test_div_one(self): x = create_node("Placeholder", name="x") @@ -127,12 +103,8 @@ def test_div_one(self): div = create_node("Div", name="div", inputs=["x", "one"]) graph = self.create_graph([x, one, div]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("div", names) @@ -142,12 +114,8 @@ def test_neg_neg(self): neg2 = create_node("Neg", name="neg2", inputs=["neg1"]) graph = self.create_graph([x, neg1, neg2]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("neg1", names) self.assertNotIn("neg2", names) @@ -158,12 +126,8 @@ def test_logical_not_not(self): not2 = create_node("LogicalNot", name="not2", inputs=["not1"]) graph = self.create_graph([x, not1, not2]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("not1", names) self.assertNotIn("not2", names) @@ -174,20 +138,14 @@ def test_equal_same(self): eq = create_node("Equal", name="eq", inputs=["x", "x"]) graph = self.create_graph([x, eq]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["eq_bool"] - ) - + optimized_graph = self._run_optimization(graph, protected_nodes=["eq_bool"]) trues = [ n - for n in optimizer.graph_def.node + for n in optimized_graph.node if n.op == "Const" and n.name == "eq_bool" ] self.assertEqual(len(trues), 1) - self.assertNotIn("eq", {n.name for n in optimizer.graph_def.node}) + self.assertNotIn("eq", {n.name for n in optimized_graph.node}) def test_select_same_branch(self): cond = create_node("Placeholder", name="cond") @@ -195,12 +153,8 @@ def test_select_same_branch(self): sel = create_node("Select", name="sel", inputs=["cond", "x", "x"]) graph = self.create_graph([cond, x, sel]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("sel", names) @@ -210,13 +164,8 @@ def test_no_simplify_add_nonzero(self): add = create_node("Add", name="add", inputs=["x", "y"]) graph = self.create_graph([x, y, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - # Should remain unchanged - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("add", names) def test_sub_same(self): @@ -225,14 +174,8 @@ def test_sub_same(self): sub = create_node("Sub", name="sub", inputs=["x", "x"]) graph = self.create_graph([x, sub]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["sub_zero"] - ) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph, protected_nodes=["sub_zero"]) + names = {n.name for n in optimized_graph.node} self.assertIn("sub_zero", names) self.assertNotIn("sub", names) @@ -243,14 +186,8 @@ def test_add_neg(self): add = create_node("Add", name="add", inputs=["x", "neg"]) graph = self.create_graph([x, neg, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["add_zero"] - ) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph, protected_nodes=["add_zero"]) + names = {n.name for n in optimized_graph.node} self.assertIn("add_zero", names) self.assertNotIn("add", names) @@ -259,12 +196,8 @@ def test_mul_same(self): mul = create_node("Mul", name="mul", inputs=["x", "x"]) graph = self.create_graph([x, mul]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - ops = {n.op for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + ops = {n.op for n in optimized_graph.node} self.assertIn("Square", ops) self.assertNotIn("Mul", ops) @@ -274,14 +207,8 @@ def test_div_same(self): div = create_node("Div", name="div", inputs=["x", "x"]) graph = self.create_graph([x, div]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["div_one"] - ) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph, protected_nodes=["div_one"]) + names = {n.name for n in optimized_graph.node} self.assertIn("div_one", names) self.assertNotIn("div", names) @@ -291,12 +218,8 @@ def test_pow_one(self): pow_node = create_node("Pow", name="pow", inputs=["x", "one"]) graph = self.create_graph([x, one, pow_node]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("pow", names) @@ -306,12 +229,8 @@ def test_pow_two(self): pow_node = create_node("Pow", name="pow", inputs=["x", "two"]) graph = self.create_graph([x, two, pow_node]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - ops = {n.op for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + ops = {n.op for n in optimized_graph.node} self.assertIn("Square", ops) self.assertNotIn("Pow", ops) @@ -322,15 +241,8 @@ def test_logical_and_false(self): and_node = create_node("LogicalAnd", name="and_node", inputs=["x", "false"]) graph = self.create_graph([x, false_node, and_node]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["and_node_bool"] - ) - - consts = [n for n in optimizer.graph_def.node if n.op == "Const"] - # Should have and_node_bool (False) + optimized_graph = self._run_optimization(graph, protected_nodes=["and_node_bool"]) + consts = [n for n in optimized_graph.node if n.op == "Const"] has_false = any( n.name == "and_node_bool" and n.attr["value"].tensor.bool_val[0] == False for n in consts @@ -344,14 +256,8 @@ def test_logical_or_true(self): or_node = create_node("LogicalOr", name="or_node", inputs=["x", "true"]) graph = self.create_graph([x, true_node, or_node]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["or_node_bool"] - ) - - consts = [n for n in optimizer.graph_def.node if n.op == "Const"] + optimized_graph = self._run_optimization(graph, protected_nodes=["or_node_bool"]) + consts = [n for n in optimized_graph.node if n.op == "Const"] has_true = any( n.name == "or_node_bool" and n.attr["value"].tensor.bool_val[0] == True for n in consts @@ -366,12 +272,8 @@ def test_add_zero_broadcast_positive(self): add = create_node("Add", name="add", inputs=["x", "zero"]) graph = self.create_graph([x, zero, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("x", names) self.assertNotIn("add", names) @@ -384,13 +286,8 @@ def test_add_zero_broadcast_negative(self): add = create_node("Add", name="add", inputs=["x", "zero"]) graph = self.create_graph([x, zero, add]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform(optimizer, auto_cleanup=True) - - # Should NOT simplify - names = {n.name for n in optimizer.graph_def.node} + optimized_graph = self._run_optimization(graph) + names = {n.name for n in optimized_graph.node} self.assertIn("add", names) def test_mul_zero_broadcast(self): @@ -402,15 +299,8 @@ def test_mul_zero_broadcast(self): mul = create_node("Mul", name="mul", inputs=["x", "zero"]) graph = self.create_graph([x, zero, mul]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["mul_zero"] - ) - - # Should simplify to a [2, 2] zero constant - folded = [n for n in optimizer.graph_def.node if n.name == "mul_zero"] + optimized_graph = self._run_optimization(graph, protected_nodes=["mul_zero"]) + folded = [n for n in optimized_graph.node if n.name == "mul_zero"] self.assertEqual(len(folded), 1) shape = [d.size for d in folded[0].attr["value"].tensor.tensor_shape.dim] self.assertEqual(shape, [2, 2]) @@ -424,15 +314,8 @@ def test_logical_and_broadcast_negative(self): and_node = create_node("LogicalAnd", name="and_node", inputs=["x", "false"]) graph = self.create_graph([x, false_node, and_node]) - optimizer = GraphOptimizer(graph) - optimizer.load_state(graph) - simplify_pass = AlgebraicSimplifyPass() - simplify_pass.transform( - optimizer, auto_cleanup=True, protected_nodes=["and_node_bool"] - ) - - # Should simplify to a [2] False constant - folded = [n for n in optimizer.graph_def.node if n.name == "and_node_bool"] + optimized_graph = self._run_optimization(graph, protected_nodes=["and_node_bool"]) + folded = [n for n in optimized_graph.node if n.name == "and_node_bool"] self.assertEqual(len(folded), 1) shape = [d.size for d in folded[0].attr["value"].tensor.tensor_shape.dim] self.assertEqual(shape, [2]) diff --git a/transforms/__init__.py b/transforms/__init__.py index d7cbd9d..b89aee4 100644 --- a/transforms/__init__.py +++ b/transforms/__init__.py @@ -27,7 +27,6 @@ from .scalar import ( CSEPass, ConstantFoldPass, - AlgebraicSimplifyPass, ) # Combine transforms @@ -44,7 +43,6 @@ # Scalar 'CSEPass', 'ConstantFoldPass', - 'AlgebraicSimplifyPass', # Combine 'ConcatCombinePass', # Vectorize diff --git a/transforms/scalar/__init__.py b/transforms/scalar/__init__.py index 040c809..ba376f5 100644 --- a/transforms/scalar/__init__.py +++ b/transforms/scalar/__init__.py @@ -6,7 +6,7 @@ 类似 LLVM 的 InstCombine、DCE、CSE 等 Pass。 包含的 Pass： -- algebraic_simplify.py : 代数恒等式化简（包括 Identity 折叠、算术/逻辑/比较恒等变换） +- algebraic_simplify/ : 代数恒等式化简（包括 Identity 折叠、算术/逻辑/比较恒等变换） - cse.py : 公共子表达式消除（签名去重） 特点： @@ -17,10 +17,9 @@ from .cse import CSEPass from .constant_fold import ConstantFoldPass -from .algebraic_simplify import AlgebraicSimplifyPass +from .algebraic_simplify import * __all__ = [ 'CSEPass', 'ConstantFoldPass', - 'AlgebraicSimplifyPass', ] diff --git a/transforms/scalar/algebraic_simplify.py b/transforms/scalar/algebraic_simplify.py deleted file mode 100644 index 76c94af..0000000 --- a/transforms/scalar/algebraic_simplify.py +++ /dev/null @@ -1,423 +0,0 @@ -""" -Algebraic Simplify Pass -======================= - -Purpose: --------- -Performs algebraic simplification by applying identity laws, zero-element elimination, -and inverse operation cancellation on graph operations. This includes transforming -operations like `Add(x, 0) → x`, `Mul(x, 1) → x`, `Neg(Neg(x)) → x`, etc. - -This pass generalizes `IdentityEliminationPass` by covering arithmetic, logical, and -comparison identities beyond pure Identity nodes. - -Algorithm: ----------- -1. Define patterns for common algebraic identities where one or more inputs are - constants or repeated variables. -2. Match these patterns in the graph. -3. Replace matched subgraphs with simplified expressions according to algebra rules. -4. Run iteratively until no more simplifications apply (convergence). - -Supported identities include: -- Add(x, 0) → x ; Add(0, x) → x -- Sub(x, 0) → x -- Mul(x, 1) → x ; Mul(1, x) → x -- Mul(x, 0) → 0 (with care for broadcasting) -- Div(x, 1) → x -- Neg(Neg(x)) → x -- LogicalNot(LogicalNot(x)) → x -- Abs(Abs(x)) → Abs(x) -- Square(Sqrt(x)) → x (for nonnegative x, in practice applied if domain not violated) -- Sqrt(Square(x)) → Abs(x) -- Equal(x, x) → True -- NotEqual(x, x) → False -- Less(x, x) → False -- Greater(x, x) → False -- LessEqual(x, x) → True -- GreaterEqual(x, x) → True -- And(x, True) → x ; And(True, x) → x -- Or(x, False) → x ; Or(False, x) → x -- Select(cond, x, x) → x - -Complexity: ------------ -- Time: O(N) per iteration for N nodes, typically converges in few iterations. -- Space: O(1) auxiliary space per pattern match. - -Example: --------- -Example 1 - Add zero: - Original: y = Add(x, Const(0)) - Optimized: y = x - -Example 2 - Double negation: - Original: y = Neg(Neg(x)) - Optimized: y = x - -Example 3 - Compare equal: - Original: y = Equal(a, a) - Optimized: y = Const(True) - -Relationships: --------------- -- Runs after `ConstantFoldPass` (to fold constants before simplifying forms). -- Runs before `IdentityEliminationPass` (to reduce cases like Identity(Add(x,0))). -- Helps `CSEPass` by producing simpler, more canonical expressions. -""" - -from __future__ import annotations - -from graph_optimizer.core import ( - Op, - PassRegistry, - PatternRewritePass, - Any, - RewriteResult, -) -from graph_optimizer.utils.graph_utils import create_node, create_const_node -from graph_optimizer.utils.logger import logger as logging -import numpy as np - - -@PassRegistry.register("algebraic_simplify", opt_level=1, priority=7) -class AlgebraicSimplifyPass(PatternRewritePass): - """ - Applies algebraic identities to simplify expressions. - """ - - def __init__(self): - # We'll handle multiple patterns manually in _rewrite - pattern = Any(alias="op") # fallback, we check inside - super().__init__(pattern, self._rewrite, name="AlgebraicSimplify") - - def _rewrite(self, match, optimizer): - node = match.matched_nodes["op"] - op_type = node.op - inputs = list(node.input) - name = node.name - - def _mapped_result(target_name): - return RewriteResult(new_nodes=[], node_mapping={name: target_name}) - - def _new_node_result(new_node): - return RewriteResult( - new_nodes=[new_node], node_mapping={name: new_node.name} - ) - - # Helper to create True/False const - def _bool_const(val): - return _new_node_result( - create_const_node(name + "_bool", value=val, dtype="bool", shape=[]) - ) - - # Helper to get node object ignoring output index - def _get_node(name): - real_name = name.split(":")[0] - return optimizer.nodes.get(real_name) - - # Helper to check if a node is Const with given value (broadcast-safe) - def _is_const(node_name, value): - node = _get_node(node_name) - if node is None: - return False - if node.op != "Const": - return False - val = optimizer.get_node_attr(node, "value") - # Check if all elements are equal to the target value - return np.all(np.equal(val, value)) - - # Helper to get shape of a node - def _get_shape(node_name): - node = _get_node(node_name) - if node is None: - return None - # Check for shape attribute (Placeholder, etc.) - if "shape" in node.attr: - return [d.size for d in node.attr["shape"].shape.dim] - # Check for Const value shape - if node.op == "Const" and "value" in node.attr: - tensor = node.attr["value"].tensor - if tensor.HasField("tensor_shape"): - return [d.size for d in tensor.tensor_shape.dim] - return None - - # Helper to check if a node is definitely scalar - def _is_scalar(node_name): - shape = _get_shape(node_name) - return shape == [] - - # Helper to compute broadcast shape of two shapes - def _get_broadcast_shape(s1, s2): - if s1 is None or s2 is None: - return None - if s1 == s2: - return s1 - if not s1: - return s2 - if not s2: - return s1 - - # Simple broadcasting logic - len1, len2 = len(s1), len(s2) - max_len = max(len1, len2) - result = [] - for i in range(max_len): - d1 = s1[len1 - 1 - i] if i < len1 else 1 - d2 = s2[len2 - 1 - i] if i < len2 else 1 - if d1 == d2: - result.append(d1) - elif d1 == 1: - result.append(d2) - elif d2 == 1: - result.append(d1) - else: - return None # Incompatible - return result[::-1] - - # Helper to check if simplification is shape-preserving - def _is_shape_preserving(source_shape, target_shape): - # If both are unknown, assume it's safe (common in simple tests) - if source_shape is None and target_shape is None: - return True - if source_shape is None or target_shape is None: - return False - return source_shape == target_shape - - # Rule: Add(x, 0) or Add(0, x) - if op_type == "Add": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - - if _is_const(left, 0) and _is_shape_preserving(s_res, s_right): - return _mapped_result(right) - if _is_const(right, 0) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # Add(x, Neg(x)) -> 0 or Add(Neg(x), x) -> 0 - # Note: This is a simplified check for Neg(x) - for l, r in [(left, right), (right, left)]: - rn = _get_node(r) - if rn and rn.op == "Neg" and rn.input[0] == l: - s = _get_shape(l) - if s is not None: - source = _get_node(l) - dtype = source.attr.get("dtype", "float32") if source else "float32" - return _new_node_result( - create_const_node(name + "_zero", value=0, dtype=dtype, shape=s) - ) - - # Rule: Sub(x, 0) → x - if op_type == "Sub": - left, right = inputs[0], inputs[1] - if _is_const(right, 0) and ( - _is_scalar(right) or _get_shape(right) == _get_shape(left) - ): - return _mapped_result(left) - # Sub(x, x) → 0 - if left == right: - s = _get_shape(left) - if s is not None: - source = _get_node(left) - dtype = source.attr.get("dtype", "float32") if source else "float32" - return _new_node_result( - create_const_node(name + "_zero", value=0, dtype=dtype, shape=s) - ) - - # Rule: Mul(x, 1) or Mul(1, x) - if op_type == "Mul": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - - if _is_const(left, 1) and _is_shape_preserving(s_res, s_right): - return _mapped_result(right) - if _is_const(right, 1) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # Mul(x, 0) → 0 - if _is_const(left, 0) or _is_const(right, 0): - if s_res is not None: - source_name = right if _is_const(left, 0) else left - source = _get_node(source_name) - dtype = source.attr.get("dtype", "float32") if source else "float32" - return _new_node_result( - create_const_node( - name + "_zero", value=0, dtype=dtype, shape=s_res - ) - ) - # Mul(x, x) -> Square(x) - if left == right: - return _new_node_result( - create_node("Square", name + "_sq", inputs=[left]) - ) - - # Rule: Div(x, 1) → x - if op_type == "Div": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - if _is_const(right, 1) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # Div(x, x) -> 1 - if left == right: - s = _get_shape(left) - if s is not None: - source = _get_node(left) - dtype = source.attr.get("dtype", "float32") if source else "float32" - return _new_node_result( - create_const_node(name + "_one", value=1, dtype=dtype, shape=s) - ) - - # Rule: Neg(Neg(x)) → x - if op_type == "Neg": - inp = _get_node(inputs[0]) - if inp and inp.op == "Neg": - return _mapped_result(inp.input[0]) - - # Rule: LogicalNot(LogicalNot(x)) → x - if op_type == "LogicalNot": - inp = _get_node(inputs[0]) - if inp and inp.op == "LogicalNot": - return _mapped_result(inp.input[0]) - - # Rule: Abs(Abs(x)) → Abs(x) - if op_type == "Abs": - inp = _get_node(inputs[0]) - if inp and inp.op == "Abs": - orig = _get_node(inp.input[0]) - if orig: - return _new_node_result( - create_node("Abs", name + "_abs", inputs=[orig.name]) - ) - - # Rule: Square(Sqrt(x)) → x (domain assumed ok) - if op_type == "Square": - inp = _get_node(inputs[0]) - if inp and inp.op == "Sqrt": - return _mapped_result(inp.input[0]) - - # Rule: Sqrt(Square(x)) → Abs(x) - if op_type == "Sqrt": - inp = _get_node(inputs[0]) - if inp and inp.op == "Square": - orig = _get_node(inp.input[0]) - if orig: - return _new_node_result( - create_node("Abs", name + "_abs", inputs=[orig.name]) - ) - - # Rule: Pow(x, 1) -> x - if op_type == "Pow": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - if _is_const(right, 1) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # Pow(x, 2) -> Square(x) - if _is_const(right, 2) and _is_shape_preserving(s_res, s_left): - return _new_node_result( - create_node("Square", name + "_sq", inputs=[left]) - ) - - # Helper for comparison results - def _comparison_const(val): - # Equal(x, x) -> True should have same shape as x (or broadcasted shape) - # If x is [2, 2], result is [2, 2] of True - s = _get_shape(inputs[0]) - if s is None: - return None # Safer to skip if shape unknown - return _new_node_result( - create_const_node(name + "_bool", value=val, dtype="bool", shape=s) - ) - - # Rule: Equal(x, x) → True - if op_type == "Equal": - left, right = inputs[0], inputs[1] - if left == right: - return _comparison_const(True) - - # Rule: NotEqual(x, x) → False - if op_type == "NotEqual": - left, right = inputs[0], inputs[1] - if left == right: - return _comparison_const(False) - - # Rule: Less(x, x) → False ; Greater(x, x) → False - if op_type in ("Less", "Greater") and inputs[0] == inputs[1]: - return _comparison_const(False) - - # Rule: LessEqual(x, x) → True ; GreaterEqual(x, x) → True - if op_type in ("LessEqual", "GreaterEqual") and inputs[0] == inputs[1]: - return _comparison_const(True) - - # Rule: And(x, True) → x ; And(True, x) → x - if op_type == "LogicalAnd": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - - if _is_const(left, True) and _is_shape_preserving(s_res, s_right): - return _mapped_result(right) - if _is_const(right, True) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # LogicalAnd(x, x) -> x - if left == right: - return _mapped_result(left) - # LogicalAnd(x, False) -> False - if _is_const(left, False) or _is_const(right, False): - if s_res is not None: - return _new_node_result( - create_const_node(name + "_bool", value=False, dtype="bool", shape=s_res) - ) - - # Rule: Or(x, False) → x ; Or(False, x) → x - if op_type == "LogicalOr": - left, right = inputs[0], inputs[1] - s_left, s_right = _get_shape(left), _get_shape(right) - s_res = _get_broadcast_shape(s_left, s_right) - - if _is_const(left, False) and _is_shape_preserving(s_res, s_right): - return _mapped_result(right) - if _is_const(right, False) and _is_shape_preserving(s_res, s_left): - return _mapped_result(left) - # LogicalOr(x, x) -> x - if left == right: - return _mapped_result(left) - # LogicalOr(x, True) -> True - if _is_const(left, True) or _is_const(right, True): - if s_res is not None: - return _new_node_result( - create_const_node(name + "_bool", value=True, dtype="bool", shape=s_res) - ) - - # Rule: Select(cond, x, x) → x - if op_type == "Select": - if len(inputs) >= 3 and inputs[1] == inputs[2]: - return _mapped_result(inputs[1]) - - # Rule: Identity(x) -> x (bypass or collapse nested Identity) - if op_type == "Identity": - # Skip if protected/output node - if ( - hasattr(optimizer, "protected_nodes") - and name in optimizer.protected_nodes - ): - return None - # Skip ReadVariableOp - if "ReadVariableOp" in name: - return None - # Skip colocation constraint - if "_class" in node.attr: - return None - # Collapse nested Identity - inp_node = _get_node(inputs[0]) - if inp_node and inp_node.op == "Identity": - inner_input = inp_node.input[0] - new_node = create_node( - "Identity", name + "_collapsed", inputs=[inner_input] - ) - return _new_node_result(new_node) - # Bypass single Identity - return _mapped_result(inputs[0]) - - return None diff --git a/transforms/scalar/algebraic_simplify/__init__.py b/transforms/scalar/algebraic_simplify/__init__.py new file mode 100644 index 0000000..cf2f013 --- /dev/null +++ b/transforms/scalar/algebraic_simplify/__init__.py @@ -0,0 +1,4 @@ + +from .arithmetic import * +from .logical_comparison import * +from .other import * diff --git a/transforms/scalar/algebraic_simplify/arithmetic.py b/transforms/scalar/algebraic_simplify/arithmetic.py new file mode 100644 index 0000000..5aaf552 --- /dev/null +++ b/transforms/scalar/algebraic_simplify/arithmetic.py @@ -0,0 +1,352 @@ + +from __future__ import annotations + +from __future__ import annotations + +from graph_optimizer.core import ( + Any, + CommutativeOp, + Op, + PassRegistry, + PatternRewritePass, + RewriteResult, +) +from graph_optimizer.utils.graph_utils import create_node, create_const_node +from graph_optimizer.utils.graph_utils import get_node_shape, is_const_with_value, get_broadcast_shape, is_shape_preserving + +# ============================================================================== +# Helper functions from the original pass +# ============================================================================== + +def _get_node(optimizer, name): + real_name = name.split(":")[0] + return optimizer.nodes.get(real_name) + +# ============================================================================== +# Arithmetic Simplification Patterns +# ============================================================================== + +# Rule: Add(x, 0) -> x +@PassRegistry.register("arithmetic_simplify_add_zero", opt_level=1, priority=7) +class AddZeroPass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "Add", + Any(alias="x"), + Op("Const", alias="zero"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="AddZero") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + zero = match.matched_nodes["zero"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, zero, 0): + s_x = get_node_shape(optimizer, x) + s_zero = get_node_shape(optimizer, zero) + s_res = get_broadcast_shape(s_x, s_zero) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Add(x, Neg(x)) -> 0 +@PassRegistry.register("arithmetic_simplify_add_neg", opt_level=1, priority=7) +class AddNegPass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "Add", + Any(alias="x"), + Op("Neg", Any(alias="neg_x")), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="AddNeg") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + neg_x = match.matched_nodes["neg_x"] + root = match.matched_nodes["root"] + + if x.name == neg_x.name: + shape = get_node_shape(optimizer, x) + if shape is not None: + dtype = optimizer.get_node_attr(x, "dtype", "float32") + zero_const = create_const_node(root.name + "_zero", value=0, dtype=dtype, shape=shape) + return RewriteResult(new_nodes=[zero_const], node_mapping={root.name: zero_const.name}) + return None + +# Rule: Sub(x, 0) -> x +@PassRegistry.register("arithmetic_simplify_sub_zero", opt_level=1, priority=7) +class SubZeroPass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Sub", + Any(alias="x"), + Op("Const", alias="zero"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="SubZero") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + zero = match.matched_nodes["zero"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, zero, 0): + s_x = get_node_shape(optimizer, x) + s_zero = get_node_shape(optimizer, zero) + if s_x == s_zero or s_zero == []: + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Sub(x, x) -> 0 +@PassRegistry.register("arithmetic_simplify_sub_self", opt_level=1, priority=7) +class SubSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Sub", + Any(alias="x"), + Any(alias="y"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="SubSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + + if x.name == y.name: + shape = get_node_shape(optimizer, x) + if shape is not None: + dtype = optimizer.get_node_attr(x, "dtype", "float32") + zero_const = create_const_node(root.name + "_zero", value=0, dtype=dtype, shape=shape) + return RewriteResult(new_nodes=[zero_const], node_mapping={root.name: zero_const.name}) + return None + +# Rule: Mul(x, 1) -> x +@PassRegistry.register("arithmetic_simplify_mul_one", opt_level=1, priority=7) +class MulOnePass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "Mul", + Any(alias="x"), + Op("Const", alias="one"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="MulOne") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + one = match.matched_nodes["one"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, one, 1): + s_x = get_node_shape(optimizer, x) + s_one = get_node_shape(optimizer, one) + s_res = get_broadcast_shape(s_x, s_one) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Mul(x, 0) -> 0 +@PassRegistry.register("arithmetic_simplify_mul_zero", opt_level=1, priority=7) +class MulZeroPass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "Mul", + Any(alias="x"), + Op("Const", alias="zero"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="MulZero") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + zero = match.matched_nodes["zero"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, zero, 0): + s_x = get_node_shape(optimizer, x) + s_zero = get_node_shape(optimizer, zero) + s_res = get_broadcast_shape(s_x, s_zero) + if s_res is not None: + dtype = optimizer.get_node_attr(x, "dtype", "float32") + zero_const = create_const_node(root.name + "_zero", value=0, dtype=dtype, shape=s_res) + return RewriteResult(new_nodes=[zero_const], node_mapping={root.name: zero_const.name}) + return None + +# Rule: Mul(x, x) -> Square(x) +@PassRegistry.register("arithmetic_simplify_mul_self", opt_level=1, priority=7) +class MulSelfToSquarePass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Mul", + Any(alias="x"), + Any(alias="y"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="MulSelfToSquare") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + + if x.name == y.name: + square_node = create_node("Square", root.name + "_sq", inputs=[x.name]) + return RewriteResult(new_nodes=[square_node], node_mapping={root.name: square_node.name}) + return None + +# Rule: Div(x, 1) -> x +@PassRegistry.register("arithmetic_simplify_div_one", opt_level=1, priority=7) +class DivOnePass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Div", + Any(alias="x"), + Op("Const", alias="one"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="DivOne") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + one = match.matched_nodes["one"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, one, 1): + s_x = get_node_shape(optimizer, x) + s_one = get_node_shape(optimizer, one) + s_res = get_broadcast_shape(s_x, s_one) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Div(x, x) -> 1 +@PassRegistry.register("arithmetic_simplify_div_self", opt_level=1, priority=7) +class DivSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Div", + Any(alias="x"), + Any(alias="y"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="DivSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + + if x.name == y.name: + shape = get_node_shape(optimizer, x) + if shape is not None: + dtype = optimizer.get_node_attr(x, "dtype", "float32") + one_const = create_const_node(root.name + "_one", value=1, dtype=dtype, shape=shape) + return RewriteResult(new_nodes=[one_const], node_mapping={root.name: one_const.name}) + return None + +# Rule: Pow(x, 1) -> x +@PassRegistry.register("arithmetic_simplify_pow_one", opt_level=1, priority=7) +class PowOnePass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Pow", + Any(alias="x"), + Op("Const", alias="one"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="PowOne") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + one = match.matched_nodes["one"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, one, 1): + s_x = get_node_shape(optimizer, x) + s_one = get_node_shape(optimizer, one) + s_res = get_broadcast_shape(s_x, s_one) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Pow(x, 2) -> Square(x) +@PassRegistry.register("arithmetic_simplify_pow_two", opt_level=1, priority=7) +class PowTwoToSquarePass(PatternRewritePass): + def __init__(self): + self.pattern = Op( + "Pow", + Any(alias="x"), + Op("Const", alias="two"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="PowTwoToSquare") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + two = match.matched_nodes["two"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, two, 2): + s_x = get_node_shape(optimizer, x) + s_two = get_node_shape(optimizer, two) + s_res = get_broadcast_shape(s_x, s_two) + if is_shape_preserving(s_res, s_x): + square_node = create_node("Square", root.name + "_sq", inputs=[x.name]) + return RewriteResult(new_nodes=[square_node], node_mapping={root.name: square_node.name}) + return None + +# Rule: Square(Sqrt(x)) -> x +@PassRegistry.register("arithmetic_simplify_square_sqrt", opt_level=1, priority=7) +class SquareSqrtPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Square", Op("Sqrt", Any(alias="x"), alias="sqrt"), alias="root") + super().__init__(self.pattern, self._rewrite, name="SquareSqrt") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + +# Rule: Sqrt(Square(x)) -> Abs(x) +@PassRegistry.register("arithmetic_simplify_sqrt_square", opt_level=1, priority=7) +class SqrtSquareToAbsPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Sqrt", Op("Square", Any(alias="x"), alias="square"), alias="root") + super().__init__(self.pattern, self._rewrite, name="SqrtSquareToAbs") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + abs_node = create_node("Abs", root.name + "_abs", inputs=[x.name]) + return RewriteResult(new_nodes=[abs_node], node_mapping={root.name: abs_node.name}) + +# Rule: Mul(x, -1) -> Neg(x) +@PassRegistry.register("arithmetic_simplify_mul_neg_one", opt_level=1, priority=7) +class MulNegOneToNegPass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "Mul", + Any(alias="x"), + Op("Const", alias="neg_one"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="MulNegOneToNeg") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + neg_one = match.matched_nodes["neg_one"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, neg_one, -1): + s_x = get_node_shape(optimizer, x) + s_neg_one = get_node_shape(optimizer, neg_one) + s_res = get_broadcast_shape(s_x, s_neg_one) + if is_shape_preserving(s_res, s_x): + neg_node = create_node("Neg", root.name + "_neg", inputs=[x.name]) + return RewriteResult(new_nodes=[neg_node], node_mapping={root.name: neg_node.name}) + return None diff --git a/transforms/scalar/algebraic_simplify/logical_comparison.py b/transforms/scalar/algebraic_simplify/logical_comparison.py new file mode 100644 index 0000000..4f64b04 --- /dev/null +++ b/transforms/scalar/algebraic_simplify/logical_comparison.py @@ -0,0 +1,235 @@ + +from __future__ import annotations + +from __future__ import annotations + +from graph_optimizer.core import ( + Any, + CommutativeOp, + Op, + PassRegistry, + PatternRewritePass, + RewriteResult, +) +from graph_optimizer.utils.graph_utils import create_const_node, get_node_shape, is_const_with_value, get_broadcast_shape, is_shape_preserving + +# ============================================================================== +# Helper functions +# ============================================================================== + +def _comparison_const(optimizer, root_node, value): + """Helper to create a boolean constant with the same shape as the input.""" + # The result of a comparison should have a shape determined by broadcasting the inputs. + # For self-comparisons (e.g., Equal(x,x)), the shape is simply the shape of x. + input_node = optimizer.nodes.get(root_node.input[0].split(":")[0]) + if input_node: + shape = get_node_shape(optimizer, input_node) + if shape is not None: + const_node = create_const_node(root_node.name + "_bool", value=value, dtype="bool", shape=shape) + return RewriteResult(new_nodes=[const_node], node_mapping={root_node.name: const_node.name}) + return None + +# ============================================================================== +# Logical and Comparison Simplification Patterns +# ============================================================================== + +# Rule: LogicalNot(LogicalNot(x)) -> x +@PassRegistry.register("logical_simplify_double_not", opt_level=1, priority=7) +class DoubleNotPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("LogicalNot", Op("LogicalNot", Any(alias="x"), alias="inner_not"), alias="root") + super().__init__(self.pattern, self._rewrite, name="DoubleNot") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + +# Rule: Equal(x, x) -> True +@PassRegistry.register("comparison_simplify_equal_self", opt_level=1, priority=7) +class EqualSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Equal", Any(alias="x"), Any(alias="y"), alias="root") + super().__init__(self.pattern, self._rewrite, name="EqualSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + if x.name == y.name: + return _comparison_const(optimizer, root, True) + return None + +# Rule: NotEqual(x, x) -> False +@PassRegistry.register("comparison_simplify_not_equal_self", opt_level=1, priority=7) +class NotEqualSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("NotEqual", Any(alias="x"), Any(alias="y"), alias="root") + super().__init__(self.pattern, self._rewrite, name="NotEqualSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + if x.name == y.name: + return _comparison_const(optimizer, root, False) + return None + +# Rules: Less(x, x) -> False, Greater(x, x) -> False +@PassRegistry.register("comparison_simplify_less_greater_self", opt_level=1, priority=7) +class LessGreaterSelfPass(PatternRewritePass): + def __init__(self): + # This pattern is a bit broader, we'll check the op type inside + self.pattern = Any(alias="root") + super().__init__(self.pattern, self._rewrite, name="LessGreaterSelf") + + def _rewrite(self, match, optimizer): + root = match.matched_nodes["root"] + if root.op in ("Less", "Greater") and len(root.input) == 2 and root.input[0] == root.input[1]: + return _comparison_const(optimizer, root, False) + return None + +# Rules: LessEqual(x, x) -> True, GreaterEqual(x, x) -> True +@PassRegistry.register("comparison_simplify_le_ge_self", opt_level=1, priority=7) +class LessEqualGreaterEqualSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Any(alias="root") + super().__init__(self.pattern, self._rewrite, name="LessEqualGreaterEqualSelf") + + def _rewrite(self, match, optimizer): + root = match.matched_nodes["root"] + if root.op in ("LessEqual", "GreaterEqual") and len(root.input) == 2 and root.input[0] == root.input[1]: + return _comparison_const(optimizer, root, True) + return None + +# Rule: LogicalAnd(x, True) -> x +@PassRegistry.register("logical_simplify_and_true", opt_level=1, priority=7) +class LogicalAndTruePass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "LogicalAnd", + Any(alias="x"), + Op("Const", alias="true_const"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="LogicalAndTrue") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + true_const = match.matched_nodes["true_const"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, true_const, True): + s_x = get_node_shape(optimizer, x) + s_true = get_node_shape(optimizer, true_const) + s_res = get_broadcast_shape(s_x, s_true) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: LogicalAnd(x, x) -> x +@PassRegistry.register("logical_simplify_and_self", opt_level=1, priority=7) +class LogicalAndSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("LogicalAnd", Any(alias="x"), Any(alias="y"), alias="root") + super().__init__(self.pattern, self._rewrite, name="LogicalAndSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + if x.name == y.name: + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: LogicalAnd(x, False) -> False +@PassRegistry.register("logical_simplify_and_false", opt_level=1, priority=7) +class LogicalAndFalsePass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "LogicalAnd", + Any(alias="x"), + Op("Const", alias="false_const"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="LogicalAndFalse") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + false_const = match.matched_nodes["false_const"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, false_const, False): + s_x = get_node_shape(optimizer, x) + s_false = get_node_shape(optimizer, false_const) + s_res = get_broadcast_shape(s_x, s_false) + if s_res is not None: + new_const = create_const_node(root.name + "_bool", value=False, dtype="bool", shape=s_res) + return RewriteResult(new_nodes=[new_const], node_mapping={root.name: new_const.name}) + return None + +# Rule: LogicalOr(x, False) -> x +@PassRegistry.register("logical_simplify_or_false", opt_level=1, priority=7) +class LogicalOrFalsePass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "LogicalOr", + Any(alias="x"), + Op("Const", alias="false_const"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="LogicalOrFalse") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + false_const = match.matched_nodes["false_const"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, false_const, False): + s_x = get_node_shape(optimizer, x) + s_false = get_node_shape(optimizer, false_const) + s_res = get_broadcast_shape(s_x, s_false) + if is_shape_preserving(s_res, s_x): + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: LogicalOr(x, x) -> x +@PassRegistry.register("logical_simplify_or_self", opt_level=1, priority=7) +class LogicalOrSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("LogicalOr", Any(alias="x"), Any(alias="y"), alias="root") + super().__init__(self.pattern, self._rewrite, name="LogicalOrSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + if x.name == y.name: + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: LogicalOr(x, True) -> True +@PassRegistry.register("logical_simplify_or_true", opt_level=1, priority=7) +class LogicalOrTruePass(PatternRewritePass): + def __init__(self): + self.pattern = CommutativeOp( + "LogicalOr", + Any(alias="x"), + Op("Const", alias="true_const"), + alias="root" + ) + super().__init__(self.pattern, self._rewrite, name="LogicalOrTrue") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + true_const = match.matched_nodes["true_const"] + root = match.matched_nodes["root"] + + if is_const_with_value(optimizer, true_const, True): + s_x = get_node_shape(optimizer, x) + s_true = get_node_shape(optimizer, true_const) + s_res = get_broadcast_shape(s_x, s_true) + if s_res is not None: + new_const = create_const_node(root.name + "_bool", value=True, dtype="bool", shape=s_res) + return RewriteResult(new_nodes=[new_const], node_mapping={root.name: new_const.name}) + return None diff --git a/transforms/scalar/algebraic_simplify/other.py b/transforms/scalar/algebraic_simplify/other.py new file mode 100644 index 0000000..dbb9d70 --- /dev/null +++ b/transforms/scalar/algebraic_simplify/other.py @@ -0,0 +1,88 @@ + +from __future__ import annotations + +from __future__ import annotations + +from graph_optimizer.core import ( + Any, + Op, + PassRegistry, + PatternRewritePass, + RewriteResult, +) +from graph_optimizer.utils.graph_utils import create_node + +# ============================================================================== +# Miscellaneous Simplification Patterns +# ============================================================================== + +# Rule: Neg(Neg(x)) -> x +@PassRegistry.register("other_simplify_double_neg", opt_level=1, priority=7) +class DoubleNegPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Neg", Op("Neg", Any(alias="x"), alias="inner_neg"), alias="root") + super().__init__(self.pattern, self._rewrite, name="DoubleNeg") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + +# Rule: Abs(Abs(x)) -> Abs(x) +@PassRegistry.register("other_simplify_double_abs", opt_level=1, priority=7) +class DoubleAbsPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Abs", Op("Abs", Any(alias="x"), alias="inner_abs"), alias="root") + super().__init__(self.pattern, self._rewrite, name="DoubleAbs") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + # The result should be a new Abs node pointing to the original input x + new_abs = create_node("Abs", root.name + "_abs", inputs=[x.name]) + return RewriteResult(new_nodes=[new_abs], node_mapping={root.name: new_abs.name}) + +# Rule: Select(cond, x, x) -> x +@PassRegistry.register("other_simplify_select_self", opt_level=1, priority=7) +class SelectSelfPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Select", Any(alias="cond"), Any(alias="x"), Any(alias="y"), alias="root") + super().__init__(self.pattern, self._rewrite, name="SelectSelf") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + y = match.matched_nodes["y"] + root = match.matched_nodes["root"] + + if x.name == y.name: + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) + return None + +# Rule: Identity(x) -> x (Bypass) +# This is a very common and important simplification. +@PassRegistry.register("other_simplify_identity_bypass", opt_level=1, priority=1) # High priority +class IdentityBypassPass(PatternRewritePass): + def __init__(self): + self.pattern = Op("Identity", Any(alias="x"), alias="root") + super().__init__(self.pattern, self._rewrite, name="IdentityBypass") + + def _rewrite(self, match, optimizer): + x = match.matched_nodes["x"] + root = match.matched_nodes["root"] + + # Safety checks from the original implementation + if ( + hasattr(optimizer, "protected_nodes") + and root.name in optimizer.protected_nodes + ): + return None + if "ReadVariableOp" in root.name: + return None + if "_class" in root.attr: + return None + + # Check if the input is another Identity node (handled by iterative nature of passes) + # If we map Identity -> its input, and its input is also an Identity, the next iteration + # of this pass will handle it. No need for special nested logic. + + return RewriteResult(new_nodes=[], node_mapping={root.name: x.name}) diff --git a/utils/graph_utils.py b/utils/graph_utils.py index 211e549..6b51c78 100644 --- a/utils/graph_utils.py +++ b/utils/graph_utils.py @@ -453,3 +453,80 @@ def build_consumer_index(graph_def: tf.GraphDef) -> Dict[str, list]: base_name = extract_base_name(input_name) consumers[base_name].append(node.name) return consumers + + +def get_node_shape(optimizer, node_or_name): + """Returns the output shape of a node, if available.""" + node = ( + optimizer.nodes.get(node_or_name) + if isinstance(node_or_name, str) + else node_or_name + ) + if node is None: + return None + + if "_output_shapes" in node.attr: + shapes = node.attr["_output_shapes"].list.shape + if shapes: + return [d.size for d in shapes[0].dim] + + if "shape" in node.attr: + return [d.size for d in node.attr["shape"].shape.dim] + + if node.op == "Const" and "value" in node.attr: + tensor = node.attr["value"].tensor + if tensor.HasField("tensor_shape"): + return [d.size for d in tensor.tensor_shape.dim] + + return None + + +def is_const_with_value(optimizer, node_or_name, value): + """Checks if a node is a Const with a specific value.""" + node = ( + optimizer.nodes.get(node_or_name) + if isinstance(node_or_name, str) + else node_or_name + ) + if node is None or node.op != "Const": + return False + + try: + val = tensor_util.MakeNdarray(node.attr["value"].tensor) + return np.all(np.equal(val, value)) + except Exception: + return False + + +def get_broadcast_shape(s1, s2): + """Computes the broadcast shape of two shapes.""" + if s1 is None or s2 is None: + return None + if s1 == s2: + return s1 + + if not s1: return s2 + if not s2: return s1 + + len1, len2 = len(s1), len(s2) + max_len = max(len1, len2) + result = [] + for i in range(max_len): + d1 = s1[len1 - 1 - i] if i < len1 else 1 + d2 = s2[len2 - 1 - i] if i < len2 else 1 + if d1 == d2 or d1 == -1 or d2 == -1: + result.append(d1 if d1 != 1 else d2) + elif d1 == 1: + result.append(d2) + elif d2 == 1: + result.append(d1) + else: + return None + return result[::-1] + + +def is_shape_preserving(source_shape, target_shape): + """Checks if a simplification is shape-preserving.""" + if source_shape is None or target_shape is None: + return True + return source_shape == target_shape