A new catalyst decomposition at MLIR based on the experimental graph-based decomposition system #2029

maliasadi · 2025-09-04T16:46:30Z

Catalyst uses the new graph-based decomposition system, resulting in the same final gate set as PennyLane.
A new MLIR decompose-lowering will rewrite the circuit and apply decomposition rules in the core Catalyst compiler.
Circuits containing structure (for and while loops) can be decomposed with Catalyst present.
Decomposition rules with classical StableHlo ops can be compiled [sc-100673]
Decomposition rules can include dynamically allocated wires and still work with Catalyst [WIP]
Decomposition rules can be manually added to operations via add_decomps and fixed/alt_decomps.

qml.capture.enable()
qml.decomposition.enable_graph()

@qml.qjit
@partial(qml.transforms.decompose, gate_set={qml.H, qml.CNOT, qml.ControlledPhaseShift})
@qml.qnode(qml.device("null.qubit"))
def circuit():
    [qml.Hadamard(i) for i in range(4)]
    qml.QFT(wires=range(4))
    return qml.state()

…r e2e testing from v3)

codecov · 2025-09-04T17:02:14Z

Codecov Report

❌ Patch coverage is 99.20635% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 97.46%. Comparing base (0a78ffe) to head (f1d0fb5).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
frontend/catalyst/from_plxpr/decompose.py	98.68%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2029      +/-   ##
==========================================
+ Coverage   97.45%   97.46%   +0.01%     
==========================================
  Files          90       91       +1     
  Lines       10476    10591     +115     
  Branches      972      990      +18     
==========================================
+ Hits        10209    10323     +114     
+ Misses        213      212       -1     
- Partials       54       56       +2

☔ 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 · 2025-09-09T21:15:07Z

Just a question: what is this PR's relationship to #1962 ?

maliasadi · 2025-09-10T19:23:25Z

Just a question: what is this PR's relationship to #1962 ?

@paul0403 The main difference is the scope! This PR enables decomposition at the MLIR level and improves interoperability with other MLIR/xDSL passes. #1962 is a temporary patch to apply decomposition at PLxPR that is still aligned with the compiler-specific decomposition in this PR. But for the compiler-specific decomposition, we don't use the graph-based decomposition to find the solutions because

the experimental graph-based system hasn't been maintained w/ capture enabled and from_plxpr
the compiler-specific decomp mainly deals with high-level operators and templates that each have a single decomposition rule, so using the graph wouldn’t provide much benefit in most cases

rniczh

Awesome 🎉 just small changes and questions

frontend/catalyst/from_plxpr/decompose.py

frontend/catalyst/from_plxpr/from_plxpr.py

frontend/catalyst/from_plxpr/decompose.py

frontend/catalyst/from_plxpr/from_plxpr.py

frontend/catalyst/passes/builtin_passes.py

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

mlir/include/Quantum/Transforms/Passes.td

mlir/lib/Quantum/Transforms/decompose_lowering.cpp

runtime/include/RuntimeCAPI.h

paul0403

Looks good from my side, seems like all the "hacky" bits are isolated in the new decompose.py and there are plans to use the upstream PL decomposition interpreter, so the main catalyst from_plxpr program flow is not seeing anything increasing hacky-ness 😆

Approving pending Astral's comments regarding decomp logic (which I'm not super familiar with).

frontend/catalyst/from_plxpr/decompose.py

frontend/catalyst/from_plxpr/from_plxpr.py

frontend/catalyst/from_plxpr/decompose.py

… present in MLIR (#2001) **Context:** Introduces a new MLIR pass that enables user-defined decomposition of quantum gate operations. The pass discovers decomposition functions in the module and uses MLIR's inlining infrastructure to replace target quantum operations with their decomposed implementations. **Description of the Change:** This PR will **recognize** the `func.func` as decomposition if they have name start with the pattern `<Gate>_rule[.]*` where the `<Gate>` is a `quantum.custom` op it expect to replace (It's not what generated from frontend). Also you can just mark the attributes `catalyst.decomposition` and `catalyst.decomposition.target_op` at decomposition function, it still works. And this pass just **discover** those decomposition functions and **replace** the target operation with `call` to the function and rely on the Inliner interface to inline these decomposition function if needed. The stages of the `--decompose-lowering`: 1. `decompose_lowering.cpp`: Main pass that orchestrates the decomposition process 2. `DecomposeLoweringPatterns`: Pattern rewriting logic that replaces `quantum.custom` operations with function call 3. `QuantumInlinerInterface`: Dialect interface that enables MLIR's inliner to process quantum operations correctly 4. Remove unused decomposition function (if the inliner decide not to inline the certain function, then it will leave a `func.call`, so the func is used, shouldn't be removed) The type signature between decomposition function provided from the frontend is different to the `quantum.custom` op, `OpSignatureAnalyzer` does the trick to get the enough information to generate a function call to decomposition function. Current right now, frontend choose to generate the decomposition function with type: ``` (qreg, param*, inWires*, inCtrlWires*?, inCtrlValues*?) -> qreg ``` We need to figure out the information that need to pass to create the function call: qreg, `in wires indices`, and `in ctrl wires`. `OpSignatureAnalyzer` will use the target `quantum.custom` op (that one we need to replace) to traverse the qubit, and find out the qreg, and wire indices. The traversal logic based on the following assumption for the quantum dialect: 1. The ordering of qubits is preserved in quantum instructions 2. For other instructions, they will use a register, so will hit an extract before them It promise that every gates, the result (qubits) will match to the operands (qubits) at the same index. That's the important assumptions to support the traversal logic here. Traverse the qubit, match the index, keep going up until reach the `qauntum.extract` operation. After then, we use those information to generate a function call and rely on the Inliner in MLIR infra to replace it. **Benefits:** The decision to use MLIR's inlining infrastructure (`InlinerPass`) is to provide future-proofing and flexibility. **Possible Drawbacks:** **Related GitHub Issues:** [sc-98593] --------- Co-authored-by: Ali Asadi <[email protected]>

…talyst into feature/handle_gdecomp_v4

astralcai

Thank you for answering all my questions.

…andle_gdecomp_v4

maliasadi added 2 commits September 4, 2025 12:29

Graph-based Catalyst decomposition at MLIR (cherry-picking commits fo…

81042ff

…r e2e testing from v3)

Merge with main

579c3f8

Tidy up

02be034

maliasadi force-pushed the feature/handle_gdecomp_v4 branch from f22bcd7 to 02be034 Compare September 5, 2025 18:01

maliasadi and others added 13 commits September 5, 2025 14:04

Merge branch 'main' into feature/handle_gdecomp_v4

97a1b17

Add py example

c20baaf

cherry-picking decomp_gateset commits

b6f4155

Tidy up

821069d

Merge branch 'main' into feature/handle_gdecomp_v4

bcc0f0e

Make the visibility of decomp rules to public

891dc5a

Tidy up

2617208

Fix a couple of more issues

599baf0

Fix func-redefined

5465b93

Update comiled name of rules

d4dcb19

Update

e6bf756

apply multiple decomp pass

fbf5911

pylint: disable=too-many-instance-attributes

6253c90

maliasadi added 4 commits September 9, 2025 18:23

Add decompose-lowering to the pass pipeline

bd5fc7c

provide support for decomp to apply after/before other passes

a22a4c5

code format

6539064

Merge with main

d668929

maliasadi added 2 commits September 11, 2025 09:44

Merge branch 'main' into feature/handle_gdecomp_v4

b4737a7

Tidy up

5ce3c00

rniczh reviewed Sep 11, 2025

View reviewed changes

Tidy up

c2f5ea2

astralcai self-requested a review September 15, 2025 20:48

Update tests

b2237c0

astralcai reviewed Sep 18, 2025

View reviewed changes

frontend/catalyst/from_plxpr/decompose.py Outdated Show resolved Hide resolved

astralcai reviewed Sep 18, 2025

View reviewed changes

frontend/catalyst/from_plxpr/decompose.py Outdated Show resolved Hide resolved

astralcai reviewed Sep 18, 2025

View reviewed changes

frontend/catalyst/from_plxpr/from_plxpr.py Outdated Show resolved Hide resolved

maliasadi added 5 commits September 18, 2025 16:33

Update orders of catalyst decomps

54e539d

Update tests

4f8c6ee

Merge with main

06beb8a

Support multi-qubit decomp rules

fdaef6b

Update changelog

9c2681c

paul0403 reviewed Sep 19, 2025

View reviewed changes

paul0403 approved these changes Sep 19, 2025

View reviewed changes

astralcai reviewed Sep 22, 2025

View reviewed changes

frontend/catalyst/from_plxpr/decompose.py Outdated Show resolved Hide resolved

rniczh and others added 9 commits September 23, 2025 14:28

Apply code review suggestions

2651f9e

Merge branch 'feature/handle_gdecomp_v4' of github.com:PennyLaneAI/ca…

2591daf

…talyst into feature/handle_gdecomp_v4

Update support with dynamic allocations

4273792

Update tests

e879db8

Update

12279d0

Fix the issue with dynamic qubit allocs

e7d071f

Update lit tests

b28f1d6

Tidy up

b014b12

astralcai approved these changes Sep 25, 2025

View reviewed changes

maliasadi added 2 commits September 25, 2025 12:26

Update tests

974eea3

Merge branch 'main' of github.com:PennyLaneAI/catalyst into feature/h…

afa2848

…andle_gdecomp_v4

This was referenced Sep 25, 2025

🚧 A compiler native Decomposition Graph Solver in C++ #2065

Draft

Add rulesPath option to DecomposeLoweringPass #2062

Draft

maliasadi added 2 commits September 26, 2025 09:33

StableHLO Update

7bffa1f

Merge with main

f1d0fb5

maliasadi merged commit a13db03 into main Sep 26, 2025
38 checks passed

maliasadi deleted the feature/handle_gdecomp_v4 branch September 26, 2025 14:36

Uh oh!

A new catalyst decomposition at MLIR based on the experimental graph-based decomposition system #2029

A new catalyst decomposition at MLIR based on the experimental graph-based decomposition system #2029

Uh oh!

Conversation

maliasadi commented Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

paul0403 commented Sep 9, 2025

Uh oh!

maliasadi commented Sep 10, 2025

Uh oh!

rniczh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

paul0403 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

astralcai left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

maliasadi commented Sep 4, 2025 •

edited

Loading

codecov bot commented Sep 4, 2025 •

edited

Loading

paul0403 left a comment •

edited

Loading