Add Shardy dialect support for JAX 0.8.2+ compatibility #22930
Conversation
This file contains hidden or 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
robtaylor
requested review from
ScottTodd,
benvanik and
stellaraccident
as code owners
Author
Areas that may need attention:
The implementation has been tested locally with JAX 0.8.2 on macOS ARM64, but additional review and testing is needed before this is ready for merge. |
robtaylor
force-pushed
the
add-shardy-dialect-support
branch
from
7092ca5 to
046c926
Compare
JAX 0.8.2+ uses the Shardy partitioner by default (jax_use_shardy_partitioner=True),
which emits MLIR bytecode containing the sdy (Shardy) dialect. Without proper
dialect registration, IREE fails with "dialect 'sdy' does not implement the
bytecode interface".
This PR adds:
1. Shardy submodule (openxla/shardy) with CMake build support
- build_tools/third_party/shardy/ contains CMake build files since
upstream Shardy only has Bazel
2. New IREE input plugin at compiler/plugins/input/Shardy/
- Registers the sdy dialect via IREE's plugin architecture
- Provides StripShardyDialect pass to remove sdy ops/attributes for
single-device execution (sdy ops are metadata-only sharding annotations)
3. New IREE_INPUT_SHARDY CMake option (ON by default)
- Enables/disables Shardy dialect support in the compiler
4. Test for Shardy integration (test_shardy.py)
- Verifies JAX works with Shardy enabled on IREE PJRT backends
Technical details:
- ShardySession plugin class with DefaultActivated policy
- Dialect registration via mlir::sdy::registerAllDialects()
- Input conversion pipeline strips sdy.sharding attributes
- Properly integrated via iree_compiler_register_plugin() to ensure
symbols are included in libIREECompiler
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.5 <[email protected]>
Signed-off-by: Rob Taylor <[email protected]>
robtaylor
force-pushed
the
add-shardy-dialect-support
branch
from
046c926 to
5253b62
Compare
Author
|
I've manually reviewed this PR, and it looks good to me. |
- Fix iterator invalidation bug in StripShardyDialect.cpp by collecting ops first then erasing in reverse order - Add warning for unexpected Shardy op patterns that can't be handled - Implement pass registration in registerShardyInputConversionPasses() - Add StableHLO dependency check in CMake - Update copyright years to 2025 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <[email protected]> Signed-off-by: Rob Taylor <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
JAX 0.8.2+ uses the Shardy partitioner by default (
jax_use_shardy_partitioner=True), which emits MLIR bytecode containing thesdy(Shardy) dialect. Without proper dialect registration, IREE fails with:This PR adds full Shardy dialect support to IREE.
Changes
third_party/shardy): Add openxla/shardy as a git submodulebuild_tools/third_party/shardy/): CMake build files since upstream Shardy only has Bazelcompiler/plugins/input/Shardy/):StripShardyDialectpass to remove sdy ops/attributes for single-device executioniree_compiler_register_plugin()to ensure symbols are includedIREE_INPUT_SHARDY(ON by default)test_shardy.pyverifies JAX works with Shardy enabledTechnical Details
The implementation follows IREE's plugin architecture patterns:
ShardySessionplugin class withDefaultActivatedpolicymlir::sdy::registerAllDialects()sdy.shardingattributes (metadata-only for single-device)libIREECompilerTest plan
test_shardy.pyadded to test runner🤖 Generated with Claude Code