Migrate to new buffer deallocation pipeline #1778
Draft
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
…oved. Intended replacement is --buffer-deallocation-pipeline. [mlir][bufferization] Remove buffer-deallocation pass llvm/llvm-project#126366, https://discourse.llvm.org/t/psa-bufferization-new-buffer-deallocation-pipeline/73375. However, we encountered a few issues. See detailed discord discussion at https://discord.com/channels/636084430946959380/642426447167881246/1376919538301403276 The heart of the new deallocation pass, --ownership-based-buffer-deallocation, blanket-ly fails for ops with unknown memory effects: https://github.com/llvm/llvm-project/blob/da4958ae2b384c2a027cf20c67b7e211d39fcbfe/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp#L523 To solve this issue, the suggestion was to add memory effects to custom operations. The migration guide also suggests that bufferizable ops no longer implement the bufferizesToAllocation method, so we remove them. This was supposed to be done alongside the llvm update in #1752; However, soon it became clear that this migration to the new buffer deallocation is very complicated, and should be its own story. The llvm update in #1752 thus did not finish this migration. This PR records the work that was already done on this back in #1752.
|
I am opening this PR simply to record the WIP progress of what we have so far for the new dealloc. This can be beneficial when we revisit. The huge diff is just from #1752 itself, and will go away after it's merged. Please leave this PR in draft to save CI. Thanks! |
Merged
paul0403
added a commit
that referenced
this pull request
Jun 3, 2025
**Context:** We update the llvm version tagged by jax 0.6.0: ``` mhlo=617a9361d186199480c080c9e8c474a5e30c22d1 llvm=179d30f8c3fddd3c85056fd2b8e877a4a8513158 ``` We also update Enzyme to the latest version, which is 0.0.180, at commit `db0181320d6e425ee963bd496ed0d8dbb615be18` **Description of the Change:** Firstly, jax recently moved from the Google github organization to its own jax-ml organization. This means the urls, and the retrieval method for the underlying llvm and mhlo git commit tags, needs to be updated. (Thanks @mehrdad2m !) Now on to the actual changes. I will list the changes in increasing complexity. 1. The new enzyme cmake target is `EnzymeStatic-21` (from 20) 2. Enzyme works with a later llvm then our target, so it has some llvm intrinsics unknown to the one we are targeting. We patch them away. They do not concern us since they are all intrinsics for nvidia backends. 3. `applyPatternsAndFoldGreedily` is removed. Drop-in replacement is `applyPatternsGreedily`. llvm/llvm-project#104649, llvm/llvm-project#126701 4. ops with `CallOpInterface` must have two new optional attributes `arg_attrs` and `res_attrs` llvm/llvm-project#123176 5. `CallInterfaceCallable` objects now must be directly casted to the callee `SymbolRefAttr`, i.e. `callee.get<SymbolRefAttr>()` -> `cast<SymbolRefAttr>(callee)` llvm/llvm-project@35e8989 6. The `lookupOrCreateFn` family of functions now return `FailureOr<funcop>` instead of just `funcop`, so a `.value()` needs to be used to retrieve the underlying `funcop`. llvm/llvm-project@e84f6b6 7. The cpp api for `OneShotBufferizePassOptions` no longer needs complicated lambdas for the type converter options. They can be set with the `mlir::bufferization::LayoutMapOption::IdentityLayoutMap` options directly. 8. The individual `match` and `rewrite` methods in pattern rewrites are removed. Use the two-in-one `matchAndRewrite` instead. llvm/llvm-project#129861 9. For rewrite patterns with 1-to-N convertions, a new `macthAndRewrite` overload with `OneToNOpAdaptor` must be used. For us, this is only the `catalyst.list*` ops. llvm/llvm-project#116470 10. The lowering of `cf::AssertOp` to llvm was split from the overall`--covert-cf-to-llvm` pass. We need to manually call this separate pattern for cf.assert duriing quantum to llvm dialect lowering, where we also convert cf to llvm. https://github.com/llvm/llvm-project/pull/120431/files 11. The new mhlo depends on a [shardy](https://github.com/openxla/shardy) dialect. Shardy is built with bazel, not cmake. Building shardy ourselves would be very difficult (not having bazel in our build ecosystem is a hard constraint, cc @mlxd ), and also not necessary (we just use mhlo for their "standard" passes). We thus patch out all shardy components. 12. Three necessary passes were removed in mhlo: `mhlo-legalize-control-flow`, `mhlo-legalize-to-std`, `hlo-legalize-sort` tensorflow/mlir-hlo@4a640be#diff-ef0d7e30da19a396ba036405a9ef636f8b1be194618b0a90f4602671fc2ef34d tensorflow/mlir-hlo@2a5e267#diff-f8c7cb07b43593403e00e0dbf9983f0186b4eb70368cc99af3b924061f1ea46f - Alongside the removal of `mhlo-legalize-to-std`, the cmake target `MhloToStandard` was removed too. We simply patch them back for now. **For the above two points, note that there will be an overall migration to the stablehlo repo, as mhlo is sunseting. Therefore, spending too much time on this isn't necessary, so we just patch.** 13. The new pattern applicator (`applyPatternsGreedily`) is more aggressive in dead code elimination, and is eliminating dead `Value`s in the adjoint gradient method. The `nodealloc` function we generate for adjoint gradient lowering used to only return the qreg, not the expval result. This causes the expval op to be eliminated since it has no users. This further causes wrong gradient results, since the entire program, all ops included (regardless of dead or not), impacts the gradient through chain rule. To avoid this, we return the expval result as well. In doing this, we implicitly assume that differentiated qnodes can only return expval. Although this assumption is true and also restricted by frontend, ideally we should not have it hard coded. We leave this as a TODO for a future feature. 14. The old `--buffer-deallocation` pass is removed. Intended replacement is `--buffer-deallocation-pipeline`. This migration is very complicated. We simply add back the old buffer deallocation pass in the catalyst dialect as a util for now. We will revisit this in #1778 . mlir lit test updates: 1. `bufferization.to_tensor/memref` updated assembly format 2. gradient adjoint lowering test returns both qreg and expval 3. Some inverse unrealized conversion cast pairs are canceled by the new pattern rewriter. 4. `llvm.mlir.undef` is deprecated, use `llvm.mlir.poison` instead. llvm/llvm-project#125629 **Benefits:** Up to date with upstream versions. [sc-92017] --------- Co-authored-by: Tzung-Han Juang <tzunghan.juang@gmail.com> Co-authored-by: Ritu Thombre <42207923+ritu-thombre99@users.noreply.github.com> Co-authored-by: Mehrdad Malekmohammadi <mehrdad.malek@xanadu.ai> Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: Joey Carter <joseph.carter@xanadu.ai>
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.
As part of the llvm update, the old --buffer-deallocation pass is removed.
Intended replacement is --buffer-deallocation-pipeline. llvm/llvm-project#126366, https://discourse.llvm.org/t/psa-bufferization-new-buffer-deallocation-pipeline/73375. However, we encountered a few issues. See detailed discord discussion at https://discord.com/channels/636084430946959380/642426447167881246/1376919538301403276
The heart of the new deallocation pass, --ownership-based-buffer-deallocation, blanket-ly fails for ops with unknown memory effects: https://github.com/llvm/llvm-project/blob/da4958ae2b384c2a027cf20c67b7e211d39fcbfe/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp#L523 To solve this issue, the suggestion was to add memory effects to custom operations. The migration guide also suggests that bufferizable ops no longer implement the bufferizesToAllocation method, so we remove them.
This was supposed to be done alongside the llvm update in #1752; However, soon it became clear that this migration to the new buffer deallocation is very complicated, and should be its own story.
The llvm update in #1752 thus did not finish this migration. This PR records the work that was already done on this back in #1752.
List of memory effects added to the ops back in #1752 :
quantum dialect:
quantum.init: MemAlloc
quantum.finalize: MemFree
quantum.device: MemAlloc
quantum.device_release: MemFree
(!) quantum.alloc already has a NoMemoryEffect. I left it alone.
quantum.dealloc: MemFree
quantum.set_state and set_basis_state: MemWrite
quantum.custom, multi_rz and unitary: MemRead
Originally I had both read and write here, but this caused passes like cancel inverses to fail to remove some gates without uses after pattern rewrite.
quantum.gphase: MemRead, MemWrite
quantum.adjoint: MemRead, MemWrite
quantum.compbasis and namedobs: MemRead
quantum.measure: MemRead, MemWrite
quantum.sample, counts, state, probs: MemAlloc, MemRead
These need to allocate a classical array to store the results
quantum.expval, var: MemRead
mbqc dialect:
mbqc.measure_in_basis: MemRead, MemWrite
catalyst dialect:
catalyst.list_init: MemAlloc
catalyst.list_dealloc: MemFree
catalyst.list_push: MemWrite, MemAlloc
catalyst.list_pop: MemRead, MemFree
catalyst.list_load_data: MemRead
catalyst.print: MemRead, MemAlloc (alloc is needed to print strings)
catalyst.assert: MemWrite (write assertion error message)
gradient dialect:
I am not super sure regarding them, so I just turn on all effects to be safe.
gradient.adjiont: MemRead, MemWrite, MemAlloc, MemFree
gradient.backprop: MemRead, MemWrite, MemAlloc, MemFree
gradient.jvp: MemRead, MemWrite, MemAlloc, MemFree
gradient.vjp: MemRead, MemWrite, MemAlloc, MemFree
Note that these are not tested/thoroughly considered.
See conversations in #1752 for some of the initial discussions on these memory effects.