Skip to content

[qref 3] Migrate capture frontend to produce reference semantics MLIR#2663

Draft
paul0403 wants to merge 47 commits intomainfrom
paul0403/qref_frontend_main
Draft

[qref 3] Migrate capture frontend to produce reference semantics MLIR#2663
paul0403 wants to merge 47 commits intomainfrom
paul0403/qref_frontend_main

Conversation

@paul0403
Copy link
Copy Markdown
Member

@paul0403 paul0403 commented Apr 8, 2026

Context:
Migrate capture frontend to produce reference semantics MLIR.

This is the main PR. Each item will come into this PR as its own PR.

Description of the Change:

  • Expose python bindings for operations in the qref MLIR dialect, and define their corresponding jax primitives and lowering rules on these primitives

  • The conversion from plxpr to catalyst jaxpr now produces qref primitives whenever possible. This includes allocation, deallocation, extract (now qref.get; note that reference semantics does not have an insert op), observables, gates and adjoints

  • For the following primitives, the conversion from the PL version to catalyst version is no longer done. Instead, we define lowering rule on the PL primitives directly: for_loop, while_loop, if, adjoint. This is because with reference semantics, their MLIR operation now take in and return purely classical values, so the special handling regarding the quantum values in their operands is not needed anymore.

  • At the beginning of the default pipeline, the following passes are added (in this order):

    • --canonicalize
    • --verify-no-quantum-use-after-free
    • --convert-to-value-semantics
    • --canonicalize
  • Tests are edited correspondingly, mostly to update their expected results to check for reference semantics

  1. In pytest/from_plxpr/test_from_plxpr.py:
    • Tests comparing catxpr generated from capture and legacy frontends are removed. They are converted to proper lit tests, checking for the qref MLIR generated from the capture path. As a result, the compare_call_jaxpr test util has been removed. Note that the checks for execution results on these tests are kept.
    • Tests checking for value semantics catalyst primitives have been updated to check for the corresponding qref primitives.

Benefits:

  1. Significantly, and I mean truly significantly, simplifies the capture frontend. Many plxpr primitives are used directly, which is a big step towards lowering plxpr and removing catxpr from the pipeline altogether.
  2. Many tech debts are automatically fixed. Here's an (incomplete) list:
  1. Greatly improve data flow around control flow/adjoint/subroutine regions. We no longer have to insert all qubits into their registers and let these regions take in the registers. Frontend just generates mlir where the regions see qref values from above via closure, and the --convert-to-value-semantics pass will only create value semantics regions that take in individual qubit values, instead of entire register values.

Related GitHub Issues:
closes #2526 [sc-112704]
closes #2527 [sc-112706]

@paul0403 paul0403 changed the title Migrate capture frontend to produce reference semantics MLIR [qref 3] Migrate capture frontend to produce reference semantics MLIR Apr 8, 2026
@paul0403 paul0403 added the wip PRs that are a Work-In-Progress label Apr 8, 2026
Comment thread doc/releases/changelog-dev.md Outdated
Comment thread .github/workflows/check-catalyst.yaml Fixed
Comment thread .github/workflows/check-catalyst.yaml Outdated
Comment on lines +533 to +538
- name: Run qref frontend lit tests
run: |
for file in ./frontend/test/lit/test_qref/*; do\
echo "Testing $$file"\;
python3 "$$file" | ./llvm-build/bin/FileCheck "$$file";\
done
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I added a new CI step (under "frontend tests (gcc)") that only runs the qref lit tests I added in these PRs.

This is because the regular tests will fail until the entire migration is complete. In the meantime, this CI step can be used to verify the qref implementation itself.

This step will be removed once the entire migration is complete: it is just a subset of make lit, running frontend filecheck lit tests only on my newly added qref test files.

paul0403 and others added 7 commits April 16, 2026 10:03
…reference semantics (#2664)

**Context:**
Migrate alloc, dealloc and observable plxpr conversion to reference
semantics.

Note: this PR only handles the initial allocation from the device
initialization. It does not deal with dynamically allocated qubits from
`qml.allocate()` within a circuit. These will be handled after the gates
are handled, so there's a reasonable way to test them.

- [x] Alloc
- [x] Dealloc
- [x] AllocQubit (not generated by frontend)
- [x] DeallocQubit (not generated by frontend)
- [x] Compbasis
- [x] NamedObs
- [x] Hermitian

**Related GitHub Issues:**
[sc-115868]
Comment thread frontend/catalyst/__init__.py Outdated
Comment thread frontend/catalyst/pipelines.py Outdated
paul0403 and others added 3 commits April 22, 2026 09:49
…tics (#2672)

**Context:**
Migrate gate-like ops' plxpr conversion to reference semantics.

This PR also handles dynamic allocations and deallocations from
`qml.allocate()`.

Note: adjointed gates are not handled in this PR. They will be handled
in a future PR that specifically handles the adjoint region from plxpr.

- [x] CustomOp
- [x] MultiRZOp
- [x] PCPhaseOp
- [x] PauliRotOp
- [x] GlobalPhaseOp
- [x] QubitUnitaryOp
- [x] MeasureOp
- [x] SetStateOp
- [x] SetBasisStateOp

**Related GitHub Issues:**
[sc-115869]
@paul0403 paul0403 added author:build-wheels Run the wheel building workflows on this Pull Request reviewer:require-wheels Pull Requests will need wheel building job successful before being merged labels Apr 23, 2026
paul0403 and others added 15 commits April 23, 2026 10:46
…#2694)

**Context:**
Migrate for loop's plxpr conversion to reference semantics

**Description of the Change:**
All quantum reference values are captured now via closure from above
automatically by jax infra.
No need to keep track of quantum value args to the for loop anymore
:rocket: :fire: :partying_face:

**Benefits:**
Much cleaner conversion of plxpr to catalys jaxpr. 

[sc-115870]
…cs (#2717)

**Context:**
Migrate while loop's plxpr conversion to reference semantics

**Description of the Change:**
Delete the conversion from plxpr while loop primitive to catalyst while
loop primitive, so we just use the lowering of plxpr while loop
primitive directly.

**Benefits:**
Much cleaner conversion of plxpr to catalys jaxpr.

[sc-115871]
Comment on lines +237 to +238
pass
# precompile_decomp_rules() # pragma: no cover
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

TODO: reenable precompiled decomp rules when the entire migration is compelte

Comment thread .github/workflows/check-catalyst.yaml
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

author:build-wheels Run the wheel building workflows on this Pull Request reviewer:require-wheels Pull Requests will need wheel building job successful before being merged wip PRs that are a Work-In-Progress

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Capture path does not insert qubits before observables with dynamic indices Capture path does not insert qubits before gates with dynamic indices

3 participants