From 9904441d189b13b388768d8825d643091e31888a Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 10:59:02 -0500 Subject: [PATCH 01/15] [NFC] Restructure the passes module --- frontend/catalyst/passes.py | 534 --------------------- frontend/catalyst/passes/__init__.py | 53 ++ frontend/catalyst/passes/builtin_passes.py | 242 ++++++++++ frontend/catalyst/passes/pass_api.py | 295 ++++++++++++ 4 files changed, 590 insertions(+), 534 deletions(-) delete mode 100644 frontend/catalyst/passes.py create mode 100644 frontend/catalyst/passes/__init__.py create mode 100644 frontend/catalyst/passes/builtin_passes.py create mode 100644 frontend/catalyst/passes/pass_api.py diff --git a/frontend/catalyst/passes.py b/frontend/catalyst/passes.py deleted file mode 100644 index 498bba016a..0000000000 --- a/frontend/catalyst/passes.py +++ /dev/null @@ -1,534 +0,0 @@ -# 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. - -""" -This module contains Python decorators for enabling and configuring -individual Catalyst MLIR compiler passes. - -.. note:: - - Unlike PennyLane :doc:`circuit transformations `, - the QNode itself will not be changed or transformed by applying these - decorators. - - As a result, circuit inspection tools such as :func:`~.draw` will continue - to display the circuit as written in Python. - - Instead, these compiler passes are applied at the MLIR level, which occurs - outside of Python during compile time. To inspect the compiled MLIR from - Catalyst, use :func:`~.get_compilation_stage` with - ``stage="QuantumCompilationPass"``. - -""" - -import copy -import functools -from importlib.metadata import entry_points -from pathlib import Path -from typing import TypeAlias - -import pennylane as qml - -from catalyst.tracing.contexts import EvaluationContext - -PipelineDict: TypeAlias = dict[str, dict[str, str]] - - -class Pass: - """Class intended to hold options for passes""" - - def __init__(self, name, *options, **valued_options): - 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 - - def __repr__(self): - return ( - self.name - + " ".join(f"--{option}" for option in self.options) - + " ".join(f"--{option}={value}" for option, value in self.valued_options) - ) - - -class PassPlugin(Pass): - """Class intended to hold options for pass plugins""" - - 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) - - -def dictionary_to_list_of_passes(pass_pipeline: PipelineDict): - """Convert dictionary of passes into list of passes""" - - if pass_pipeline == None: - return [] - - if type(pass_pipeline) != dict: - return pass_pipeline - - 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 - - -## API ## -# pylint: disable=line-too-long -@functools.singledispatch -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 `_ 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 `. - - 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 gaurantee 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=None): - if not isinstance(qnode, qml.QNode): - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - clone = copy.copy(qnode) - clone.__name__ += "_transformed" - - @functools.wraps(clone) - def wrapper(*args, **kwargs): - if EvaluationContext.is_tracing(): - passes = kwargs.pop("pass_pipeline", []) - passes += dictionary_to_list_of_passes(pass_pipeline) - kwargs["pass_pipeline"] = passes - return clone(*args, **kwargs) - - return wrapper - - return _decorator - - -def cancel_inverses(qnode=None): - """ - Specify that the ``-removed-chained-self-inverse`` MLIR compiler pass - for cancelling two neighbouring self-inverse - gates should be applied to the decorated QNode during :func:`~.qjit` - compilation. - - The full list of supported gates are as follows: - - One-bit Gates: - :class:`qml.Hadamard `, - :class:`qml.PauliX `, - :class:`qml.PauliY `, - :class:`qml.PauliZ ` - - Two-bit Gates: - :class:`qml.CNOT `, - :class:`qml.CY `, - :class:`qml.CZ `, - :class:`qml.SWAP ` - - Three-bit Gates: - - :class:`qml.Toffoli ` - - .. note:: - - Unlike PennyLane :doc:`circuit transformations `, - the QNode itself will not be changed or transformed by applying these - decorators. - - As a result, circuit inspection tools such as :func:`~.draw` will continue - to display the circuit as written in Python. - - To instead view the optimized circuit, the MLIR must be viewed - after the ``"QuantumCompilationPass"`` stage via the - :func:`~.get_compilation_stage` function. - - Args: - fn (QNode): the QNode to apply the cancel inverses compiler pass to - - Returns: - ~.QNode: - - **Example** - - .. code-block:: python - - from catalyst.debug import get_compilation_stage - from catalyst.passes import cancel_inverses - - dev = qml.device("lightning.qubit", wires=1) - - @qjit(keep_intermediate=True) - @cancel_inverses - @qml.qnode(dev) - def circuit(x: float): - qml.RX(x, wires=0) - qml.Hadamard(wires=0) - qml.Hadamard(wires=0) - return qml.expval(qml.PauliZ(0)) - - >>> circuit(0.54) - Array(0.85770868, dtype=float64) - - Note that the QNode will be unchanged in Python, and will continue - to include self-inverse gates when inspected with Python (for example, - with :func:`~.draw`). - - To instead view the optimized circuit, the MLIR must be viewed - after the ``"QuantumCompilationPass"`` stage: - - >>> print(get_compilation_stage(circuit, stage="QuantumCompilationPass")) - module @circuit { - func.func public @jit_circuit(%arg0: tensor) -> tensor attributes {llvm.emit_c_interface} { - %0 = call @circuit(%arg0) : (tensor) -> tensor - return %0 : tensor - } - func.func private @circuit(%arg0: tensor) -> tensor attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage, qnode} { - quantum.device["catalyst/utils/../lib/librtd_lightning.dylib", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"] - %0 = quantum.alloc( 1) : !quantum.reg - %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit - %extracted = tensor.extract %arg0[] : tensor - %out_qubits = quantum.custom "RX"(%extracted) %1 : !quantum.bit - %2 = quantum.namedobs %out_qubits[ PauliZ] : !quantum.obs - %3 = quantum.expval %2 : f64 - %from_elements = tensor.from_elements %3 : tensor - %4 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit - quantum.dealloc %4 : !quantum.reg - quantum.device_release - return %from_elements : tensor - } - func.func @setup() { - quantum.init - return - } - func.func @teardown() { - quantum.finalize - return - } - } - - It can be seen that both Hadamards have been cancelled, and the measurement - directly follows the ``RX`` gate: - - .. code-block:: mlir - - %out_qubits = quantum.custom "RX"(%extracted) %1 : !quantum.bit - %2 = quantum.namedobs %out_qubits[ PauliZ] : !quantum.obs - %3 = quantum.expval %2 : f64 - """ - if not isinstance(qnode, qml.QNode): - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - clone = copy.copy(qnode) - clone.__name__ += "_cancel_inverses" - - @functools.wraps(clone) - def wrapper(*args, **kwargs): - pass_pipeline = kwargs.pop("pass_pipeline", []) - pass_pipeline.append(Pass("remove-chained-self-inverse")) - kwargs["pass_pipeline"] = pass_pipeline - return clone(*args, **kwargs) - - return wrapper - - -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 ` - 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 - - @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): - - if not isinstance(qnode, qml.QNode): - # Technically, this apply pass is general enough that it can apply to - # classical functions too. However, since we lack the current infrastructure - # to denote a function, let's limit it to qnodes - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - def qnode_call(*args, **kwargs): - pass_pipeline = kwargs.get("pass_pipeline", []) - pass_pipeline.append(Pass(pass_name, *flags, **valued_options)) - kwargs["pass_pipeline"] = pass_pipeline - return qnode(*args, **kwargs) - - return qnode_call - - return decorator - - -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 ` - 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 - - @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): - if not isinstance(qnode, qml.QNode): - # Technically, this apply pass is general enough that it can apply to - # classical functions too. However, since we lack the current infrastructure - # to denote a function, let's limit it to qnodes - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - def qnode_call(*args, **kwargs): - pass_pipeline = kwargs.get("pass_pipeline", []) - pass_pipeline.append(PassPlugin(path_to_plugin, pass_name, *flags, **valued_options)) - kwargs["pass_pipeline"] = pass_pipeline - return qnode(*args, **kwargs) - - return qnode_call - - return decorator - - -def merge_rotations(qnode=None): - """ - Specify that the ``-merge-rotations`` MLIR compiler pass - for merging roations (peephole) will be applied. - - The full list of supported gates are as follows: - - :class:`qml.RX `, - :class:`qml.CRX `, - :class:`qml.RY `, - :class:`qml.CRY `, - :class:`qml.RZ `, - :class:`qml.CRZ `, - :class:`qml.PhaseShift `, - :class:`qml.ControlledPhaseShift `, - :class:`qml.MultiRZ `. - - - .. note:: - - Unlike PennyLane :doc:`circuit transformations `, - the QNode itself will not be changed or transformed by applying these - decorators. - - As a result, circuit inspection tools such as :func:`~.draw` will continue - to display the circuit as written in Python. - - To instead view the optimized circuit, the MLIR must be viewed - after the ``"QuantumCompilationPass"`` stage via the - :func:`~.get_compilation_stage` function. - - Args: - fn (QNode): the QNode to apply the cancel inverses compiler pass to - - Returns: - ~.QNode: - - **Example** - - In this example the three :class:`qml.RX ` will be merged in a single - one with the sum of angles as parameter. - - .. code-block:: python - - from catalyst.debug import get_compilation_stage - from catalyst.passes import merge_rotations - - dev = qml.device("lightning.qubit", wires=1) - - @qjit(keep_intermediate=True) - @merge_rotations - @qml.qnode(dev) - def circuit(x: float): - qml.RX(x, wires=0) - qml.RX(0.1, wires=0) - qml.RX(x**2, wires=0) - return qml.expval(qml.PauliZ(0)) - - >>> circuit(0.54) - Array(0.5965506257017892, dtype=float64) - """ - if not isinstance(qnode, qml.QNode): - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - clone = copy.copy(qnode) - clone.__name__ += "_merge_rotations" - - @functools.wraps(clone) - def wrapper(*args, **kwargs): - pass_pipeline = kwargs.pop("pass_pipeline", []) - pass_pipeline.append(Pass("merge-rotations")) - kwargs["pass_pipeline"] = pass_pipeline - return clone(*args, **kwargs) - - return wrapper - - -def _API_name_to_pass_name(): - return { - "cancel_inverses": "remove-chained-self-inverse", - "merge_rotations": "merge-rotations", - "ions_decomposition": "ions-decomposition", - } - - -def ions_decomposition(qnode=None): # pragma: nocover - """Apply decomposition pass at the MLIR level""" - - if not isinstance(qnode, qml.QNode): - raise TypeError(f"A QNode is expected, got the classical function {qnode}") - - @functools.wraps(qnode) - def wrapper(*args, **kwargs): - pass_pipeline = kwargs.pop("pass_pipeline", []) - pass_pipeline.append(Pass("ions-decomposition")) - kwargs["pass_pipeline"] = pass_pipeline - return qnode(*args, **kwargs) - - return wrapper diff --git a/frontend/catalyst/passes/__init__.py b/frontend/catalyst/passes/__init__.py new file mode 100644 index 0000000000..20c59233d1 --- /dev/null +++ b/frontend/catalyst/passes/__init__.py @@ -0,0 +1,53 @@ +# 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. + +""" +This module provides access to Catalyst's compiler transformation infrastructure, including the +use of Python decorators to configure and schedule individual built-in compiler passes, as well +as load and run external MLIR passes from plugins. + +.. note:: + + Unlike PennyLane :doc:`circuit transformations `, + the QNode itself will not be changed or transformed by applying these + decorators. + + As a result, circuit inspection tools such as :func:`~.draw` will continue + to display the circuit as written in Python. + + Instead, these compiler passes are applied at the MLIR level, which occurs + outside of Python during compile time. To inspect the compiled MLIR from + Catalyst, use :func:`~.get_compilation_stage` with + ``stage="QuantumCompilationPass"``. + +""" + +from catalyst.passes.builtin_passes import cancel_inverses, merge_rotations +from catalyst.passes.pass_api import ( + Pass, + PassPlugin, + apply_pass, + apply_pass_plugin, + pipeline, +) + +__all__ = ( + "cancel_inverses", + "merge_rotations", + "Pass", + "PassPlugin", + "apply_pass", + "apply_pass_plugin", + "pipeline", +) diff --git a/frontend/catalyst/passes/builtin_passes.py b/frontend/catalyst/passes/builtin_passes.py new file mode 100644 index 0000000000..02b527c0b9 --- /dev/null +++ b/frontend/catalyst/passes/builtin_passes.py @@ -0,0 +1,242 @@ +# 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. + +"""This module exposes built-in Catalyst MLIR passes to the frontend.""" + +import copy +import functools + +import pennylane as qml + +from catalyst.passes.pass_api import Pass + + +## API ## +def cancel_inverses(qnode=None): + """ + Specify that the ``-removed-chained-self-inverse`` MLIR compiler pass + for cancelling two neighbouring self-inverse + gates should be applied to the decorated QNode during :func:`~.qjit` + compilation. + + The full list of supported gates are as follows: + + One-bit Gates: + :class:`qml.Hadamard `, + :class:`qml.PauliX `, + :class:`qml.PauliY `, + :class:`qml.PauliZ ` + + Two-bit Gates: + :class:`qml.CNOT `, + :class:`qml.CY `, + :class:`qml.CZ `, + :class:`qml.SWAP ` + + Three-bit Gates: + - :class:`qml.Toffoli ` + + .. note:: + + Unlike PennyLane :doc:`circuit transformations `, + the QNode itself will not be changed or transformed by applying these + decorators. + + As a result, circuit inspection tools such as :func:`~.draw` will continue + to display the circuit as written in Python. + + To instead view the optimized circuit, the MLIR must be viewed + after the ``"QuantumCompilationPass"`` stage via the + :func:`~.get_compilation_stage` function. + + Args: + fn (QNode): the QNode to apply the cancel inverses compiler pass to + + Returns: + ~.QNode: + + **Example** + + .. code-block:: python + + from catalyst.debug import get_compilation_stage + from catalyst.passes import cancel_inverses + + dev = qml.device("lightning.qubit", wires=1) + + @qjit(keep_intermediate=True) + @cancel_inverses + @qml.qnode(dev) + def circuit(x: float): + qml.RX(x, wires=0) + qml.Hadamard(wires=0) + qml.Hadamard(wires=0) + return qml.expval(qml.PauliZ(0)) + + >>> circuit(0.54) + Array(0.85770868, dtype=float64) + + Note that the QNode will be unchanged in Python, and will continue + to include self-inverse gates when inspected with Python (for example, + with :func:`~.draw`). + + To instead view the optimized circuit, the MLIR must be viewed + after the ``"QuantumCompilationPass"`` stage: + + >>> print(get_compilation_stage(circuit, stage="QuantumCompilationPass")) + module @circuit { + func.func public @jit_circuit(%arg0: tensor) -> tensor attributes {llvm.emit_c_interface} { + %0 = call @circuit(%arg0) : (tensor) -> tensor + return %0 : tensor + } + func.func private @circuit(%arg0: tensor) -> tensor attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage, qnode} { + quantum.device["catalyst/utils/../lib/librtd_lightning.dylib", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"] + %0 = quantum.alloc( 1) : !quantum.reg + %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit + %extracted = tensor.extract %arg0[] : tensor + %out_qubits = quantum.custom "RX"(%extracted) %1 : !quantum.bit + %2 = quantum.namedobs %out_qubits[ PauliZ] : !quantum.obs + %3 = quantum.expval %2 : f64 + %from_elements = tensor.from_elements %3 : tensor + %4 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit + quantum.dealloc %4 : !quantum.reg + quantum.device_release + return %from_elements : tensor + } + func.func @setup() { + quantum.init + return + } + func.func @teardown() { + quantum.finalize + return + } + } + + It can be seen that both Hadamards have been cancelled, and the measurement + directly follows the ``RX`` gate: + + .. code-block:: mlir + + %out_qubits = quantum.custom "RX"(%extracted) %1 : !quantum.bit + %2 = quantum.namedobs %out_qubits[ PauliZ] : !quantum.obs + %3 = quantum.expval %2 : f64 + """ + if not isinstance(qnode, qml.QNode): + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + clone = copy.copy(qnode) + clone.__name__ += "_cancel_inverses" + + @functools.wraps(clone) + def wrapper(*args, **kwargs): + pass_pipeline = kwargs.pop("pass_pipeline", []) + pass_pipeline.append(Pass("remove-chained-self-inverse")) + kwargs["pass_pipeline"] = pass_pipeline + return clone(*args, **kwargs) + + return wrapper + + +def merge_rotations(qnode=None): + """ + Specify that the ``-merge-rotations`` MLIR compiler pass + for merging roations (peephole) will be applied. + + The full list of supported gates are as follows: + + :class:`qml.RX `, + :class:`qml.CRX `, + :class:`qml.RY `, + :class:`qml.CRY `, + :class:`qml.RZ `, + :class:`qml.CRZ `, + :class:`qml.PhaseShift `, + :class:`qml.ControlledPhaseShift `, + :class:`qml.MultiRZ `. + + + .. note:: + + Unlike PennyLane :doc:`circuit transformations `, + the QNode itself will not be changed or transformed by applying these + decorators. + + As a result, circuit inspection tools such as :func:`~.draw` will continue + to display the circuit as written in Python. + + To instead view the optimized circuit, the MLIR must be viewed + after the ``"QuantumCompilationPass"`` stage via the + :func:`~.get_compilation_stage` function. + + Args: + fn (QNode): the QNode to apply the cancel inverses compiler pass to + + Returns: + ~.QNode: + + **Example** + + In this example the three :class:`qml.RX ` will be merged in a single + one with the sum of angles as parameter. + + .. code-block:: python + + from catalyst.debug import get_compilation_stage + from catalyst.passes import merge_rotations + + dev = qml.device("lightning.qubit", wires=1) + + @qjit(keep_intermediate=True) + @merge_rotations + @qml.qnode(dev) + def circuit(x: float): + qml.RX(x, wires=0) + qml.RX(0.1, wires=0) + qml.RX(x**2, wires=0) + return qml.expval(qml.PauliZ(0)) + + >>> circuit(0.54) + Array(0.5965506257017892, dtype=float64) + """ + if not isinstance(qnode, qml.QNode): + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + clone = copy.copy(qnode) + clone.__name__ += "_merge_rotations" + + @functools.wraps(clone) + def wrapper(*args, **kwargs): + pass_pipeline = kwargs.pop("pass_pipeline", []) + pass_pipeline.append(Pass("merge-rotations")) + kwargs["pass_pipeline"] = pass_pipeline + return clone(*args, **kwargs) + + return wrapper + + +def ions_decomposition(qnode=None): # pragma: nocover + """Apply decomposition pass at the MLIR level.""" + + if not isinstance(qnode, qml.QNode): + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + @functools.wraps(qnode) + def wrapper(*args, **kwargs): + pass_pipeline = kwargs.pop("pass_pipeline", []) + pass_pipeline.append(Pass("ions-decomposition")) + kwargs["pass_pipeline"] = pass_pipeline + return qnode(*args, **kwargs) + + return wrapper diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py new file mode 100644 index 0000000000..a50588ed46 --- /dev/null +++ b/frontend/catalyst/passes/pass_api.py @@ -0,0 +1,295 @@ +# 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 copy +import functools +from importlib.metadata import entry_points +from pathlib import Path +from typing import TypeAlias + +import pennylane as qml + +from catalyst.tracing.contexts import EvaluationContext + +PipelineDict: TypeAlias = dict[str, dict[str, str]] + + +## API ## +@functools.singledispatch +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 `_ + 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 `. + + 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=None): + if not isinstance(qnode, qml.QNode): + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + clone = copy.copy(qnode) + clone.__name__ += "_transformed" + + @functools.wraps(clone) + def wrapper(*args, **kwargs): + if EvaluationContext.is_tracing(): + passes = kwargs.pop("pass_pipeline", []) + passes += dictionary_to_list_of_passes(pass_pipeline) + kwargs["pass_pipeline"] = passes + return clone(*args, **kwargs) + + return wrapper + + return _decorator + + +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 ` 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 + + @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): + + if not isinstance(qnode, qml.QNode): + # Technically, this apply pass is general enough that it can apply to + # classical functions too. However, since we lack the current infrastructure + # to denote a function, let's limit it to qnodes + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + def qnode_call(*args, **kwargs): + pass_pipeline = kwargs.get("pass_pipeline", []) + pass_pipeline.append(Pass(pass_name, *flags, **valued_options)) + kwargs["pass_pipeline"] = pass_pipeline + return qnode(*args, **kwargs) + + return qnode_call + + return decorator + + +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 ` 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 + + @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): + if not isinstance(qnode, qml.QNode): + # Technically, this apply pass is general enough that it can apply to + # classical functions too. However, since we lack the current infrastructure + # to denote a function, let's limit it to qnodes + raise TypeError(f"A QNode is expected, got the classical function {qnode}") + + def qnode_call(*args, **kwargs): + pass_pipeline = kwargs.get("pass_pipeline", []) + pass_pipeline.append(PassPlugin(path_to_plugin, pass_name, *flags, **valued_options)) + kwargs["pass_pipeline"] = pass_pipeline + return qnode(*args, **kwargs) + + return qnode_call + + return decorator + + +class Pass: + """Class intended to hold options for passes.""" + + def __init__(self, name, *options, **valued_options): + 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 + + def __repr__(self): + return ( + self.name + + " ".join(f"--{option}" for option in self.options) + + " ".join(f"--{option}={value}" for option, value in self.valued_options) + ) + + +class PassPlugin(Pass): + """Class intended to hold options for pass plugins.""" + + 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): + """Convert dictionary of passes into list of passes.""" + + if pass_pipeline == None: + return [] + + if type(pass_pipeline) != dict: + return pass_pipeline + + 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 + + +def _API_name_to_pass_name(): + return { + "cancel_inverses": "remove-chained-self-inverse", + "merge_rotations": "merge-rotations", + "ions_decomposition": "ions-decomposition", + } From 5759186f9cc1d885592410ca54b2acc13ef19b40 Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 10:59:49 -0500 Subject: [PATCH 02/15] Don't export functions from the passes module --- frontend/catalyst/__init__.py | 9 +++------ frontend/catalyst/passes/pass_api.py | 4 ++-- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/frontend/catalyst/__init__.py b/frontend/catalyst/__init__.py index e0a39e175d..250782195d 100644 --- a/frontend/catalyst/__init__.py +++ b/frontend/catalyst/__init__.py @@ -81,7 +81,7 @@ "mlir_quantum._mlir_libs._quantumDialects.mitigation" ) -from catalyst import debug, logging +from catalyst import debug, logging, passes from catalyst.api_extensions import * from catalyst.api_extensions import __all__ as _api_extension_list from catalyst.autograph import * @@ -89,7 +89,7 @@ from catalyst.compiler import CompileOptions from catalyst.debug.assertion import debug_assert from catalyst.jit import QJIT, qjit -from catalyst.passes import Pass, PassPlugin, apply_pass, apply_pass_plugin, pipeline +from catalyst.passes import pipeline from catalyst.utils.exceptions import ( AutoGraphError, CompileError, @@ -187,11 +187,8 @@ "debug_assert", "CompileOptions", "debug", - "apply_pass", - "apply_pass_plugin", + "passes", "pipeline", - "Pass", - "PassPlugin", *_api_extension_list, *_autograph_functions, ) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index a50588ed46..4e9aab3f55 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -151,7 +151,7 @@ def apply_pass(pass_name: str, *flags, **valued_options): .. code-block:: python - @apply_pass("merge-rotations") + @passes.apply_pass("merge-rotations") @qml.qnode(qml.device("lightning.qubit", wires=1)) def qnode(): return qml.state() @@ -198,7 +198,7 @@ def apply_pass_plugin(path_to_plugin: str | Path, pass_name: str, *flags, **valu from standalone import getStandalonePluginAbsolutePath - @apply_pass_plugin(getStandalonePluginAbsolutePath(), "standalone-switch-bar-foo") + @passes.apply_pass_plugin(getStandalonePluginAbsolutePath(), "standalone-switch-bar-foo") @qml.qnode(qml.device("lightning.qubit", wires=1)) def qnode(): return qml.state() From ee3b79a2e2742e0959e7865de0cb5d15ba9cc331 Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 11:07:35 -0500 Subject: [PATCH 03/15] Remove default arg value for non-optional argument --- frontend/catalyst/passes/builtin_passes.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/passes/builtin_passes.py b/frontend/catalyst/passes/builtin_passes.py index 02b527c0b9..d929f498f0 100644 --- a/frontend/catalyst/passes/builtin_passes.py +++ b/frontend/catalyst/passes/builtin_passes.py @@ -23,7 +23,7 @@ ## API ## -def cancel_inverses(qnode=None): +def cancel_inverses(qnode): """ Specify that the ``-removed-chained-self-inverse`` MLIR compiler pass for cancelling two neighbouring self-inverse @@ -149,7 +149,7 @@ def wrapper(*args, **kwargs): return wrapper -def merge_rotations(qnode=None): +def merge_rotations(qnode): """ Specify that the ``-merge-rotations`` MLIR compiler pass for merging roations (peephole) will be applied. @@ -226,7 +226,7 @@ def wrapper(*args, **kwargs): return wrapper -def ions_decomposition(qnode=None): # pragma: nocover +def ions_decomposition(qnode): # pragma: nocover """Apply decomposition pass at the MLIR level.""" if not isinstance(qnode, qml.QNode): From 76c880b99feeed568d3efc304547591c725ea4b8 Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 11:23:46 -0500 Subject: [PATCH 04/15] Do not export pipelines in two namespaces --- frontend/catalyst/__init__.py | 2 +- frontend/catalyst/passes/__init__.py | 9 +-------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/frontend/catalyst/__init__.py b/frontend/catalyst/__init__.py index 250782195d..67cfb788a1 100644 --- a/frontend/catalyst/__init__.py +++ b/frontend/catalyst/__init__.py @@ -89,7 +89,7 @@ from catalyst.compiler import CompileOptions from catalyst.debug.assertion import debug_assert from catalyst.jit import QJIT, qjit -from catalyst.passes import pipeline +from catalyst.passes.pass_api import pipeline from catalyst.utils.exceptions import ( AutoGraphError, CompileError, diff --git a/frontend/catalyst/passes/__init__.py b/frontend/catalyst/passes/__init__.py index 20c59233d1..dd6217736a 100644 --- a/frontend/catalyst/passes/__init__.py +++ b/frontend/catalyst/passes/__init__.py @@ -34,13 +34,7 @@ """ from catalyst.passes.builtin_passes import cancel_inverses, merge_rotations -from catalyst.passes.pass_api import ( - Pass, - PassPlugin, - apply_pass, - apply_pass_plugin, - pipeline, -) +from catalyst.passes.pass_api import Pass, PassPlugin, apply_pass, apply_pass_plugin __all__ = ( "cancel_inverses", @@ -49,5 +43,4 @@ "PassPlugin", "apply_pass", "apply_pass_plugin", - "pipeline", ) From 72ce096392a61a15744bd452a327c6813a1a40ed Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 11:24:01 -0500 Subject: [PATCH 05/15] Fix import of util function --- frontend/catalyst/qfunc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/qfunc.py b/frontend/catalyst/qfunc.py index de13bd5443..e928197c20 100644 --- a/frontend/catalyst/qfunc.py +++ b/frontend/catalyst/qfunc.py @@ -49,7 +49,7 @@ from catalyst.jax_primitives import quantum_kernel_p from catalyst.jax_tracer import Function, trace_quantum_function from catalyst.logging import debug_logger -from catalyst.passes import dictionary_to_list_of_passes +from catalyst.passes.pass_api import dictionary_to_list_of_passes from catalyst.tracing.contexts import EvaluationContext from catalyst.tracing.type_signatures import filter_static_args from catalyst.utils.exceptions import CompileError From 81edaf74353747e264c0a6d645cf443cc3e5a17f Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 11:44:44 -0500 Subject: [PATCH 06/15] Fix tests --- frontend/test/pytest/test_mlir_plugin_interface.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py index fe20d97bab..91208a5d0d 100644 --- a/frontend/test/pytest/test_mlir_plugin_interface.py +++ b/frontend/test/pytest/test_mlir_plugin_interface.py @@ -25,7 +25,11 @@ def test_path_does_not_exists(): """Test what happens when a pass_plugin is given an path that does not exist""" with pytest.raises(FileNotFoundError, match="does not exist"): - catalyst.apply_pass_plugin("this-path-does-not-exist", "this-pass-also-doesnt-exists") + catalyst.passes.apply_pass_plugin( + "this-path-does-not-exist", "this-pass-also-doesnt-exists" + ) with pytest.raises(FileNotFoundError, match="does not exist"): - catalyst.apply_pass_plugin(Path("this-path-does-not-exist"), "this-pass-also-doesnt-exists") + catalyst.passes.apply_pass_plugin( + Path("this-path-does-not-exist"), "this-pass-also-doesnt-exists" + ) From b2a4f86dba3abc22e85b2cc122c0b4b3ec4ddc2c Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 12:55:00 -0500 Subject: [PATCH 07/15] Remove unused single_dispatch decorator. --- frontend/catalyst/passes/pass_api.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index 4e9aab3f55..a6e4743384 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -26,7 +26,6 @@ ## API ## -@functools.singledispatch def pipeline(pass_pipeline: PipelineDict): """Configures the Catalyst MLIR pass pipeline for quantum circuit transformations for a QNode within a qjit-compiled program. From be0e2f13715e9e6cf8acd660f934ec3ff47c3813 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 12:55:40 -0500 Subject: [PATCH 08/15] Make qnodes decorated with apply_pass_plugin and apply_pass AOT jittable --- frontend/catalyst/passes/pass_api.py | 4 ++- .../test/pytest/test_mlir_plugin_interface.py | 33 +++++++++++++++++++ .../Quantum/Transforms/DisentangleSWAP.cpp | 2 +- 3 files changed, 37 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index a6e4743384..39c86d46b9 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -168,6 +168,7 @@ def decorator(qnode): # to denote a function, let's limit it to qnodes raise TypeError(f"A QNode is expected, got the classical function {qnode}") + @functools.wraps(qnode) def qnode_call(*args, **kwargs): pass_pipeline = kwargs.get("pass_pipeline", []) pass_pipeline.append(Pass(pass_name, *flags, **valued_options)) @@ -220,6 +221,7 @@ def decorator(qnode): # to denote a function, let's limit it to qnodes raise TypeError(f"A QNode is expected, got the classical function {qnode}") + @functools.wraps(qnode) def qnode_call(*args, **kwargs): pass_pipeline = kwargs.get("pass_pipeline", []) pass_pipeline.append(PassPlugin(path_to_plugin, pass_name, *flags, **valued_options)) @@ -234,7 +236,7 @@ def qnode_call(*args, **kwargs): class Pass: """Class intended to hold options for passes.""" - def __init__(self, name, *options, **valued_options): + def __init__(self, name: str, *options: list[str], **valued_options: dict[str, str]): self.options = options self.valued_options = valued_options if "." in name: diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py index 91208a5d0d..7d94e59927 100644 --- a/frontend/test/pytest/test_mlir_plugin_interface.py +++ b/frontend/test/pytest/test_mlir_plugin_interface.py @@ -15,7 +15,9 @@ """Testing interface around main plugin functionality""" from pathlib import Path +from tempfile import NamedTemporaryFile +import pennylane as qml import pytest import catalyst @@ -33,3 +35,34 @@ def test_path_does_not_exists(): catalyst.passes.apply_pass_plugin( Path("this-path-does-not-exist"), "this-pass-also-doesnt-exists" ) + + +def test_pass_can_aot_compile(): + """Can we AOT compile when using apply_pass?""" + + @qml.qjit(target="mlir") + @catalyst.passes.apply_pass("some-pass") + @qml.qnode(qml.device("null.qubit", wires=1)) + def example(): + return qml.state() + + assert example.mlir + + +@pytest.mark.skip() +def test_pass_plugin_can_aot_compile(): + """Can we AOT compile when using apply_pass_plugin? + + We can't properly test this because tmp needs to be a valid MLIR plugin. + And therefore can only be tested when a valid MLIR plugin exists in the path. + """ + + with NamedTemporaryFile() as tmp: + + @qml.qjit(target="mlir") + @catalyst.passes.apply_pass_plugin(Path(tmp.name), "some-pass") + @qml.qnode(qml.device("null.qubit", wires=1)) + def example(): + return qml.state() + + assert example.mlir diff --git a/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp b/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp index 22f702c8dd..7b0880b548 100644 --- a/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp +++ b/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp @@ -498,7 +498,7 @@ struct DisentangleSWAPPass : public impl::DisentangleSWAPPassBasegetRegion(0).front().getOperations()) { if (auto func = dyn_cast(nestedOp)) { disentangleSWAPs(func); - } + } } } }; From b601f30e9e0a4427e1b9b745f844f3148c013da3 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 14:22:53 -0500 Subject: [PATCH 09/15] Add support for options. --- frontend/catalyst/jax_primitives_utils.py | 3 ++- frontend/catalyst/passes/pass_api.py | 9 +++++++-- frontend/test/lit/test_mlir_plugin.py | 17 +++++++++++++++++ .../test/pytest/test_mlir_plugin_interface.py | 18 ++++++++++++++++++ 4 files changed, 44 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/jax_primitives_utils.py b/frontend/catalyst/jax_primitives_utils.py index 696392ed7f..778025e1cc 100644 --- a/frontend/catalyst/jax_primitives_utils.py +++ b/frontend/catalyst/jax_primitives_utils.py @@ -274,8 +274,9 @@ 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() apply_registered_pass_op = ApplyRegisteredPassOp( - result=transform_mod_type, target=target, pass_name=_pass.name + result=transform_mod_type, target=target, pass_name=_pass.name, options=options ) target = apply_registered_pass_op.result transform_yield_op = YieldOp(operands_=[]) # pylint: disable=unused-variable diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index 39c86d46b9..7e2ff7c3e4 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -250,11 +250,16 @@ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, s self.name = name + def get_options(self): + retval = " ".join(f"{str(option)}" for option in self.options) + retval2 = " ".join(f"{str(key)}={str(value)}" for key, value in self.valued_options.items()) + return " ".join([retval, retval2]) + def __repr__(self): return ( self.name - + " ".join(f"--{option}" for option in self.options) - + " ".join(f"--{option}={value}" for option, value in self.valued_options) + + " ".join(f"--{str(option)}=true" for option in self.options) + + " ".join([f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()]) ) diff --git a/frontend/test/lit/test_mlir_plugin.py b/frontend/test/lit/test_mlir_plugin.py index 44a0d4f488..26b8daa2f8 100644 --- a/frontend/test/lit/test_mlir_plugin.py +++ b/frontend/test/lit/test_mlir_plugin.py @@ -63,6 +63,8 @@ def module(): import platform from pathlib import Path +import pennylane as qml +import catalyst from catalyst.compiler import CompileOptions, Compiler from catalyst.utils.filesystem import WorkspaceManager from catalyst.utils.runtime_environment import get_bin_path @@ -96,3 +98,18 @@ def module(): custom_compiler = Compiler(options) _, mlir_string = custom_compiler.run_from_ir(mlir_module, "test", workspace) print(mlir_string) + + +def test_pass_options(): + """Is the option in the generated MLIR?""" + + @qml.qjit(target="mlir") + # CHECK: options = "an-option maxValue=1" + @catalyst.passes.apply_pass("some-pass", "an-option", maxValue=1) + @qml.qnode(qml.device("null.qubit", wires=1)) + def example(): + return qml.state() + + print(example.mlir) + +test_pass_options() diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py index 7d94e59927..20bee0f63c 100644 --- a/frontend/test/pytest/test_mlir_plugin_interface.py +++ b/frontend/test/pytest/test_mlir_plugin_interface.py @@ -66,3 +66,21 @@ def example(): return qml.state() assert example.mlir + +def test_get_options(): + """ + ApplyRegisteredPassOp expects options to be a single StringAttr + which follows the same format as the one used with mlir-opt. + + https://mlir.llvm.org/docs/Dialects/Transform/#transformapply_registered_pass-transformapplyregisteredpassop + + Options passed to a pass are specified via the syntax {option1=value1 option2=value2 ...}, + i.e., use space-separated key=value pairs for each option. + + https://mlir.llvm.org/docs/Tutorials/MlirOpt/#running-a-pass-with-options + + However, experimentally we found that single-options also work without values. + """ + assert catalyst.passes.Pass("example-pass", "single-option").get_options() == "single-option" + assert catalyst.passes.Pass("example-pass", "an-option", "bn-option").get_options() == "an-option bn-option" + assert catalyst.passes.Pass("example-pass", option=True).get_options() == "option=True" From 65ee6ff3d5677bef8af585057dcd4801c82fa026 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 14:44:21 -0500 Subject: [PATCH 10/15] Documentation --- frontend/catalyst/passes/pass_api.py | 30 +++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index 7e2ff7c3e4..c27c278161 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -234,7 +234,12 @@ def qnode_call(*args, **kwargs): class Pass: - """Class intended to hold options for passes.""" + """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. + """ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, str]): self.options = options @@ -251,6 +256,19 @@ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, s self.name = name def get_options(self): + """ + ApplyRegisteredPassOp expects options to be a single StringAttr + which follows the same format as the one used with mlir-opt. + + https://mlir.llvm.org/docs/Dialects/Transform/#transformapply_registered_pass-transformapplyregisteredpassop + + Options passed to a pass are specified via the syntax {option1=value1 option2=value2 ...}, + i.e., use space-separated key=value pairs for each option. + + https://mlir.llvm.org/docs/Tutorials/MlirOpt/#running-a-pass-with-options + + However, experimentally we found that single-options also work without values. + """ retval = " ".join(f"{str(option)}" for option in self.options) retval2 = " ".join(f"{str(key)}={str(value)}" for key, value in self.valued_options.items()) return " ".join([retval, retval2]) @@ -258,13 +276,19 @@ def get_options(self): def __repr__(self): return ( self.name - + " ".join(f"--{str(option)}=true" for option in self.options) + + " ".join(f"--{str(option)}" for option in self.options) + " ".join([f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()]) ) class PassPlugin(Pass): - """Class intended to hold options for pass plugins.""" + """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 + """ def __init__( self, path: Path, name: str, *options: list[str], **valued_options: dict[str, str] From f26e3c209a8efce45583f12401a1fcfe794179d7 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 15:35:22 -0500 Subject: [PATCH 11/15] fix --- frontend/catalyst/passes/pass_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index c27c278161..4d797088c5 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -271,7 +271,7 @@ def get_options(self): """ retval = " ".join(f"{str(option)}" for option in self.options) retval2 = " ".join(f"{str(key)}={str(value)}" for key, value in self.valued_options.items()) - return " ".join([retval, retval2]) + return " ".join([retval, retval2]).strip() def __repr__(self): return ( From 71500bab827d3b1a4780f8ec17a6e85e87206353 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 16:25:37 -0500 Subject: [PATCH 12/15] Documentation --- frontend/catalyst/passes/pass_api.py | 32 ++++++++++++++++++- frontend/test/lit/test_mlir_plugin.py | 2 ++ .../test/pytest/test_mlir_plugin_interface.py | 6 +++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index 4d797088c5..09784ce7be 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -239,6 +239,20 @@ class Pass: :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]): @@ -277,7 +291,9 @@ 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()]) + + " ".join( + [f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()] + ) ) @@ -288,6 +304,20 @@ class PassPlugin(Pass): 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__( diff --git a/frontend/test/lit/test_mlir_plugin.py b/frontend/test/lit/test_mlir_plugin.py index 26b8daa2f8..ccc2e04b6d 100644 --- a/frontend/test/lit/test_mlir_plugin.py +++ b/frontend/test/lit/test_mlir_plugin.py @@ -64,6 +64,7 @@ def module(): from pathlib import Path import pennylane as qml + import catalyst from catalyst.compiler import CompileOptions, Compiler from catalyst.utils.filesystem import WorkspaceManager @@ -112,4 +113,5 @@ def example(): print(example.mlir) + test_pass_options() diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py index 20bee0f63c..d476fc3032 100644 --- a/frontend/test/pytest/test_mlir_plugin_interface.py +++ b/frontend/test/pytest/test_mlir_plugin_interface.py @@ -67,6 +67,7 @@ def example(): assert example.mlir + def test_get_options(): """ ApplyRegisteredPassOp expects options to be a single StringAttr @@ -82,5 +83,8 @@ def test_get_options(): However, experimentally we found that single-options also work without values. """ assert catalyst.passes.Pass("example-pass", "single-option").get_options() == "single-option" - assert catalyst.passes.Pass("example-pass", "an-option", "bn-option").get_options() == "an-option bn-option" + assert ( + catalyst.passes.Pass("example-pass", "an-option", "bn-option").get_options() + == "an-option bn-option" + ) assert catalyst.passes.Pass("example-pass", option=True).get_options() == "option=True" From 0f876ab81e9d9b716758c6aa6df638f73ac035a1 Mon Sep 17 00:00:00 2001 From: David Ittah Date: Thu, 9 Jan 2025 16:00:43 -0500 Subject: [PATCH 13/15] Export ions_decomposition in the passes module --- frontend/catalyst/passes/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/passes/__init__.py b/frontend/catalyst/passes/__init__.py index dd6217736a..6d3a705593 100644 --- a/frontend/catalyst/passes/__init__.py +++ b/frontend/catalyst/passes/__init__.py @@ -33,7 +33,11 @@ """ -from catalyst.passes.builtin_passes import cancel_inverses, merge_rotations +from catalyst.passes.builtin_passes import ( + cancel_inverses, + ions_decomposition, + merge_rotations, +) from catalyst.passes.pass_api import Pass, PassPlugin, apply_pass, apply_pass_plugin __all__ = ( From 95ab10cd73f0dc4f5489d30e349a1d3bb2dbc1cd Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 16:43:14 -0500 Subject: [PATCH 14/15] String description at the top --- frontend/catalyst/passes/pass_api.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py index 09784ce7be..cab5d5750f 100644 --- a/frontend/catalyst/passes/pass_api.py +++ b/frontend/catalyst/passes/pass_api.py @@ -271,6 +271,8 @@ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, s def get_options(self): """ + Stringify options according to what mlir-opt expects. + ApplyRegisteredPassOp expects options to be a single StringAttr which follows the same format as the one used with mlir-opt. @@ -281,7 +283,7 @@ def get_options(self): https://mlir.llvm.org/docs/Tutorials/MlirOpt/#running-a-pass-with-options - However, experimentally we found that single-options also work without values. + Experimentally we found that single-options also work without values. """ retval = " ".join(f"{str(option)}" for option in self.options) retval2 = " ".join(f"{str(key)}={str(value)}" for key, value in self.valued_options.items()) From 02cdfd801677544645265a997fe4237e23816318 Mon Sep 17 00:00:00 2001 From: Erick Ochoa Lopez Date: Thu, 9 Jan 2025 16:44:38 -0500 Subject: [PATCH 15/15] Add top level description --- frontend/test/pytest/test_mlir_plugin_interface.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py index d476fc3032..115c65cd1f 100644 --- a/frontend/test/pytest/test_mlir_plugin_interface.py +++ b/frontend/test/pytest/test_mlir_plugin_interface.py @@ -70,6 +70,8 @@ def example(): def test_get_options(): """ + Test get_options from Pass + ApplyRegisteredPassOp expects options to be a single StringAttr which follows the same format as the one used with mlir-opt.