Migrating to One-Shot Bufferization #1027
Closed
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
tzunghanjuang
force-pushed
the
bufferization-update
branch
2 times, most recently
from
2c43ec4 to
dce6cc2
Compare
tzunghanjuang
force-pushed
the
bufferization-update
branch
from
dce6cc2 to
6af3eef
Compare
Contributor
|
This is awesome! Keep up the good work and let me know if you need some help :) |
Contributor
|
Hello. You may have forgotten to update the changelog!
|
Co-authored-by: erick-xanadu <[email protected]>
Co-authored-by: erick-xanadu <[email protected]>
erick-xanadu
approved these changes
Oct 15, 2024
Contributor
erick-xanadu
left a comment
erick-xanadu
left a comment
There was a problem hiding this comment.
👍 @dime10 , I'll leave you some time to ask questions if you have any. If not let me know and I can merge.
Contributor
|
Adding do-not merge label to wait after the release. |
paul0403
added a commit
that referenced
this pull request
May 1, 2025
**Context:** This work is based on #1027 As part of the mlir update, the bufferization of the custom catalyst dialects need to migrate to the new one-shot bufferization interface, as opposed to the old pattern-rewrite style bufferization passes. See more context in #1027. As an example, here is how the new bufferization interface is used for mlir's core `arith` dialect: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h https://github.com/llvm/llvm-project/blob/7ee0097b486b31be8b9a1750b2cd47580efd9587/mlir/lib/Dialect/Arith/IR/ArithDialect.cpp#L54 **Description of the Change:** On the current mlir commit we track, both the old and new bufferization styles exist. The old pattern rewrite style is deprecated. To ease the workflow organization, we migrate one dialect at a time. This PR migrates the `Quantum` dialect's bufferization to the new one-shot interface. Note that the new one-shot interface is supposed to be called only once in the pipeline. However, because we haven't migrated all the dialects yet, we simply swap out the old `--quantum--bufferize` pass in-place, with the new one-shot bufferization pass running on the quantum dialect only. **Benefits:** Align with mlir practices; one step closer to updating mlir. [sc-71487] --------- Co-authored-by: Tzung-Han Juang <[email protected]> Co-authored-by: paul0403 <[email protected]> Co-authored-by: Paul <[email protected]> Co-authored-by: David Ittah <[email protected]>
paul0403
added a commit
that referenced
this pull request
May 12, 2025
**Context:** This work is based on #1027 . As part of the mlir update, the bufferization of the custom catalyst dialects need to migrate to the new one-shot bufferization interface, as opposed to the old pattern-rewrite style bufferization passes. See more context in #1027. The `Quantum` dialect was migrated in #1686 . **Description of the Change:** MIgrate `Catalyst` dialect to one-shot bufferization. **Benefits:** Align with mlir practices; one step closer to updating mlir. [sc-71487] --------- Co-authored-by: Tzung-Han Juang <[email protected]>
paul0403
added a commit
that referenced
this pull request
May 23, 2025
**Context:** This work is based on #1027 . As part of the mlir update, the bufferization of the custom catalyst dialects need to migrate to the new one-shot bufferization interface, as opposed to the old pattern-rewrite style bufferization passes. See more context in #1027. The `Quantum` dialect was migrated in #1686 . The `Catalyst` dialect was migrated in #1708 . Note that #1139 refactors the gradient dialect's bufferization into preprocess, bufferization, and postprocess. Only the middle bufferization stage is supposed to be replaced by one-shot bufferization. **Description of the Change:** Migrate `Gradient` dialect to one-shot bufferization. **Benefits:** Align with mlir practices; one step closer to updating mlir. [sc-71487] --------- Co-authored-by: Tzung-Han Juang <[email protected]>
paul0403
added a commit
that referenced
this pull request
May 26, 2025
…1751) **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] --------- Co-authored-by: Tzung-Han Juang <[email protected]>
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:
MLIR removed old bufferization (
Conversion-Based bufferization) passes and introduced newOne-Shot-Bufferizationpass. [Link]Note that
Ownership-based Buffer Deallocationwill be handled in a different PR.Only update llvm and mhlo.Jax sticks with 0.4.28.Description of the Change:
Newly Added Bufferization Interface:
Comparing implemented interface methods with BufferizationOpInterface.td.
Operation List:
https://github.com/PennyLaneAI/catalyst/blob/8d0324db672d507d651b46fd65e9058455a9bb03/mlir/lib/Bufferization.md
Major Updates:
one-shot-bufferizeinterface forcatalyst-bufferize,gradient-bufferize, andquantum-bufferize.declarePromisedInterfaces.FunctionOPInterfacelikeForwardOPandReverseOpneed to accessOneShotModuleBufferize.Gradient.ReturnOPisReturnLikenow to support bufferization.function-boundary-type-conversiontoidentity-layout-mapto avoid invalid type conversion betweenMemref<T>andMemref<T, Stride<...>>allow-return-allocs-from-loopsflag to avoidscp.forandscp.whileto avoid arg and yield types mismatchcopy-before-writeflagRequired LLVM and MHLO Patches:
OneShotModuleBufferizesupportFunctionOpInterfaceinstead ofFuncOP. [MLIR] MakeOneShotModuleBufferizeuseOpInterfacellvm/llvm-project#110322FunctionOpInterfacepatch to mhlo.FunctionArgTypeConverterFnunderBufferizationOptionstakesFunctionOpInterfaceinstead offunc::FuncOpopenxla/xla#17935Related GitHub Issues:
[sc-71487]
[Discourse discussion about migration]
[Discord discussion about FunctionOpInterface and FuncOp]
[Discord discussion about memory layout]
[LLVM Upstream]
[LLVM PR1]
[LLVM PR2]
llvm/llvm-project@7f04a8a
youtube link with discussion about one-shot bufferizer
pdf with info