[XLA:cpu] allow overriding of MLIR pass pipeline via debug flags #8543
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
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_598624926
branch
3 times, most recently
from
f4eae4b to
7fc319b
Compare
Use a backend-specific flag to allow overriding the MLIR pass pipeline when compiling to CPU via the new runtime pass. This is helpful to debug the pipeline, in particular the performance of the generated code, without having to recompile XLA every time. The feature is marked as experimental, hidden deep inside debug options and obviously disabled in absence of the option. Also update the signatures of `create_compilation_pipeline` and `create_specialization_pipeline` in `JitOptions` to return a status. This allows for a more graceful failure when MLIR pass pipeline cannot be constructed. CPU pipeline constructor was already aborting with a fatal error, update that to returning internal error. This makes the JIT friendlier to library users such as JAX. PiperOrigin-RevId: 599480560
copybara-service
bot
force-pushed
the
test_598624926
branch
from
7fc319b to
f685679
Compare
copybara-service
bot
closed this pull request by merging all changes
into
main
in
f685679
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.
[XLA:cpu] allow overriding of MLIR pass pipeline via debug flags
Use a backend-specific flag to allow overriding the MLIR pass pipeline when
compiling to CPU via the new runtime pass. This is helpful to debug the
pipeline, in particular the performance of the generated code, without having
to recompile XLA every time. The feature is marked as experimental, hidden deep
inside debug options and obviously disabled in absence of the option.
Also update the signatures of
create_compilation_pipelineandcreate_specialization_pipelineinJitOptionsto return a status. Thisallows for a more graceful failure when MLIR pass pipeline cannot be
constructed. CPU pipeline constructor was already aborting with a fatal error,
update that to returning internal error. This makes the JIT friendlier to
library users such as JAX.