Use pass_name from pennylane transform (#2149)

**Context:**

Trying to improve our integration between pennylane transforms and
catalyst passes.

See #2206 for where we are able to update all built-in catalyst passes
to be transforms instead. That PR helps validate we can make this change
with no breaking behaviour.

Going forward, we can setup a plan to get rid of the catalyst-specific
pass handling infrastructure.

Also note that this PR targets the PennyLane branch, so we can confirm
we won't need anymore changes there to get catalyst working correctly
before approving and merging.

**Description of the Change:**

Uses the `transform.pass_name` property as the higher priority source of
the `pass_name`.

Updates the `pass_pipeline` from the `quantum_kernel_p` primitive to be
a tuple of `TransformContainer` objects instead of `Pass`, but leaves in
the logic to handle `tuple[Pass,...]` for safer backwards compatibility.

Also makes using transforms with pass names by splitting a `QNode`'s
transform program into "things at the start of the program that are tape
transforms" and "things at the end of the program that are MLIR passes".
This will allow us to get rid of `PassPipelineWrapper` use for both
systems and soften the change from old frontend to new frontend.

**Benefits:**

Unified handling of transforms across pennylane and catalyst. 

**Possible Drawbacks:**
 


**Related GitHub Issues:**

Depends on PennyLaneAI/pennylane#8539
[sc-103775]

---------

Co-authored-by: Paul <[email protected]>

Author: albi3ro

.dep-versions

@@ -10,7 +10,7 @@ enzyme=v0.0.203
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt'
-pennylane=0.44.0-dev42
+pennylane=0.44.0-dev44
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -70,6 +70,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* Catalyst can now use the new `pass_name` property of pennylane transform objects. Passes can now
+  be created using `qml.transform(pass_name=pass_name)` instead of `PassPipelineWrapper`.
+  [(#2149](https://github.com/PennyLaneAI/catalyst/pull/2149)
+
 * An error is now raised if a transform is applied inside a QNode when program capture is enabled.
   [(#2256)](https://github.com/PennyLaneAI/catalyst/pull/2256)
 

doc/requirements.txt

@@ -33,4 +33,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.44.0-dev16
 pennylane-lightning==0.44.0-dev16
-pennylane==0.44.0-dev42
+pennylane==0.44.0-dev44

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -29,11 +29,9 @@
 from pennylane.capture.expand_transforms import ExpandTransformsInterpreter
 from pennylane.capture.primitives import jacobian_prim as pl_jac_prim
 from pennylane.capture.primitives import transform_prim
-from pennylane.transforms import cancel_inverses as pl_cancel_inverses
 from pennylane.transforms import commute_controlled as pl_commute_controlled
 from pennylane.transforms import decompose as pl_decompose
 from pennylane.transforms import merge_amplitude_embedding as pl_merge_amplitude_embedding
-from pennylane.transforms import merge_rotations as pl_merge_rotations
 from pennylane.transforms import single_qubit_fusion as pl_single_qubit_fusion
 from pennylane.transforms import unitary_to_rot as pl_unitary_to_rot
 
@@ -48,7 +46,6 @@
     qdealloc_p,
     quantum_kernel_p,
 )
-from catalyst.passes.pass_api import Pass
 from catalyst.utils.patching import Patcher
 
 from .qfunc_interpreter import PLxPRToQuantumJaxprInterpreter
@@ -286,7 +283,9 @@ def handle_qnode(
         # Fallback to the legacy decomposition if the graph-based decomposition failed
         if not graph_succeeded:
             # Remove the decompose-lowering pass from the pipeline
-            self._pass_pipeline = [p for p in self._pass_pipeline if p.name != "decompose-lowering"]
+            self._pass_pipeline = [
+                p for p in self._pass_pipeline if p.pass_name != "decompose-lowering"
+            ]
             closed_jaxpr = _apply_compiler_decompose_to_plxpr(
                 inner_jaxpr=closed_jaxpr.jaxpr,
                 consts=closed_jaxpr.consts,
@@ -334,11 +333,9 @@ def calling_convention(*args):
 # otherwise their value will be None. The second value indicates if the transform
 # requires decomposition to be supported by Catalyst.
 transforms_to_passes = {
-    pl_cancel_inverses: ("cancel-inverses", False),
     pl_commute_controlled: (None, False),
     pl_decompose: (None, False),
     pl_merge_amplitude_embedding: (None, True),
-    pl_merge_rotations: ("merge-rotations", False),
     pl_single_qubit_fusion: (None, False),
     pl_unitary_to_rot: (None, False),
 }
@@ -349,6 +346,47 @@ def register_transform(pl_transform, pass_name, decomposition):
     transforms_to_passes[pl_transform] = (pass_name, decomposition)
 
 
+def _handle_decompose_transform(self, inner_jaxpr, consts, non_const_args, tkwargs):
+    if not self.requires_decompose_lowering:
+        self.requires_decompose_lowering = True
+    else:
+        raise NotImplementedError("Multiple decomposition transforms are not yet supported.")
+
+    next_eval = copy(self)
+    # Update the decompose_gateset to be used by the quantum kernel primitive
+    # TODO: we originally wanted to treat decompose_gateset as a queue of
+    # gatesets to be used by the decompose-lowering pass at MLIR
+    # but this requires a C++ implementation of the graph-based decomposition
+    # which doesn't exist yet.
+    next_eval.decompose_tkwargs = tkwargs
+
+    # Note. We don't perform the compiler-specific decomposition here
+    # to be able to support multiple decomposition transforms
+    # and collect all the required gatesets
+    # as well as being able to support other transforms in between.
+
+    # The compiler specific transformation will be performed
+    # in the qnode handler.
+
+    # Add the decompose-lowering pass to the start of the pipeline
+    t = qml.transform(pass_name="decompose-lowering")
+    pass_container = qml.transforms.core.TransformContainer(t)
+    next_eval._pass_pipeline.insert(0, pass_container)
+
+    # We still need to construct and solve the graph based on
+    # the current jaxpr based on the current gateset
+    # but we don't rewrite the jaxpr at this stage.
+
+    # gds_interpreter = DecompRuleInterpreter(*targs, **tkwargs)
+
+    # def gds_wrapper(*args):
+    #     return gds_interpreter.eval(inner_jaxpr, consts, *args)
+
+    # final_jaxpr = jax.make_jaxpr(gds_wrapper)(*args)
+    # return self.eval(final_jaxpr.jaxpr, consts, *non_const_args)
+    return next_eval.eval(inner_jaxpr, consts, *non_const_args)
+
+
 # pylint: disable=too-many-arguments
 @WorkflowInterpreter.register_primitive(transform_prim)
 def handle_transform(
@@ -375,45 +413,11 @@ def handle_transform(
         and transform._plxpr_transform.__name__ == "decompose_plxpr_to_plxpr"
         and qml.decomposition.enabled_graph()
     ):
-        # Handle the conversion from plxpr to Catalyst jaxpr for a PL transform.
-        if not self.requires_decompose_lowering:
-            self.requires_decompose_lowering = True
-        else:
-            raise NotImplementedError("Multiple decomposition transforms are not yet supported.")
-
-        next_eval = copy(self)
-        # Update the decompose_gateset to be used by the quantum kernel primitive
-        # TODO: we originally wanted to treat decompose_gateset as a queue of
-        # gatesets to be used by the decompose-lowering pass at MLIR
-        # but this requires a C++ implementation of the graph-based decomposition
-        # which doesn't exist yet.
-        next_eval.decompose_tkwargs = tkwargs
-
-        # Note. We don't perform the compiler-specific decomposition here
-        # to be able to support multiple decomposition transforms
-        # and collect all the required gatesets
-        # as well as being able to support other transforms in between.
-
-        # The compiler specific transformation will be performed
-        # in the qnode handler.
-
-        # Add the decompose-lowering pass to the start of the pipeline
-        next_eval._pass_pipeline.insert(0, Pass("decompose-lowering"))
+        return _handle_decompose_transform(self, inner_jaxpr, consts, non_const_args, tkwargs)
 
-        # We still need to construct and solve the graph based on
-        # the current jaxpr based on the current gateset
-        # but we don't rewrite the jaxpr at this stage.
-
-        # gds_interpreter = DecompRuleInterpreter(*targs, **tkwargs)
-
-        # def gds_wrapper(*args):
-        #     return gds_interpreter.eval(inner_jaxpr, consts, *args)
-
-        # final_jaxpr = jax.make_jaxpr(gds_wrapper)(*args)
-        # return self.eval(final_jaxpr.jaxpr, consts, *non_const_args)
-        return next_eval.eval(inner_jaxpr, consts, *non_const_args)
-
-    catalyst_pass_name = transforms_to_passes.get(transform, (None,))[0]
+    catalyst_pass_name = transform.pass_name
+    if catalyst_pass_name is None:
+        catalyst_pass_name = transforms_to_passes.get(transform, (None,))[0]
     if catalyst_pass_name is None:
         # Use PL's ExpandTransformsInterpreter to expand this and any embedded
         # transform according to PL rules. It works by overriding the primitive
@@ -435,7 +439,9 @@ def wrapper(*args):
 
     # Apply the corresponding Catalyst pass counterpart
     next_eval = copy(self)
-    next_eval._pass_pipeline.insert(0, Pass(catalyst_pass_name, *targs, **tkwargs))
+    t = qml.transform(pass_name=catalyst_pass_name)
+    bound_pass = qml.transforms.core.TransformContainer(t, args=targs, kwargs=tkwargs)
+    next_eval._pass_pipeline.insert(0, bound_pass)
     return next_eval.eval(inner_jaxpr, consts, *non_const_args)
 
 

frontend/catalyst/jax_primitives_utils.py

@@ -325,6 +325,16 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.ctx.module_context = self.old_module_context
 
 
+def _lowered_options(args, kwargs):
+    lowered_options = {}
+    for arg in args:
+        lowered_options[str(arg)] = get_mlir_attribute_from_pyval(True)
+    for option, value in kwargs.items():
+        mlir_option = str(option).replace("_", "-")
+        lowered_options[mlir_option] = get_mlir_attribute_from_pyval(value)
+    return lowered_options
+
+
 def transform_named_sequence_lowering(jax_ctx: mlir.LoweringRuleContext, pipeline):
     """Generate a transform module embedded in the current module and schedule
     the transformations in pipeline"""
@@ -371,11 +381,16 @@ def transform_named_sequence_lowering(jax_ctx: mlir.LoweringRuleContext, pipelin
         with ir.InsertionPoint(bb_named_sequence):
             target = bb_named_sequence.arguments[0]
             for _pass in pipeline:
-                options = _pass.get_options()
+                if isinstance(_pass, qml.transforms.core.TransformContainer):
+                    options = _lowered_options(_pass.args, _pass.kwargs)
+                    name = _pass.pass_name
+                else:
+                    options = _pass.get_options()
+                    name = _pass.name
                 apply_registered_pass_op = ApplyRegisteredPassOp(
                     result=transform_mod_type,
                     target=target,
-                    pass_name=_pass.name,
+                    pass_name=name,
                     options=options,
                     dynamic_options={},
                 )
@@ -387,7 +402,7 @@ def transform_named_sequence_lowering(jax_ctx: mlir.LoweringRuleContext, pipelin
                         is_xdsl_pass,
                     )
 
-                    if is_xdsl_pass(_pass.name):
+                    if is_xdsl_pass(name):
                         uses_xdsl_passes = True
                         apply_registered_pass_op.operation.attributes["catalyst.xdsl_pass"] = (
                             ir.UnitAttr.get()

frontend/catalyst/qfunc.py

@@ -285,18 +285,23 @@ def __call__(self, *args, **kwargs):
 
         assert isinstance(self, qml.QNode)
 
+        new_transform_program, new_pipeline = _extract_passes(self.transform_program)
         # Update the qnode with peephole pipeline
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline = dictionary_to_list_of_passes(pass_pipeline)
+        old_pipeline = kwargs.pop("pass_pipeline", None)
+        processed_old_pipeline = tuple(dictionary_to_list_of_passes(old_pipeline))
+        pass_pipeline = processed_old_pipeline + new_pipeline
+        new_qnode = copy(self)
+        # pylint: disable=attribute-defined-outside-init, protected-access
+        new_qnode._transform_program = new_transform_program
 
         # Mid-circuit measurement configuration/execution
-        fn_result = configure_mcm_and_try_one_shot(self, args, kwargs, pass_pipeline)
+        fn_result = configure_mcm_and_try_one_shot(new_qnode, args, kwargs, pass_pipeline)
 
         # If the qnode is failed to execute as one-shot, fn_result will be None
         if fn_result is not None:
             return fn_result
 
-        new_device = copy(self.device)
+        new_device = copy(new_qnode.device)
         qjit_device = QJITDevice(new_device)
 
         static_argnums = kwargs.pop("static_argnums", ())
@@ -307,11 +312,11 @@ def __call__(self, *args, **kwargs):
 
         def _eval_quantum(*args, **kwargs):
             trace_result = trace_quantum_function(
-                self.func,
+                new_qnode.func,
                 qjit_device,
                 args,
                 kwargs,
-                self,
+                new_qnode,
                 static_argnums,
                 debug_info,
             )
@@ -655,3 +660,22 @@ def wrap_single_shot_qnode(*_):
         return _finalize_output(out, ctx)
 
     return one_shot_wrapper
+
+
+def _extract_passes(transform_program):
+    """Extract transforms with pass names from the end of the TransformProgram."""
+    tape_transforms = []
+    pass_pipeline = []
+    i = len(transform_program)
+    for t in reversed(transform_program):
+        if t.pass_name is None:
+            break
+        i -= 1
+    pass_pipeline = transform_program[i:]
+    tape_transforms = transform_program[:i]
+    for t in tape_transforms:
+        if t.transform is None:
+            raise ValueError(
+                f"{t} without a tape definition occurs before tape transform {tape_transforms[-1]}."
+            )
+    return qml.transforms.core.TransformProgram(tape_transforms), tuple(pass_pipeline)

frontend/test/lit/test_decomposition.py

@@ -46,6 +46,7 @@ def wrapper():
             error_msg = str(e)
             if (
                 "Unsupported type annotation None for parameter pauli_word" in error_msg
+                or "Unsupported type annotation <class 'str'> for parameter pauli_word" in error_msg
                 or "index is out of bounds for axis" in error_msg
             ):
                 print(f"# SKIPPED {test_func.__name__}: PauliRot type annotation issue")

frontend/test/lit/test_peephole_optimizations.py

@@ -87,6 +87,42 @@ def test_pipeline_lowering_workflow(x):
 test_pipeline_lowering()
 
 
+def test_transform_lowering():
+    """
+    Basic pipeline lowering on one qnode.
+    """
+
+    @qjit(keep_intermediate=True)
+    @qml.transforms.merge_rotations
+    @qml.transforms.cancel_inverses
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def test_pipeline_lowering_workflow(x):
+        qml.RX(x, wires=[0])
+        qml.Hadamard(wires=[1])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    # CHECK: pipeline=(<cancel_inverses((), {})>, <merge_rotations((), {})>)
+    print_jaxpr(test_pipeline_lowering_workflow, 1.2)
+
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "cancel-inverses" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print_mlir(test_pipeline_lowering_workflow, 1.2)
+
+    # CHECK: {{%.+}} = call @test_pipeline_lowering_workflow_0(
+    # CHECK: func.func public @test_pipeline_lowering_workflow_0(
+    # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    test_pipeline_lowering_workflow(42.42)
+    flush_peephole_opted_mlir_to_iostream(test_pipeline_lowering_workflow)
+
+
+test_transform_lowering()
+
+
 def test_pipeline_lowering_keep_original():
     """
     Test when the pipelined qnode and the original qnode are both used,

frontend/test/pytest/from_plxpr/test_capture_integration.py

@@ -1048,7 +1048,7 @@ def circuit(x: float):
         assert jnp.allclose(circuit(0.1), capture_result)
 
     @pytest.mark.usefixtures("use_capture")
-    def test_pass_with_options(self, backend):
+    def test_pass_with_options_patch(self, backend):
         """Test the integration for a circuit with a pass that takes in options."""
 
         @qml.transform
@@ -1071,6 +1071,25 @@ def captured_circuit():
             in capture_mlir
         )
 
+    @pytest.mark.usefixtures("use_capture")
+    def test_pass_with_options(self, backend):
+        """Test the integration for a circuit with a pass that takes in options."""
+
+        my_pass = qml.transform(pass_name="my-pass")
+
+        @qjit(target="mlir")
+        @partial(my_pass, my_option="my_option_value", my_other_option=False)
+        @qml.qnode(qml.device(backend, wires=1))
+        def captured_circuit():
+            return qml.expval(qml.PauliZ(0))
+
+        capture_mlir = captured_circuit.mlir
+        assert 'transform.apply_registered_pass "my-pass"' in capture_mlir
+        assert (
+            'with options = {"my-option" = "my_option_value", "my-other-option" = false}'
+            in capture_mlir
+        )
+
     def test_transform_cancel_inverses_workflow(self, backend):
         """Test the integration for a circuit with a 'cancel_inverses' transform."""