diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 0000000..1e2f017 --- /dev/null +++ b/.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`. diff --git a/core.py b/core.py index 25d8232..b04ce93 100644 --- a/core.py +++ b/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( diff --git a/transforms/scalar/algebraic_simplify.py b/transforms/scalar/algebraic_simplify.py index 76c94af..39a5419 100644 --- a/transforms/scalar/algebraic_simplify.py +++ b/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"] diff --git a/transforms/scalar/constant_fold.py b/transforms/scalar/constant_fold.py index 187cff4..78c5864 100644 --- a/transforms/scalar/constant_fold.py +++ b/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.