Update Jax(LLVM, MLIR) and Enzyme #428

rmoyard · 2024-01-04T20:55:59Z

Description of the Change:
NullOp -> ZeroOp
llvm/llvm-project#67183

CodeGenOpt -> CodeGenOptLevel
https://github.com/llvm/llvm-project/pull/66295/files

Converters are const
llvm/llvm-project@ce25459

Remove typed pointer
llvm/llvm-project#71285

thlo gml_st on ice
tensorflow/mlir-hlo@7077cec

ByteCode is a new interface
llvm/llvm-project@837d1ce

getArgsOperandMutable was added CallOpInterface
llvm/llvm-project@d790a21

Remove redundant memref.tensor_store op
https://github.com/llvm/llvm-project/pull/71010/files

Lowering of memref.realloc was extracted to its own pass
https://reviews.llvm.org/D159430

Benefits:

Possible Drawbacks:

Related GitHub Issues:

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp

github-actions · 2024-01-05T07:32:48Z

Hello. You may have forgotten to update the changelog!
Please edit doc/changelog.md on your branch with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

mlir/lib/Quantum/Transforms/quantum_to_llvm.cpp

mlir/lib/Quantum/Transforms/emit_catalyst_pyface.cpp

rmoyard · 2024-01-05T16:39:46Z

[sc-53151]

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

mlir/lib/Quantum/Transforms/emit_catalyst_pyface.cpp

mlir/lib/Quantum/Transforms/quantum_to_llvm.cpp

mlir/test/Catalyst/BufferizationTest.mlir

frontend/test/pytest/test_jax_dynamic_api.py

mlir/lib/Gradient/Transforms/GradMethods/PS_QuantumGradient.cpp

rmoyard · 2024-01-16T20:25:21Z

Blocked by EnzymeAD/Enzyme#1608

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

codecov · 2024-02-21T20:07:53Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (96fef3b) 99.52% compared to head (c550989) 99.52%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #428   +/-   ##
=======================================
  Coverage   99.52%   99.52%           
=======================================
  Files          45       45           
  Lines        7941     7942    +1     
  Branches      537      537           
=======================================
+ Hits         7903     7904    +1     
  Misses         20       20           
  Partials       18       18

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

erick-xanadu

Looks good

mlir/lib/Gradient/Transforms/ConversionPatterns.cpp

mlir/lib/Gradient/Transforms/GradMethods/PS_QuantumGradient.cpp

Co-authored-by: erick-xanadu <110487834+erick-xanadu@users.noreply.github.com>

dime10

Thanks @rmoyard for this long awaited update 💯

.dep-versions

dime10 · 2024-02-21T22:05:37Z

frontend/catalyst/compiler.py

@@ -128,7 +128,6 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
        "func.func(hlo-legalize-to-linalg)",
        "func.func(mhlo-legalize-to-std)",
        "convert-to-signless",
-        "func.func(scalarize)",


@erick-xanadu I remember you found that either of two passes was providing similar speedups, one at the hlo level and one at the linalg level. Can we replace this pass with the linalg one (since it no longer exists)?

dime10 · 2024-02-21T22:06:30Z

frontend/catalyst/compiler.py

@@ -179,6 +178,7 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
 MLIR_TO_LLVM_PASS = (
    "MLIRToLLVMDialect",
    [
+        "expand-realloc",


Small comment why this pass is necessary will be nice, since it's hard to see why pipelines are the way they are.

I will get the reference I don't remember where it comes from

https://reviews.llvm.org/D159430 This was part of the memref lowering and it was extracted to make lowering simpler

Right I meant adding it as a comment into the file

frontend/test/lit/test_aot_compilation.py

mlir/test/Catalyst/BufferizationTest.mlir

mlir/test/Quantum/WrapperTest.mlir

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

**Description of the Change:** NullOp -> ZeroOp llvm/llvm-project#67183 CodeGenOpt -> CodeGenOptLevel https://github.com/llvm/llvm-project/pull/66295/files Converters are const llvm/llvm-project@ce25459 Remove typed pointer llvm/llvm-project#71285 thlo gml_st on ice tensorflow/mlir-hlo@7077cec ByteCode is a new interface llvm/llvm-project@837d1ce getArgsOperandMutable was added CallOpInterface llvm/llvm-project@d790a21 Remove redundant memref.tensor_store op https://github.com/llvm/llvm-project/pull/71010/files Lowering of memref.realloc was extracted to its own pass https://reviews.llvm.org/D159430 **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: erick-xanadu <110487834+erick-xanadu@users.noreply.github.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com>

rmoyard added 4 commits January 4, 2024 15:55

Update Jax and Enzyme

b2ede44

Merge branch 'main' into update_jax_llvm_mlir_enzyme

8ba55aa

MLIR/LLVM update

47b3135

Merge branch 'update_jax_llvm_mlir_enzyme' of https://github.com/Penn…

48910a5

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

rmoyard commented Jan 5, 2024

View reviewed changes

mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp Outdated Show resolved Hide resolved

rmoyard commented Jan 5, 2024

View reviewed changes

mlir/lib/Quantum/Transforms/quantum_to_llvm.cpp Outdated Show resolved Hide resolved

rmoyard commented Jan 5, 2024

View reviewed changes

mlir/lib/Quantum/Transforms/emit_catalyst_pyface.cpp Outdated Show resolved Hide resolved

rmoyard added 9 commits January 5, 2024 17:24

Update

d632ac6

Format

a8d649a

Update test and materilize op

9284a4b

Merge branch 'main' into update_jax_llvm_mlir_enzyme

02ec505

Fix PS test

98ba0fb

Update conversion

b1eab8a

Merge branch 'main' into update_jax_llvm_mlir_enzyme

fd63185

Update

803f22e

Merge branch 'update_jax_llvm_mlir_enzyme' of https://github.com/Penn…

5b1abf4

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

rmoyard commented Jan 10, 2024

View reviewed changes

mlir/lib/Quantum/Transforms/emit_catalyst_pyface.cpp Show resolved Hide resolved

mlir/lib/Quantum/Transforms/quantum_to_llvm.cpp Outdated Show resolved Hide resolved

Change test

455de63

dime10 reviewed Jan 10, 2024

View reviewed changes

mlir/test/Catalyst/BufferizationTest.mlir Outdated Show resolved Hide resolved

rmoyard added 6 commits January 10, 2024 14:10

Quantum conversion tests

aad042f

Update module ctx Jax

48e51d8

Format

9bb368f

Fix lit test

8c4419a

Isort

bc22c76

Jax dynamic

7df7ee1

rmoyard commented Jan 10, 2024

View reviewed changes

frontend/test/pytest/test_jax_dynamic_api.py Outdated Show resolved Hide resolved

rmoyard commented Jan 11, 2024

View reviewed changes

mlir/lib/Gradient/Transforms/GradMethods/PS_QuantumGradient.cpp Outdated Show resolved Hide resolved

Update enzyme

e92f0c7

Reset

4b08304

rmoyard added 11 commits January 16, 2024 15:25

Merge branch 'main' into update_jax_llvm_mlir_enzyme

b355257

Merge branch 'main' into update_jax_llvm_mlir_enzyme

0296f15

Update enzyme

b611332

Fix Modifiers (control and adjoint)

8d8b3d3

Merge branch 'update_jax_llvm_mlir_enzyme' of https://github.com/Penn…

627d964

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

Merge branch 'main' into update_jax_llvm_mlir_enzyme

cceb600

Misc

a30d8d3

Merge branch 'update_jax_llvm_mlir_enzyme' of https://github.com/Penn…

399995b

…yLaneAI/catalyst into update_jax_llvm_mlir_enzyme

Remove test

29c585e

Formatting and remove test

c8c8e6e

C++ format

6e9df06

rmoyard marked this pull request as ready for review February 21, 2024 18:26

rmoyard requested review from dime10 and erick-xanadu February 21, 2024 18:27

dime10 added this to the v0.5.0 milestone Feb 21, 2024

Update lit test

2fa704c

rmoyard added the author:build-wheels Run the wheel building workflows on this Pull Request label Feb 21, 2024

dime10 added the reviewer:require-wheels Pull Requests will need wheel building job successful before being merged label Feb 21, 2024

erick-xanadu approved these changes Feb 21, 2024

View reviewed changes

mlir/lib/Gradient/Transforms/ConversionPatterns.cpp Show resolved Hide resolved

mlir/lib/Gradient/Transforms/GradMethods/PS_QuantumGradient.cpp Outdated Show resolved Hide resolved

rmoyard and others added 2 commits February 21, 2024 17:09

Trigger CI

c377524

Update mlir/lib/Gradient/Transforms/GradMethods/PS_QuantumGradient.cpp

450ab43

Co-authored-by: erick-xanadu <110487834+erick-xanadu@users.noreply.github.com>

dime10 approved these changes Feb 21, 2024

View reviewed changes

rmoyard and others added 2 commits February 21, 2024 17:25

Update .dep-versions

d7895c4

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

Merge branch 'main' into update_jax_llvm_mlir_enzyme

c550989

rmoyard merged commit 8649c9e into main Feb 22, 2024
55 of 56 checks passed

rmoyard deleted the update_jax_llvm_mlir_enzyme branch February 22, 2024 01:40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update Jax(LLVM, MLIR) and Enzyme #428

Update Jax(LLVM, MLIR) and Enzyme #428

rmoyard commented Jan 4, 2024 •

edited

Loading

github-actions bot commented Jan 5, 2024

rmoyard commented Jan 5, 2024

rmoyard commented Jan 16, 2024

codecov bot commented Feb 21, 2024 •

edited

Loading

erick-xanadu left a comment

dime10 left a comment

dime10 Feb 21, 2024

dime10 Feb 21, 2024

rmoyard Feb 21, 2024

rmoyard Feb 21, 2024

dime10 Feb 22, 2024

Update Jax(LLVM, MLIR) and Enzyme #428

Update Jax(LLVM, MLIR) and Enzyme #428

Conversation

rmoyard commented Jan 4, 2024 • edited Loading

github-actions bot commented Jan 5, 2024

rmoyard commented Jan 5, 2024

rmoyard commented Jan 16, 2024

codecov bot commented Feb 21, 2024 • edited Loading

Codecov Report

erick-xanadu left a comment

Choose a reason for hiding this comment

dime10 left a comment

Choose a reason for hiding this comment

dime10 Feb 21, 2024

Choose a reason for hiding this comment

dime10 Feb 21, 2024

Choose a reason for hiding this comment

rmoyard Feb 21, 2024

Choose a reason for hiding this comment

rmoyard Feb 21, 2024

Choose a reason for hiding this comment

dime10 Feb 22, 2024

Choose a reason for hiding this comment

rmoyard commented Jan 4, 2024 •

edited

Loading

codecov bot commented Feb 21, 2024 •

edited

Loading