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

@@ -90,6 +90,7 @@ class CompileOptions:
     checkpoint_stage: Optional[str] = ""
     disable_assertions: Optional[bool] = False
     seed: Optional[int] = None
+    circuit_transform_pipeline: Optional[dict] = None
 
     def __post_init__(self):
         # Check that async runs must not be seeded
@@ -196,7 +197,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/jit.py

@@ -38,6 +38,7 @@
 from catalyst.jax_tracer import lower_jaxpr_to_mlir, trace_to_jaxpr
 from catalyst.logging import debug_logger, debug_logger_init
 from catalyst.passes import _inject_transform_named_sequence
+from catalyst.passes import pipeline as circuit_transform_pass_pipeline
 from catalyst.qfunc import QFunc
 from catalyst.tracing.contexts import EvaluationContext
 from catalyst.tracing.type_signatures import (
@@ -83,6 +84,7 @@
     abstracted_axes=None,
     disable_assertions=False,
     seed=None,
+    circuit_transform_pipeline=None,
 ):  # pylint: disable=too-many-arguments,unused-argument
     """A just-in-time decorator for PennyLane and JAX programs using Catalyst.
 
@@ -585,7 +587,12 @@
             params = {}
             params["static_argnums"] = kwargs.pop("static_argnums", static_argnums)
             params["_out_tree_expected"] = []
-            return QFunc.__call__(qnode, *args, **dict(params, **kwargs))
+            return QFunc.__call__(
+                qnode,
+                self.compile_options.circuit_transform_pipeline,
+                *args,
+                **dict(params, **kwargs),
+            )
 
         with Patcher(
             (qml.QNode, "__call__", closure),

frontend/catalyst/passes.py

@@ -33,15 +33,69 @@
 """
 
 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}")
+
+    if pass_pipeline is None:
+        # TODO: design a default peephole pipeline
+        return fn
+
+    API_calls = API_name_to_API_calls()
+
+    fn_original_name = fn.__name__
+    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():
+        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 +191,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/catalyst/qfunc.py

@@ -56,6 +56,7 @@
 from catalyst.jax_primitives import func_p
 from catalyst.jax_tracer import trace_quantum_function
 from catalyst.logging import debug_logger
+from catalyst.passes import pipeline
 from catalyst.tracing.type_signatures import filter_static_args
 from catalyst.utils.toml import DeviceCapabilities, ProgramFeatures
 
@@ -109,9 +110,11 @@
 
     # pylint: disable=no-member
     @debug_logger
-    def __call__(self, *args, **kwargs):
+    def __call__(self, pass_pipeline, *args, **kwargs):
         assert isinstance(self, qml.QNode)
 
+        self = pipeline(pass_pipeline=pass_pipeline)(self)
+
         # Mid-circuit measurement configuration/execution
         dynamic_one_shot_called = getattr(self, "_dynamic_one_shot_called", False)
         if not dynamic_one_shot_called: