-
Notifications
You must be signed in to change notification settings - Fork 47
Update llvm version #1752
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Update llvm version #1752
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
instead, the mem write interface can copy for us, and even decide when a copy is needed!
…e_gradient_dialect
@paul0403 sorry for adding the require build wheel tag! :( |
Wait what happened? (fyi wheels succeeded at the previous commit) |
@paul0403 I added the tags, and they automatically fail if it is required. I removed them and tried re-running but they will keep failing until the wheels are built. Maybe we can remove them manually somehow. |
erick-xanadu
reviewed
Jun 3, 2025
…ate and can't be written to for 24 hrs)
dime10
approved these changes
Jun 3, 2025
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
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:
We update the llvm version tagged by jax 0.6.0:
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.
The new enzyme cmake target is
EnzymeStatic-21
(from 20)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.
applyPatternsAndFoldGreedily
is removed. Drop-in replacement isapplyPatternsGreedily
. [mlir] Enable decoupling two kinds of greedy behavior. llvm/llvm-project#104649, [mlir] Update docs for Greedy Pattern Rewrite Driver(NFC) llvm/llvm-project#126701ops with
CallOpInterface
must have two new optional attributesarg_attrs
andres_attrs
[mlir] share argument attributes interface between calls and callables llvm/llvm-project#123176CallInterfaceCallable
objects now must be directly casted to the calleeSymbolRefAttr
, i.e.callee.get<SymbolRefAttr>()
->cast<SymbolRefAttr>(callee)
llvm/llvm-project@35e8989The
lookupOrCreateFn
family of functions now returnFailureOr<funcop>
instead of justfuncop
, so a.value()
needs to be used to retrieve the underlyingfuncop
. llvm/llvm-project@e84f6b6The cpp api for
OneShotBufferizePassOptions
no longer needs complicated lambdas for the type converter options. They can be set with themlir::bufferization::LayoutMapOption::IdentityLayoutMap
options directly.The individual
match
andrewrite
methods in pattern rewrites are removed. Use the two-in-onematchAndRewrite
instead. [mlir][IR] Movematch
andrewrite
functions into separate class llvm/llvm-project#129861For rewrite patterns with 1-to-N convertions, a new
macthAndRewrite
overload withOneToNOpAdaptor
must be used. For us, this is only thecatalyst.list*
ops. [mlir][Transforms] Add 1:NmatchAndRewrite
overload llvm/llvm-project#116470The 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
The new mhlo depends on a 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.
Three necessary passes were removed in mhlo:
mhlo-legalize-control-flow
,mhlo-legalize-to-std
,hlo-legalize-sort
tensorflow/mlir-hlo@4a640be#diff-ef0d7e30da19a396ba036405a9ef636f8b1be194618b0a90f4602671fc2ef34dtensorflow/mlir-hlo@2a5e267#diff-f8c7cb07b43593403e00e0dbf9983f0186b4eb70368cc99af3b924061f1ea46f
mhlo-legalize-to-std
, the cmake targetMhloToStandard
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.
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.
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 Migrate to new buffer deallocation pipeline #1778 .
mlir lit test updates:
bufferization.to_tensor/memref
updated assembly formatllvm.mlir.undef
is deprecated, usellvm.mlir.poison
instead. [mlir][LLVM] Switchundef
forpoison
for uninitialized values llvm/llvm-project#125629Benefits:
Up to date with upstream versions.
[sc-92017]