Source code for catalyst.passes.pass_api

# Copyright 2024 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools
from importlib.metadata import entry_points
from pathlib import Path
from typing import TypeAlias

import pennylane as qml

from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
from catalyst.tracing.contexts import EvaluationContext

PipelineDict: TypeAlias = dict[str, dict[str, str]]


## API ##
class QNodeWrapper:
    """A wrapper class for QNodes that preserves the pass pipeline when called.

    This class is used internally by pass decorators to ensure that compiler passes
    are properly applied when the QNode is called.

    Args:
        qnode (QNode or callable): The QNode or wrapped QNode function to call
    """

    def __init__(self, qnode):
        self.qnode = qnode
        if isinstance(qnode, QNodeWrapper):
            self.original_qnode = qnode.original_qnode
            self.func = qnode.func
        elif isinstance(qnode, qml.QNode):
            self.original_qnode = qnode
            self.func = qnode.func
        else:
            raise TypeError(f"A QNode is expected, got the classical function {qnode}")

        # Update the wrapper to preserve the original function's metadata (signature, annotations,
        # etc.) This is crucial for qjit to properly inspect the function signature and make
        # compilation decisions.
        functools.update_wrapper(self, self.func)

    def __call__(self, *args, **kwargs):
        return self.qnode(*args, **kwargs)


class PassPipelineWrapper(QNodeWrapper):
    """A QNodeWrapper subclass that adds passes to the pipeline when called.

    This class can handle both single passes and multiple passes in a pipeline.
    For single passes, it takes a pass name and optional flags/options.
    For multiple passes, it takes a dictionary of pass configurations.
    """

    def __init__(self, qnode, pass_name_or_pipeline, *flags, **valued_options):
        super().__init__(qnode)
        self.pass_name_or_pipeline = pass_name_or_pipeline
        self.flags = flags
        self.valued_options = valued_options

    def __call__(self, *args, **kwargs):
        if EvaluationContext.is_tracing():
            pass_pipeline = kwargs.pop("pass_pipeline", [])
            pass_pipeline = (
                dictionary_to_list_of_passes(
                    self.pass_name_or_pipeline, *self.flags, **self.valued_options
                )
                + pass_pipeline
            )
            kwargs["pass_pipeline"] = pass_pipeline
        return super().__call__(*args, **kwargs)


[docs] def pipeline(pass_pipeline: PipelineDict): """Configures the Catalyst MLIR pass pipeline for quantum circuit transformations for a QNode within a qjit-compiled program. Args: pass_pipeline (dict[str, dict[str, str]]): A dictionary that specifies the pass pipeline order, and optionally arguments for each pass in the pipeline. Keys of this dictionary should correspond to names of passes found in the `catalyst.passes <https://docs.pennylane.ai/projects/catalyst/en/stable/code/__init__.html#module-catalyst.passes>`_ module, values should either be empty dictionaries (for default pass options) or dictionaries of valid keyword arguments and values for the specific pass. The order of keys in this dictionary will determine the pass pipeline. If not specified, the default pass pipeline will be applied. Returns: callable : A decorator that can be applied to a qnode. For a list of available passes, please see the :doc:`catalyst.passes module <code/passes>`. The default pass pipeline when used with Catalyst is currently empty. **Example** ``pipeline`` can be used to configure the pass pipeline order and options of a QNode within a qjit-compiled function. Configuration options are passed to specific passes via dictionaries: .. code-block:: python my_pass_pipeline = { "cancel_inverses": {}, "my_circuit_transformation_pass": {"my-option" : "my-option-value"}, } @pipeline(my_pass_pipeline) @qnode(dev) def circuit(x): qml.RX(x, wires=0) return qml.expval(qml.PauliZ(0)) @qjit def fn(x): return jnp.sin(circuit(x ** 2)) ``pipeline`` can also be used to specify different pass pipelines for different parts of the same qjit-compiled workflow: .. code-block:: python my_pipeline = { "cancel_inverses": {}, "my_circuit_transformation_pass": {"my-option" : "my-option-value"}, } my_other_pipeline = {"cancel_inverses": {}} @qjit def fn(x): circuit_pipeline = pipeline(my_pipeline)(circuit) circuit_other = pipeline(my_other_pipeline)(circuit) return jnp.abs(circuit_pipeline(x) - circuit_other(x)) .. note:: As of Python 3.7, the CPython dictionary implementation orders dictionaries based on insertion order. However, for an API guarantee of dictionary order, ``collections.OrderedDict`` may also be used. Note that the pass pipeline order and options can be configured *globally* for a qjit-compiled function, by using the ``circuit_transform_pipeline`` argument of the :func:`~.qjit` decorator. .. code-block:: python my_pass_pipeline = { "cancel_inverses": {}, "my_circuit_transformation_pass": {"my-option" : "my-option-value"}, } @qjit(circuit_transform_pipeline=my_pass_pipeline) def fn(x): return jnp.sin(circuit(x ** 2)) Global and local (via ``@pipeline``) configurations can coexist, however local pass pipelines will always take precedence over global pass pipelines. """ def _decorator(qnode): return PassPipelineWrapper(qnode, pass_pipeline) return _decorator
[docs] def apply_pass(pass_name: str, *flags, **valued_options): """Applies a single pass to the QNode, where the pass is from Catalyst or a third-party if `entry_points` has been implemented. See :doc:`the compiler plugin documentation <dev/plugins>` for more details. Args: pass_name (str): Name of the pass *flags: Pass options **valued_options: options with values Returns: Function that can be used as a decorator to a QNode. E.g., .. code-block:: python @passes.apply_pass("merge-rotations") @qml.qnode(qml.device("lightning.qubit", wires=1)) def qnode(): return qml.state() @qml.qjit(target="mlir") def module(): return qnode() """ def decorator(qnode): return PassPipelineWrapper(qnode, pass_name, *flags, **valued_options) return decorator
[docs] def apply_pass_plugin(path_to_plugin: str | Path, pass_name: str, *flags, **valued_options): """Applies a pass plugin to the QNode. See :doc:`the compiler plugin documentation <dev/plugins>` for more details. Args: path_to_plugin (str | Path): full path to plugin pass_name (str): Name of the pass *flags: Pass options **valued_options: options with values Returns: Function that can be used as a decorator to a QNode. E.g., .. code-block:: python from standalone import getStandalonePluginAbsolutePath @passes.apply_pass_plugin(getStandalonePluginAbsolutePath(), "standalone-switch-bar-foo") @qml.qnode(qml.device("lightning.qubit", wires=1)) def qnode(): return qml.state() @qml.qjit(target="mlir") def module(): return qnode() """ if not isinstance(path_to_plugin, Path): path_to_plugin = Path(path_to_plugin) if not path_to_plugin.exists(): raise FileNotFoundError(f"File '{path_to_plugin}' does not exist.") def decorator(qnode): return PassPipelineWrapper(qnode, pass_name, *flags, **valued_options) return decorator
[docs] class Pass: """Class intended to hold options for passes. :class:`Pass` will be used when generating `ApplyRegisteredPassOp`s. The attribute `pass_name` corresponds to the field `name`. The attribute `options` is generated by the `get_options` method. People working on MLIR plugins may use this or :class:`PassPlugin` to schedule their compilation pass. E.g., .. code-block:: python def an_optimization(qnode): @functools.wraps(qnode) def wrapper(*args, **kwargs): pass_pipeline = kwargs.pop("pass_pipeline", []) pass_pipeline.append(Pass("my_library.my_optimization", *args, **kwargs)) kwargs["pass_pipeline"] = pass_pipeline return qnode(*args, **kwargs) return wrapper """ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, str]): self.options = options self.valued_options = valued_options if "." in name: resolution_functions = entry_points(group="catalyst.passes_resolution") key, passname = name.split(".") resolution_function = resolution_functions[key + ".passes"] module = resolution_function.load() path, name = module.name2pass(passname) assert EvaluationContext.is_tracing() EvaluationContext.add_plugin(path) self.name = name
[docs] def get_options(self): """ Build a dictionary mapping option names to MLIR attributes. ApplyRegisteredPassOp expects options to be a dictionary from strings to attributes. See https://github.com/llvm/llvm-project/pull/143159 """ options_dict = {} for option in self.options: options_dict[str(option)] = get_mlir_attribute_from_pyval(True) for option, value in self.valued_options.items(): # MLIR options are either CamelCase # or spinal-case (kebab-case) which is not allowed in python mlir_option = str(option).replace("_", "-") options_dict[mlir_option] = get_mlir_attribute_from_pyval(value) return options_dict
def __repr__(self): return ( self.name + " ".join(f"--{str(option)}" for option in self.options) + " ".join( [f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()] ) )
[docs] class PassPlugin(Pass): """Similar to :class:`Pass` but takes into account the plugin. The plugin is used during the creation of the compilation command. E.g., --pass-plugin=path/to/plugin --dialect-plugin=path/to/plugin People working on MLIR plugins may use this or :class:`Pass` to schedule their compilation pass. E.g., .. code-block:: python def an_optimization(qnode): @functools.wraps(qnode) def wrapper(*args, **kwargs): pass_pipeline = kwargs.pop("pass_pipeline", []) pass_pipeline.append(PassPlugin(path_to_plugin, "my_optimization", *args, **kwargs)) kwargs["pass_pipeline"] = pass_pipeline return qnode(*args, **kwargs) return wrapper """ def __init__( self, path: Path, name: str, *options: list[str], **valued_options: dict[str, str] ): assert EvaluationContext.is_tracing() EvaluationContext.add_plugin(path) self.path = path super().__init__(name, *options, **valued_options)
## PRIVATE ## def dictionary_to_list_of_passes(pass_pipeline: PipelineDict | str, *flags, **valued_options): """Convert dictionary of passes or single pass name into list of passes. Args: pass_pipeline: Either a dictionary of pass configurations or a single pass name. *flags: Optional flags for single pass **valued_options: Optional valued options for single pass """ if pass_pipeline is None: return [] if isinstance(pass_pipeline, str): return [Pass(pass_pipeline, *flags, **valued_options)] if isinstance(pass_pipeline, dict): passes = [] pass_names = _API_name_to_pass_name() for API_name, pass_options in pass_pipeline.items(): name = pass_names.get(API_name, API_name) passes.append(Pass(name, **pass_options)) return passes return pass_pipeline def _API_name_to_pass_name(): return { "cancel_inverses": "remove-chained-self-inverse", "disentangle_cnot": "disentangle-CNOT", "disentangle_swap": "disentangle-SWAP", "merge_rotations": "merge-rotations", "ions_decomposition": "ions-decomposition", "to_ppr": "to-ppr", "commute_ppr": "commute-ppr", "merge_ppr_ppm": "merge-ppr-ppm", "decompose_non_clifford_ppr": "decompose-non-clifford-ppr", "decompose_clifford_ppr": "decompose-clifford-ppr", "ppm_compilation": "ppm-compilation", "t_layer_reduction": "t-layer-reduction", }