Skip to content

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 216 commits into from
Jun 3, 2025
Merged

Update llvm version #1752

merged 216 commits into from
Jun 3, 2025

Conversation

paul0403
Copy link
Member

@paul0403 paul0403 commented May 20, 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. [mlir] Enable decoupling two kinds of greedy behavior. llvm/llvm-project#104649, [mlir] Update docs for Greedy Pattern Rewrite Driver(NFC) llvm/llvm-project#126701

  4. ops with CallOpInterface must have two new optional attributes arg_attrs and res_attrs [mlir] share argument attributes interface between calls and callables 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. [mlir][IR] Move match and rewrite functions into separate class 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. [mlir][Transforms] Add 1:N matchAndRewrite overload 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 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.

  1. The new pattern applicator (applyPatternsGreedily) is more aggressive in dead code elimination,
    and is eliminating dead Values 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.

  2. 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:

  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. [mlir][LLVM] Switch undef for poison for uninitialized values llvm/llvm-project#125629

Benefits:
Up to date with upstream versions.

[sc-92017]

@paul0403 paul0403 removed the author:build-wheels Run the wheel building workflows on this Pull Request label Jun 2, 2025
@erick-xanadu erick-xanadu added reviewer:require-wheels Pull Requests will need wheel building job successful before being merged and removed reviewer:require-wheels Pull Requests will need wheel building job successful before being merged labels Jun 3, 2025
@erick-xanadu
Copy link
Contributor

@paul0403 sorry for adding the require build wheel tag! :(

@paul0403
Copy link
Member Author

paul0403 commented Jun 3, 2025

@paul0403 sorry for adding the require build wheel tag! :(

Wait what happened? (fyi wheels succeeded at the previous commit)

@erick-xanadu
Copy link
Contributor

@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.

@paul0403 paul0403 added 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 Jun 3, 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 Jun 3, 2025
paul0403 and others added 2 commits June 3, 2025 17:11
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
@paul0403 paul0403 merged commit c79ab47 into main Jun 3, 2025
38 checks passed
@paul0403 paul0403 deleted the paul0403/update_llvm branch June 3, 2025 21:35
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.

6 participants