Create a frontend UI for users to specify quantum compilation pipelines

frontend/catalyst/api_extensions/error_mitigation.py

@@ -27,7 +27,8 @@
 import pennylane as qml
 from jax._src.tree_util import tree_flatten
 
-from catalyst.jax_primitives import Folding, zne_p
+from catalyst.jax_primitives import Folding, apply_registered_pass_p, zne_p
+from catalyst.tracing.contexts import EvaluationContext
 
 
 def _is_odd_positive(numbers_list):
@@ -164,6 +165,15 @@ def __init__(
     ):
         if not isinstance(fn, qml.QNode):
             raise TypeError(f"A QNode is expected, got the classical function {fn}")
+
+        wrapped_qnode_function = fn.func
+
+        def wrapper(*args, **kwrags):
+            if EvaluationContext.is_tracing():
+                apply_registered_pass_p.bind(pass_name="lower-mitigation")
+            return wrapped_qnode_function(*args, **kwrags)
+
+        fn.func = wrapper
         self.fn = fn
         self.__name__ = f"zne.{getattr(fn, '__name__', 'unknown')}"
         self.num_folds = num_folds

frontend/catalyst/compiler.py

@@ -196,7 +196,6 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
     [
         "apply-transform-sequence",  # Run the transform sequence defined in the MLIR module
         "annotate-function",
-        "lower-mitigation",
         "lower-gradients",
         "adjoint-lowering",
         "disable-assertion",

frontend/catalyst/jax_primitives.py

@@ -514,7 +514,8 @@
             """
 
     # If there already is a apply_registered_pass,
-    # insert after the last pass in the existing pass sequence.
+    # insert before the first pass in the existing pass sequence.
+    # See comment in frontend/catalyst/passes.py/pipeline()
     # Note that ir.InsertionPoint(op) sets the insertion point to immediately BEFORE the op
     named_sequence_op_block = named_sequence_op.regions[0].blocks[0]
     first_op_in_block = named_sequence_op_block.operations[0].operation
@@ -528,16 +529,28 @@
         """
 
     if first_op_in_block.name == "transform.apply_registered_pass":
+        """
         _ = len(named_sequence_op_block.operations)
         yield_op = named_sequence_op_block.operations[_ - 1].operation
         current_last_pass = named_sequence_op_block.operations[_ - 2].operation
+
         with ir.InsertionPoint(yield_op):
             apply_registered_pass_op = ApplyRegisteredPassOp(
                 result=transform_mod_type,
                 target=current_last_pass.result,
                 pass_name=pass_name,
                 options=options,
             )
+        """
+        current_first_pass = named_sequence_op_block.operations[0].operation
+        with ir.InsertionPoint(first_op_in_block):
+            apply_registered_pass_op = ApplyRegisteredPassOp(
+                result=transform_mod_type,
+                target=named_sequence_op.regions[0].blocks[0].arguments[0],
+                pass_name=pass_name,
+                options=options,
+            )
+            current_first_pass.operands[0] = apply_registered_pass_op.result
 
     # otherwise it's the first pass, i.e. only a yield op is in the block
     # so insert right before the yield op

frontend/catalyst/passes.py

@@ -33,15 +33,65 @@
 """
 
 import copy
+import functools
 
 import pennylane as qml
 
+from catalyst.api_extensions import (
+    mitigate_with_zne as _mitigate_with_zne_api_extensions,
+)
 from catalyst.jax_primitives import apply_registered_pass_p, transform_named_sequence_p
+from catalyst.tracing.contexts import EvaluationContext
 
 
 ## API ##
 # pylint: disable=line-too-long
-def cancel_inverses(fn=None):
+def pipeline(fn=None, *, pass_pipeline=None):
+    """
+    Here are documentation words
+    """
+
+    """
+    Implementation design: it just stacks the decorators for the user automatically.
+    e.g.pass_pipeline = {
+            "mitigate_with_zne": {"scale_factors": [1, 2, 3]},
+            "cancel_inverses": {}
+        }
+    will just do 
+    fn = mitigate_with_zne(scale_factors=[1, 2, 3])(fn)
+    fn = cancel_inverses(fn)
+
+    Note that since python 3.7, dictionary order is guaranteed to be in insertion order
+    https://docs.python.org/3/library/stdtypes.html#dict.values
+    """
+    kwargs = copy.copy(locals())
+    kwargs.pop("fn")
+
+    if fn is None:
+        return functools.partial(pipeline, **kwargs)
+
+    if not isinstance(fn, qml.QNode):
+        raise TypeError(f"A QNode is expected, got the classical function {fn}")
+
+    fn_original_name = fn.__name__
+
+    API_calls = API_name_to_API_calls()
+    fn_clone = copy.copy(fn)
+    fn_clone.__name__ = fn_original_name + "_transformed"
+    # Note: we create wrappers to inject the apply_registered_pass primitive
+    # In other words, the last API call, aka the last wrapper, will be the outermost primitive
+    # Therefore when lowering these primitives to mlir (in jax_primitives.py), lower them in reverse order!
+    for decorator, decorator_args in pass_pipeline.items():
+        print("seeing pass: ", decorator)
+        if decorator not in API_calls.keys():
+            raise RuntimeError(f"{decorator} is not a valid quantum transformation pass")
+
+        fn_clone = API_calls[decorator](fn_clone, keep_original=False, **decorator_args)
+
+    return fn_clone
+
+
+def cancel_inverses(fn=None, keep_original=True):
     """
     Specify that the ``-removed-chained-self-inverse`` MLIR compiler pass
     for cancelling two neighbouring self-inverse
@@ -137,33 +187,61 @@
     if not isinstance(fn, qml.QNode):
         raise TypeError(f"A QNode is expected, got the classical function {fn}")
 
-    wrapped_qnode_function = fn.func
     funcname = fn.__name__
+    wrapped_qnode_function = fn.func
 
-    def wrapper(*args, **kwrags):
-        # TODO: hint the compiler which qnodes to run the pass on via an func attribute,
-        # instead of the qnode name. That way the clone can have this attribute and
-        # the original can just not have it.
-        # We are not doing this right now and passing by name because this would
-        # be a discardable attribute (i.e. a user/developer wouldn't know that this
-        # attribute exists just by looking at qnode's documentation)
-        # But when we add the full peephole pipeline in the future, the attribute
-        # could get properly documented.
-
-        apply_registered_pass_p.bind(
-            pass_name="remove-chained-self-inverse",
-            options=f"func-name={funcname}" + "_cancel_inverses",
-        )
-        return wrapped_qnode_function(*args, **kwrags)
+    if keep_original:
 
-    fn_clone = copy.copy(fn)
-    fn_clone.func = wrapper
-    fn_clone.__name__ = funcname + "_cancel_inverses"
+        def wrapper(*args, **kwrags):
+            # TODO: hint the compiler which qnodes to run the pass on via an func attribute,
+            # instead of the qnode name. That way the clone can have this attribute and
+            # the original can just not have it.
+            # We are not doing this right now and passing by name because this would
+            # be a discardable attribute (i.e. a user/developer wouldn't know that this
+            # attribute exists just by looking at qnode's documentation)
+            # But when we add the full peephole pipeline in the future, the attribute
+            # could get properly documented.
 
-    return fn_clone
+            if EvaluationContext.is_tracing():
+                apply_registered_pass_p.bind(
+                    pass_name="remove-chained-self-inverse",
+                    options=f"func-name={funcname}" + "_cancel_inverses",
+                )
+            return wrapped_qnode_function(*args, **kwrags)
+
+        fn_clone = copy.copy(fn)
+        fn_clone.func = wrapper
+        fn_clone.__name__ = funcname + "_cancel_inverses"
+
+        return fn_clone
+
+    else:
+
+        def wrapper(*args, **kwrags):
+            if EvaluationContext.is_tracing():
+                apply_registered_pass_p.bind(
+                    pass_name="remove-chained-self-inverse",
+                    options=f"func-name={funcname}",
+                )
+            return wrapped_qnode_function(*args, **kwrags)
+
+        fn.func = wrapper
+        return fn
+
+
+def mitigate_with_zne(*args, keep_original=False, **kwrags):
+    """
+    An alias of catalyst.mitigate_with_zne.
+    See https://docs.pennylane.ai/projects/catalyst/en/stable/code/api/catalyst.mitigate_with_zne.html
+    """
+    return _mitigate_with_zne_api_extensions(*args, **kwrags)
 
 
 ## IMPL and helpers ##
+def API_name_to_API_calls():
+    return {"cancel_inverses": cancel_inverses, "mitigate_with_zne": mitigate_with_zne}
+
+
 def _inject_transform_named_sequence():
     """
     Inject a transform_named_sequence jax primitive.

frontend/test/lit/test_peephole_optimizations.py

@@ -30,8 +30,9 @@
 from lit_util_printers import print_jaxpr, print_mlir
 
 from catalyst import qjit
+from catalyst.api_extensions.error_mitigation import polynomial_extrapolation
 from catalyst.debug import get_compilation_stage
-from catalyst.passes import cancel_inverses
+from catalyst.passes import cancel_inverses, pipeline
 
 
 def flush_peephole_opted_mlir_to_iostream(QJIT):
@@ -75,6 +76,123 @@
 test_transform_named_sequence_injection()
 
 
+#
+# pipeline
+#
+
+
+def test_pipeline_lowering():
+    my_pipeline = {
+        "cancel_inverses": {},
+        "mitigate_with_zne": {
+            "scale_factors": [1, 3, 5, 7],
+            "extrapolate": polynomial_extrapolation(2),
+            "folding": "global",
+        },
+    }
+
+    @qjit(keep_intermediate=True)
+    @pipeline(pass_pipeline=my_pipeline)
+    @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: transform_named_sequence
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   pass_name=lower-mitigation
+    # CHECK: ]
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   options=func-name=test_pipeline_lowering_workflow_transformed
+    # CHECK:   pass_name=remove-chained-self-inverse
+    # CHECK: ]
+    print_jaxpr(test_pipeline_lowering_workflow, 1.2)
+
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=test_pipeline_lowering_workflow_transformed"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "lower-mitigation" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print_mlir(test_pipeline_lowering_workflow, 1.2)
+
+    # CHECK: func.func private @test_pipeline_lowering_workflow_transformed.withMeasurements
+    # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK: func.func public @jit_zne.test_pipeline_lowering_workflow_transformed
+    # CHECK: func.func private @test_pipeline_lowering_workflow_transformed
+    # 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_pipeline_lowering()
+
+
+def test_pipeline_lowering_keep_original():
+    my_pipeline = {
+        "cancel_inverses": {},
+        "mitigate_with_zne": {
+            "scale_factors": [1, 3, 5, 7],
+            "extrapolate": polynomial_extrapolation(2),
+            "folding": "global",
+        },
+    }
+
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def f(x):
+        qml.RX(x, wires=[0])
+        qml.Hadamard(wires=[1])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    f_pipeline = pipeline(pass_pipeline=my_pipeline)(f)
+
+    @qjit(keep_intermediate=True)
+    def test_pipeline_lowering_keep_original_workflow(x):
+        return f(1.2), f_pipeline(1.2)
+
+    # CHECK: transform_named_sequence
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   pass_name=lower-mitigation
+    # CHECK: ]
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   options=func-name=f_transformed
+    # CHECK:   pass_name=remove-chained-self-inverse
+    # CHECK: ]
+    print_jaxpr(test_pipeline_lowering_keep_original_workflow, 1.2)
+
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=f_transformed"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "lower-mitigation" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print_mlir(test_pipeline_lowering_keep_original_workflow, 1.2)
+
+    # CHECK: func.func private @f_transformed.withMeasurements
+    # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK: func.func public @jit_test_pipeline_lowering_keep_original_workflow
+    # CHECK: {{%.+}} = call @f(
+    # CHECK: {{%.+}} = func.call @f_transformed.folded(
+    # CHECK: func.func private @f(
+    # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
+    # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK: func.func private @f_transformed
+    # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
+    test_pipeline_lowering_keep_original_workflow(42.42)
+    flush_peephole_opted_mlir_to_iostream(test_pipeline_lowering_keep_original_workflow)
+
+
+test_pipeline_lowering_keep_original()
+
+
 #
 # cancel_inverses
 #
@@ -132,8 +250,8 @@
 
     # CHECK: module @test_cancel_inverses_tracing_and_lowering_workflow
     # CHECK: transform.named_sequence @__transform_main
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=f_cancel_inverses"}
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=g_cancel_inverses"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=f_cancel_inverses"}
     # CHECK-NOT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=h_cancel_inverses"}
     # CHECK-NEXT: transform.yield
     print_mlir(test_cancel_inverses_tracing_and_lowering_workflow, 1.1)

frontend/test/pytest/test_mitigation.py

@@ -62,6 +62,17 @@ def mitigated_qnode(args):
 
     assert np.allclose(mitigated_qnode(params), circuit(params))
 
+    @catalyst.qjit
+    def mitigated_qnode_passes(args):
+        return catalyst.passes.mitigate_with_zne(
+            circuit,
+            scale_factors=scale_factors,
+            extrapolate=extrapolation,
+            folding=folding,
+        )(args)
+
+    assert np.allclose(mitigated_qnode_passes(params), circuit(params))
+
 
 @pytest.mark.parametrize("params", [0.1, 0.2, 0.3, 0.4, 0.5])
 @pytest.mark.parametrize("extrapolation", [quadratic_extrapolation, exponential_extrapolate])