diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 0000000..9abdc95 --- /dev/null +++ b/.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. diff --git a/core.py b/core.py index 25d8232..e444b10 100644 --- a/core.py +++ b/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( diff --git a/transforms/scalar/algebraic_simplify.py b/transforms/scalar/algebraic_simplify.py index 76c94af..e035919 100644 --- a/transforms/scalar/algebraic_simplify.py +++ b/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"] diff --git a/transforms/scalar/constant_fold.py b/transforms/scalar/constant_fold.py index 187cff4..e9fc7c1 100644 --- a/transforms/scalar/constant_fold.py +++ b/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.