Closed paul0403 closed 3 weeks ago
Attention: Patch coverage is 92.40506%
with 6 lines
in your changes missing coverage. Please review.
Project coverage is 97.89%. Comparing base (
7231b5b
) to head (6b0ad2e
).
Files | Patch % | Lines |
---|---|---|
frontend/catalyst/jax_primitives.py | 90.32% | 3 Missing and 3 partials :warning: |
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
TODO: changelog, top-level API documentation of catalyst.cancel_inverses
Note on codecov failures: this is because lit test are not covered by codecov.
Currently these tests exist:
Pytest: tests for the numerical results remaining the same before and after @cancel_inverses. Also added an additional comparison with unjit-ted default.qubit
to make sure.
The lit test checks for correct lowering to jaxpr and mlir, including:
transform.with_named_sequence
is generated for all qjits-transform-interpreter
to check for correctness during lowering, with the cancel_inverses-ed qnode both inside and outside a qjit-transform-interpreter
to check for correctness of pass actionI believe these tests provide enough usage coverage
Since the jax update #931 , which updated LLVM, the transform dialect enabled a new check that the payload module cannot be an ancestor of the transformer module.
A new commit conforms to this new requirement.
New design:
transform.named_sequence @__transform_main
in a transformer module, and put the transformer module in the top-level module (aka the payload).
The transformer module needs to be marked with attribute {transform.with_named_sequence}
The payload module does not need to be marked with any attributes.-transform-interpreter
and -transform-cleanup
passes.
We add a new pass, -apply-transform-sequence
. This pass does three things (in the following order):
Looks great! I would be ok with this being merged (with the exception of the walkers not exiting early and skipping nested operations if possible).
(bookkeeping comment) The walkers are now exited early when the desired result is seen.
Context: Implement a frontend UI to invoke the mlir quantum peephole transformation passes. Currently there is a
remove-chained-self-inverse
mlir pass in Catalyst that removes two neighbouring Hadamard gates in a circuit, but there is no way to actually invoke it from the frontend. This PR implements this frontend as a decorator on a qnode. When thecatalyst.cancel_inverses
decorator is added onto a qnode, theremove-chained-self-inverse
mlir pass will be run on that qnode.Following the decision to split #883 into two PRs, this is the frontend portion.
Description of the Change: We implment the peephole optimization library with the mlir transform dialect https://mlir.llvm.org/docs/Dialects/Transform/, https://mlir.llvm.org/docs/Tutorials/transform/
Briefly, we wish to generate the following mlir during frontend tracing and lowering
The outer module is the payload (the module to run the transform on), and the inner module is the transformer (the module that schedules what passes to run). The schedule in this example is run the
-remove-chained-self-inverse
pass on the payload module@workflow
, with pass optionfunc-name=f
.The
transform.named_sequence
must be terminated by atransform.yield
.A new pass,
ApplyTransformSequence.cpp
, will perform the following:We generate this mlir through two new jax primitives.
The first one,
transform_named_sequence_p
, will be lowered totransform.named_sequence
with atransform.yield
inside and put it in a parent transformer module marked with the unit attributetransform.with_named_sequence
. The parent transformer module is inserted into the top-level payload module.The second one,
apply_registered_pass_p
, will be lowered to atransform.apply_registered_pass
inside thetransform.named_sequence
. If onetransform.apply_registered_pass
already exists in the sequence, the new pass will be added to after the previous one.The design is as follows:
QJIT start
Capture jaxpr (
QJIT.capture()
) 1.1. Before capture, we insert thetransform_named_sequence_p
to every jaxpr produced (the purpose is so that transform dialect will not fail) 1.2. During capture, if the API@cancel_inverses
is called, inject a correspondingapply_registered_pass_p
into jaxprGenerate mlir and compile (
QJIT.generate_ir()
,QJIT.compile()
) 2.1. The jaxpr primitives are lowered to mlir. 2.2. Here theCompiler
class will run the pipeline, with the new-apply-transform-sequence
passBenefits: The user can now run quantum mlir passes from the frontend.
Possible Drawbacks:
Related GitHub Issues:
[sc-67519]