⚡ Bolt: Implement O(1) indexed matching for AlgebraicSimplify and ConstantFold

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2026-02-01 - O(1) Indexed Matching for General Passes
+**Learning:** Using a general `Any()` wildcard pattern in a rewrite pass that only targets specific operations is a performance anti-pattern. It forces the `PatternMatcher` to evaluate the rewriter for EVERY node in the graph (O(N)), bypassing the O(1) op-type index.
+**Action:** Always prefer registering specific `Op(op_type)` patterns. Upgraded `PatternRewritePass` to support multiple patterns to facilitate this for passes like `AlgebraicSimplify` and `ConstantFold`.

core.py

@@ -1516,11 +1516,29 @@ class PatternRewritePass(BasePass):
     for the actual pattern matching. Iterates until convergence (no more matches).
     """
 
-    def __init__(self, pattern, rewriter, name=None, optimizer_alias=None):
+    def __init__(
+        self,
+        pattern=None,
+        rewriter=None,
+        name=None,
+        optimizer_alias=None,
+        patterns=None,
+    ):
         # Use iterative mode - run until convergence
         super().__init__(name, optimizer_alias, iterative=True, max_iterations=100)
-        self.pattern = pattern
-        self.rewriter = trace_transformation(rewriter)
+
+        # Support both 'pattern' (single) and 'patterns' (list of Pattern or (Pattern, Rewriter))
+        self.patterns = []
+        if pattern:
+            self.patterns.append((pattern, trace_transformation(rewriter)))
+
+        if patterns:
+            for item in patterns:
+                if isinstance(item, tuple):
+                    self.patterns.append((item[0], trace_transformation(item[1])))
+                else:
+                    # Use common rewriter if only pattern is provided
+                    self.patterns.append((item, trace_transformation(rewriter)))
 
     def transform_once(
         self,
@@ -1534,9 +1552,10 @@ def transform_once(
         Returns:
             int: Number of changes made
         """
-        # Register the pattern (clear first to avoid duplicates)
+        # Register the patterns (clear first to avoid duplicates)
         optimizer.clear_transformations()
-        optimizer.add_transformation(self.pattern, self.rewriter)
+        for pattern, rewriter in self.patterns:
+            optimizer.add_transformation(pattern, rewriter)
 
         # Run one pattern matching iteration
         new_graph_def, changes = optimizer.match_patterns_once(

transforms/scalar/algebraic_simplify.py

@@ -87,9 +87,33 @@ class AlgebraicSimplifyPass(PatternRewritePass):
     """
 
     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")
+        # Use specific Op patterns for O(1) indexed matching instead of Any()
+        ops = [
+            "Add",
+            "Sub",
+            "Mul",
+            "Div",
+            "Neg",
+            "LogicalNot",
+            "Abs",
+            "Square",
+            "Sqrt",
+            "Pow",
+            "Equal",
+            "NotEqual",
+            "Less",
+            "Greater",
+            "LessEqual",
+            "GreaterEqual",
+            "LogicalAnd",
+            "LogicalOr",
+            "Select",
+            "Identity",
+        ]
+        patterns = [Op(op, alias="op") for op in ops]
+        super().__init__(
+            patterns=patterns, rewriter=self._rewrite, name="AlgebraicSimplify"
+        )
 
     def _rewrite(self, match, optimizer):
         node = match.matched_nodes["op"]

transforms/scalar/constant_fold.py

@@ -58,9 +58,55 @@ class ConstantFoldPass(PatternRewritePass):
     """
 
     def __init__(self):
-        # Matches any operation with all inputs as Const
-        pattern = Any(alias="op")
-        super().__init__(pattern, self._rewrite_constant_op, name="ConstantFold")
+        # Use specific Op patterns for O(1) indexed matching instead of Any()
+        ops = [
+            "Add",
+            "Mul",
+            "Sub",
+            "Div",
+            "Neg",
+            "Equal",
+            "NotEqual",
+            "Less",
+            "Greater",
+            "LessEqual",
+            "GreaterEqual",
+            "LogicalAnd",
+            "LogicalOr",
+            "LogicalNot",
+            "BitwiseAnd",
+            "BitwiseOr",
+            "BitwiseXor",
+            "Abs",
+            "Exp",
+            "Expm1",
+            "Log",
+            "Log1p",
+            "Sqrt",
+            "Pow",
+            "Rsqrt",
+            "Square",
+            "Sin",
+            "Cos",
+            "Tan",
+            "Asin",
+            "Acos",
+            "Atan",
+            "Atan2",
+            "Floor",
+            "Ceil",
+            "Round",
+            "Sign",
+            "Reshape",
+            "Transpose",
+            "ConcatV2",
+            "Select",
+            "Cast",
+        ]
+        patterns = [Op(op, alias="op") for op in ops]
+        super().__init__(
+            patterns=patterns, rewriter=self._rewrite_constant_op, name="ConstantFold"
+        )
 
     def _is_all_const(self, inputs, optimizer):
         """Check if all inputs are Const nodes.