Skip to content

Conversation

paul0403
Copy link
Member

@paul0403 paul0403 commented May 20, 2025

Context:
This work is based on #1027.

Now that we have migrated all the individual dialects, we should migrate the entire bufferization pipeline.

The Quantum dialect was migrated in #1686 .
The Catalyst dialect was migrated in #1708 .
The Gradient dialect was migrated in #1740 .

See more context in #1027.

Upstream changes in llvm were required for this bufferization update. As a result, the llvm version and mlir-hlo version were updated to

mhlo=25b008569f413d76cfa8f481f3a84e82b89c47f4
llvm=5f74671c85877e03622e8d308aee15ed73ccee7c

These are the versions tracked by jax 0.4.32.
These are the earliest jax-tagged versions with complete upstream bufferization changes.

Related GitHub Issues:
[sc-71487]

paul0403 added 4 commits May 21, 2025 12:05
llvm/llvm-project#107109
The core dialect conversion no longer tries to materialize ops on the fly during conversion,
but rather inserts unrealized conversion casts to be safe, and removes these inverse
cast pairs later after the entire conversion is done.
The old analysis is removed.

Personally I think this is fine, as we end our pipeline with --reconcil-unrealized-casts already.

Changes:
Catalyst/ConversionTest.mlir and Catalyst/MemrefLoadStoreLoweringTBAA.mlir: changing op order
Quantum/ConversionTest.mlir: changing op order; adding two unrealized conversion cast pairs
Gradient/ConversionTest.mlir: 5 unused Values are removed
Copy link
Contributor

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great, this looks good as is 👌 Just wondering about the other recommended passes in the bufferization docs.

Base automatically changed from paul0403/new_bufferize_gradient_dialect to main May 23, 2025 20:52
@paul0403
Copy link
Member Author

Since this one updates LLVM, let's check the wheels before merging. I will also clean the CI cache under this llvm commit before checking just for absolute certainty.

@paul0403 paul0403 added author:build-wheels Run the wheel building workflows on this Pull Request reviewer:require-wheels Pull Requests will need wheel building job successful before being merged and removed do-not-merge labels May 23, 2025
@paul0403 paul0403 removed reviewer:require-wheels Pull Requests will need wheel building job successful before being merged author:build-wheels Run the wheel building workflows on this Pull Request labels May 24, 2025
Copy link
Contributor

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀

@paul0403 paul0403 merged commit 120e5c4 into main May 26, 2025
39 of 40 checks passed
@paul0403 paul0403 deleted the paul0403/one-shot-bufferize-final branch May 26, 2025 18:19
paul0403 added a commit that referenced this pull request May 28, 2025
1. Moving ops around due to random ordering updates
2. Revert the inserted unrealized_conversion_casts in #1751
3. llvm.mlir.undef is deprecated, use llvm.mlir.poison instead
llvm/llvm-project#125629
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants