⚡ Bolt: [performance improvement] indexed pattern matching

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2026-02-03 - [Optimizing Pattern Matching with Indexed Lookups]
+**Learning:** Using a general `Any()` wildcard pattern at the root of a `PatternRewritePass` is a performance anti-pattern. It forces the `PatternMatcher` to evaluate the rewriter logic for every single node in the graph, resulting in O(N) complexity per iteration. By explicitly registering specific `Op` patterns, the matcher can leverage its O(1) `pattern_index` lookup, skipping nodes that don't match the required operation types.
+**Action:** Always prefer registering a list of specific `Op` patterns over a single `Any()` wildcard for passes that target a known set of operations. Update `PatternRewritePass` to support multiple patterns to facilitate this.

core.py

@@ -1516,11 +1516,46 @@ class PatternRewritePass(BasePass):
     for the actual pattern matching. Iterates until convergence (no more matches).
     """
 
-    def __init__(self, pattern, rewriter, name=None, optimizer_alias=None):
-        # Use iterative mode - run until convergence
+    def __init__(
+        self, pattern=None, rewriter=None, name=None, optimizer_alias=None, patterns=None
+    ):
+        """
+        Initialize a pattern rewrite pass.
+
+        Args:
+            pattern: Single Pattern object (optional if patterns is used)
+            rewriter: Rewriter function for the single pattern
+            name: Human-readable pass name
+            optimizer_alias: Short alias for node naming
+            patterns: List of Pattern objects or (Pattern, Rewriter) tuples.
+                     If Pattern objects are provided, the 'rewriter' argument
+                     will be used as the rewriter for all of them.
+        """
         super().__init__(name, optimizer_alias, iterative=True, max_iterations=100)
-        self.pattern = pattern
-        self.rewriter = trace_transformation(rewriter)
+
+        self._patterns = []
+        default_rewriter = trace_transformation(rewriter) if rewriter else None
+
+        if patterns is not None:
+            for item in patterns:
+                if isinstance(item, tuple):
+                    p, r = item
+                    self._patterns.append((p, trace_transformation(r)))
+                else:
+                    if default_rewriter is None:
+                        raise ValueError(
+                            "Rewriter must be provided if patterns contains Pattern objects"
+                        )
+                    self._patterns.append((item, default_rewriter))
+        elif pattern is not None and rewriter is not None:
+            self._patterns.append((pattern, default_rewriter))
+        else:
+            raise ValueError(
+                "Either (pattern and rewriter) or patterns must be provided"
+            )
+
+        # For backward compatibility and convenience
+        self.pattern, self.rewriter = self._patterns[0]
 
     def transform_once(
         self,
@@ -1534,9 +1569,10 @@ def transform_once(
         Returns:
             int: Number of changes made
         """
-        # Register the pattern (clear first to avoid duplicates)
+        # Register all patterns (clear first to avoid duplicates)
         optimizer.clear_transformations()
-        optimizer.add_transformation(self.pattern, self.rewriter)
+        for p, r in self._patterns:
+            optimizer.add_transformation(p, r)
 
         # Run one pattern matching iteration
         new_graph_def, changes = optimizer.match_patterns_once(

transforms/scalar/algebraic_simplify.py

@@ -87,9 +87,32 @@ 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")
+        # Explicitly register patterns for all supported operations
+        # This enables O(1) indexed matching instead of O(N) wildcard matching
+        supported_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 supported_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,56 @@ 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")
+        # Explicitly register patterns for all operations supported by constant folding
+        # This enables O(1) indexed matching instead of O(N) wildcard matching
+        supported_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 supported_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.