Skip to content

Migrate catalyst dialect to new one-shot bufferization #1708

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 37 commits into from
May 12, 2025

Conversation

paul0403
Copy link
Member

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

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]

@paul0403 paul0403 marked this pull request as ready for review May 6, 2025 20:13
Copy link

codecov bot commented May 6, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.51%. Comparing base (1749d2b) to head (a752ead).
Report is 1 commits behind head on main.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #1708   +/-   ##
=======================================
  Coverage   96.51%   96.51%           
=======================================
  Files          82       82           
  Lines        9029     9029           
  Branches      861      861           
=======================================
  Hits         8714     8714           
  Misses        258      258           
  Partials       57       57           

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

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@paul0403 paul0403 requested review from erick-xanadu and dime10 May 7, 2025 17:43
Copy link
Contributor

@erick-xanadu erick-xanadu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks very good! Just a couple of questions for clarity.

thus memory write must be true for custom call
Copy link
Contributor

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @paul0403, this is good but you are right we should be extra careful with this one given the complexities of passing memory to external functions!

Copy link
Contributor

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good 💯 Happy we went through the semantics of the custom call / callback in-depth, I think that was one of the uncertainties holding up the previous PR :)

@paul0403 paul0403 merged commit 806d6e6 into main May 12, 2025
40 checks passed
@paul0403 paul0403 deleted the paul0403/new_bufferize_catalyst_dialect branch May 12, 2025 17:09
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
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants