[Frontend] Use custom lowering rules #1152
Merged
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
erick-xanadu
force-pushed
the
eochoa/2024-09-24/use-custom-lowering-rules-instead-of-modifying-mlir
branch
from
2dff6d1
to
584c15d
Compare
dime10
reviewed
Sep 24, 2024
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1152 +/- ##
==========================================
- Coverage 97.87% 97.86% -0.01%
==========================================
Files 76 76
Lines 10850 10810 -40
Branches 1283 1281 -2
==========================================
- Hits 10619 10579 -40
Misses 179 179
Partials 52 52 ☔ View full report in Codecov by Sentry. |
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Not yet, but I'll create an issue in JAX and maybe it will get some traction? |
dime10
approved these changes
Sep 25, 2024
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
erick-xanadu
deleted the
eochoa/2024-09-24/use-custom-lowering-rules-instead-of-modifying-mlir
branch
erick-xanadu
restored the
eochoa/2024-09-24/use-custom-lowering-rules-instead-of-modifying-mlir
branch
rauletorresc
pushed a commit
that referenced
this pull request
Oct 10, 2024
**Context:** Using [`mlir.register_lowering`](https://github.com/jax-ml/jax/blob/ae86ef16c7a03409cb444da3d477b2adb8134e6f/jax/_src/interpreters/mlir.py#L814-L825) [will modify the global variable `_lowerings` or `platform_specific_lowerings`](https://github.com/jax-ml/jax/blob/ae86ef16c7a03409cb444da3d477b2adb8134e6f/jax/_src/interpreters/mlir.py#L810-L812). To avoid this, JAX provides the[ `LoweringParameters`](https://github.com/jax-ml/jax/blob/ae86ef16c7a03409cb444da3d477b2adb8134e6f/jax/_src/interpreters/mlir.py#L630-L649) structure to pass custom lowering rules. When lowering rules are passed using [`LoweringParameters`, these global variables will not be rewritten.](https://github.com/jax-ml/jax/blob/ae86ef16c7a03409cb444da3d477b2adb8134e6f/jax/_src/interpreters/mlir.py#L1745-L1751) **Description of the Change:** Use LoweringParameters instead of register_lowering **Benefits:** Avoid rewriting global data structures from JAX. **Possible Drawbacks:** None. **Related GitHub Issues:** --------- Co-authored-by: David Ittah <dime10@users.noreply.github.com>
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.
Context: Using
mlir.register_loweringwill modify the global variable_loweringsorplatform_specific_lowerings. To avoid this, JAX provides theLoweringParametersstructure to pass custom lowering rules. When lowering rules are passed usingLoweringParameters, these global variables will not be rewritten.Description of the Change: Use LoweringParameters instead of register_lowering
Benefits: Avoid rewriting global data structures from JAX.
Possible Drawbacks: None.
Related GitHub Issues: