[Relax] Refactor PatternRewriter into separate Block/Expr mutators #16730
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Prior to this commit, the `PatternRewriter` mutator handled pattern rewriting at either the expression level (`rewrite_call`) or the dataflow block level (`rewrite_bindings`). These two functionalities had different external APIs, defined diffierent member variables, and visited different IR nodes. In effect, it had two entirely independent implementations, which just happened to be implemented within the same class. This commit refactors the single `PatternRewriter` mutator into separate `BlockPatternRewriter` and `ExprPatternRewriter` mutators.
|
related to the class but not this PR necessarily would be good to think about complexity reduction for rewrite_bindings, which is source of slowdown in our overall compilation |
|
Agreed, and I think this PR helps move in that direction. By separating out implementation of I wanted to separate the two implementations before making an improvement specific to |
slyubomirsky
approved these changes
Mar 25, 2024
thaisacs
pushed a commit
to thaisacs/tvm
that referenced
this pull request
Apr 3, 2024
…pache#16730) Prior to this commit, the `PatternRewriter` mutator handled pattern rewriting at either the expression level (`rewrite_call`) or the dataflow block level (`rewrite_bindings`). These two functionalities had different external APIs, defined diffierent member variables, and visited different IR nodes. In effect, it had two entirely independent implementations, which just happened to be implemented within the same class. This commit refactors the single `PatternRewriter` mutator into separate `BlockPatternRewriter` and `ExprPatternRewriter` mutators.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Jul 8, 2024
A follow-up to apache#16730. Now that the implementations for `rewrite_call` and `rewrite_bindings` are in separate classes, they can be further split out into separate files.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Jul 10, 2024
A follow-up to apache#16730. Now that the implementations for `rewrite_call` and `rewrite_bindings` are in separate classes, they can be further split out into separate files.
Lunderberg
added a commit
to Lunderberg/tvm
that referenced
this pull request
Jul 10, 2024
A follow-up to apache#16730. Now that the implementations for `rewrite_call` and `rewrite_bindings` are in separate classes, they can be further split out into separate files.
sunggg
pushed a commit
that referenced
this pull request
Jul 24, 2024
* [TVMScript][Bugfix] Normalize relax::If with function's TIR var Prior to this commit, the branches of `relax::If` were normalized using `EraseToWellDefinedInScope`, using a fresh variable scope. While this had the intended behavior of preventing variables defined in a single branch from being usable outside of the conditional, it also caused the conditional's branches to treat function-scope symbolic variables as if they were undefined. This commit updates the `tvm::relax::Normalizer` so that `relax::If` is normalized within an inherited scope. This preserves the previous behavior for symbolic variables defined within a branch, but allows shapes within a branch to use symbolic variables defined outside of the branch. * [Relax] Canonicalize known symbolic shapes in Relax expressions Prior to this commit, known constants in Relax functions would be inlined by the `CanonicalizeBindings` pass, but only if they appeared as Relax expressions (e.g. `R.const` or `R.prim_value`). Known constants that appeared as TIR variables (e.g. symbolic shapes) would be kept as dynamic parameters, even if they were known at compile time. This commit updates the `CanonicalizeBindings` pass to identify known values of symbolic shapes, and to use these known values in shape expressions. * [Relax][Refactor] Reorganize pattern-matching A follow-up to #16730. Now that the implementations for `rewrite_call` and `rewrite_bindings` are in separate classes, they can be further split out into separate files. * [Relax][Refactor] Implement Rewriter class for pattern-rewrite Prior to this commit, the pattern to be matched and the rewrite to be performed were provided as separate arguments. This commit introduces a new class `ExprRewriter`, which contains both parts. This abstraction will make it easier to combine multiple different rewrite rules, applying them in a single pass. * lint fixes * Remove unnecessary change which broke a unit test * lint fix for import order * Add docstrings * lint fix * Lint fix * lint fixes * lint fix * Update based on review comments * Add test case for matching against arbitrary dtype * Fix breakage in unit tests One unit test that had been relying on invalid shape propagation. Another unit test that required constructed an ill-formed output to test against. * Updated base class name from ExprRewriter to PatternMatchingRewriter * lint fix
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Prior to this commit, the
PatternRewritermutator handled pattern rewriting at either the expression level (rewrite_call) or the dataflow block level (rewrite_bindings). These two functionalities had different external APIs, defined diffierent member variables, and visited different IR nodes. In effect, it had two entirely independent implementations, which just happened to be implemented within the same class.This commit refactors the single
PatternRewritermutator into separateBlockPatternRewriterandExprPatternRewritermutators.