diff --git a/frontend/catalyst/device/python_decompositions.py b/frontend/catalyst/device/python_decompositions.py index b6ba5d71fe..95f233e666 100644 --- a/frontend/catalyst/device/python_decompositions.py +++ b/frontend/catalyst/device/python_decompositions.py @@ -14,26 +14,6 @@ """ This module provides infrastructure for compile-time lowering of decomposition rules via python. - -Python decomposition wrappers should adhere to the following specifications: - The wrapper: - - Is named `{op name}_decomposition_wrapper`. - - Has a signature identical to the named parameters of the associated PL operator; dynamic - arguments may be unused, but should still be included for compatibility. - - Is able to AOT lower the decomposition rule to MLIR without invoking the compiler, e.g. - using `target="mlir"`, AOT compilation and `QJIT.mlir_module`. - See existing examples for further information. - - Returns a string representation of an MLIR module, containing a FuncOp which represents - the instantiated decomposition rule. - - The FuncOp decomposition rule in the returned string: - - Is named `{op name}_decomp_rule`. - - Is an MLIR representation of the PennyLane decomposition rule associated with the - specified operator. - - Is instantiated with the static data provided, and all other data remains dynamic. - - Is compatible with the `decompose-lowering` pass, i.e. can be mapped to the MLIR operation - it decomposes and inlined. - - Is self-contained, and does not contain any device initialization, setup/teardown etc. """ # pylint: disable=protected-access,bare-except @@ -42,25 +22,134 @@ import jax.numpy as jnp import pennylane as qp +from jax._src.lib.mlir import ir +from jaxlib.mlir.dialects.builtin import ModuleOp + +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval + +_MLIR_DTYPES = { + "i1": jnp.bool_, + "i8": jnp.int8, + "i16": jnp.int16, + "i32": jnp.int32, + "i64": jnp.int64, + "f16": jnp.float16, + "f32": jnp.float32, + "f64": jnp.float64, + "complex": jnp.complex64, + "complex": jnp.complex128, +} + + +def get_dummy_values_for_container(container): + """Given a container of python types, replace the types with corresponding dummy values.""" + dummy_args = [] + for dtype in container: + if isinstance(dtype, str): + if dtype in _MLIR_DTYPES: + count = 1 + dtype = _MLIR_DTYPES[dtype] + elif dtype.startswith("tensor"): + # tensor<{number}x{type}> + dtype = dtype.removeprefix("tensor<") + dtype = dtype.remove_suffice(">") + count, dtype = dtype.split("x") + else: + raise ValueError(f"Unknown dtype {dtype}.") + else: + count = 1 + dtype = jnp.dtype(dtype) + + dummy_args.append(jnp.zeros((count,), dtype=dtype)) + + return tuple(dummy_args) + + +def get_graph_op_id(op: qp.decomposition.CompressedResourceOp | qp.core.Operator2): + """ + Return the graph operator id for the operator2 instance `op`. + + The FuncOp decomposition rules in the returned string satisfy the following requirements: + - Are named `{rule name}_{op graph ID}`. + - Are MLIR representations of the PennyLane decomposition rules associated with the + specified operator. + - Are instantiated with the static data provided, and all other data remains dynamic. + - Are self-contained, and do not contain any device initialization, setup/teardown etc. + - Are compatible with the `decompose-lowering` and `graph-decomposition` passes, meaning + the following: + - Their `target_gate` attribute is set to the provided graph operator ID + - They have a resources attribute containing an operations attribute which maps graph + operator IDs to counts of their occurrences in the rule. + - Their arguments are mappable to the operator they decompose via `decompose-lowering`. + + Note that this function should not be updated without updating the corresponding method on the + DecomposableGate interface in mlir/lib/quantum/IR/QuantumInterfaces.cpp. + """ + if isinstance(op, qp.decomposition.CompressedResourceOp): + # NOTE: handling this the old-fashioned way, remove once Operator2 migration is complete + op_type = op.op_type + else: + op_type = op + + if issubclass(op_type, qp.core.operator.Operator): + name = op_type.__name__ + num_params = op_type.num_params + num_wires = str(op_type.num_wires) if op_type.num_wires else "0" + return name + "[" + ",".join(["f64"] * num_params) + "][" + num_wires + "]{}" + elif isinstance(op_type, qp.core.operator.Operator2): + # TODO: use real getters here + name = op_type.__name__ + dynamic_shape = op_type.getDynamicShape() + wire_lens = op_type.getWireLens() + static_data = op_type.getStaticData() + extra_data = op_type.uid + return ( + name + + ("[" + dynamic_shape + "]") + + ("[" + wire_lens + "]") + + ("{" + static_data + "}") + + ("[" + extra_data + "]") + ) + else: + raise ValueError( + "Only AbstractOperator and CompressedResourceOp types are supported for generating a " + f"graph ID, got {op} of type {type(op)}" + ) -def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: - """Generic decomposition wrapper.""" +def python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp: + """Python decomposition rule lowering.""" + # TODO update docstring device = qp.device("null.qubit", wires=sum(wire_lens)) wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) + decomp_rules = list(qp.decomposition.list_decomps(op_name)) + + # map rules to resource resources, in a more generic format + + name_to_resources = {} + for rule in decomp_rules: + # TODO: not all PL ops have been migrated to the operator 2 format expected by mlir graph + # decomp This means some rules will fail the python callback compilation. When migration is + # complete, remove the try-except. + try: + name_to_resources[rule.name] = { + get_graph_op_id(op): count + for op, count in rule.compute_resources(**static_data).gate_counts.items() + } + except: # pylint: disable=bare-except + decomp_rules.remove(rule) + def rule_to_subroutine(rule): def decomp_rule(*params, wires): rule._impl(*params, *wires, **static_data) - # TODO remove this once we have unified lowering, we should be able to set target_gate and - # stop relying on function names - decomp_rule.__name__ = op_id + "_" + rule.name + # keep the frontend name for readability, append target op_id for symbol uniqueness + decomp_rule.__name__ = rule._impl.__name__ + "_" + op_id return qp.capture.subroutine(decomp_rule) - # let this fail with the standard error message if the op is not found - subroutines = [rule_to_subroutine(rule) for rule in qp.decomposition.list_decomps(op_name)] + subroutines = [rule_to_subroutine(rule) for rule in decomp_rules] # TODO: not all PL ops have been migrated to the operator 2 format expected by mlir graph decomp # This means some rules will fail the python callback compilation. @@ -74,12 +163,33 @@ def decomp_rule(*params, wires): @qp.qnode(device=device) def circuit(): for subroutine in subroutines: - # TODO: I know this is dynamic, but we should probably have a better way of handling - # this than hard-coded dummy values. Revisit this when unifying the decomp-rule - # lowering pipeline - subroutine(*[0.5 for _ in dynamic_shape], wires=wires) + subroutine(*get_dummy_values_for_container(dynamic_shape), wires=wires) - return str(circuit.mlir_module) + module = circuit.mlir_module + + def update_funcop_attributes(op): + """Update the decomposition rule attributes if op is a decomposition rule. + + For use with module.walk + + This function updates the following attributes: + - Adds the `target_gate` attribute. + - Adds the `resources` attribute. + """ + if op.name == "func.func": + rule_name = ir.StringAttr(op.attributes["sym_name"]).value.removesuffix("_" + op_id) + if rule_name in name_to_resources: + op.attributes["resources"] = get_mlir_attribute_from_pyval( + {"operations": name_to_resources[rule_name]} + ) + op.attributes["target_gate"] = ir.StringAttr.get(op_id) + + return ir.WalkResult.ADVANCE + + with module.context: + module.operation.walk(update_funcop_attributes) + + return module except: warnings.warn( f"Python decomposition rule compilation failed for operator " @@ -88,3 +198,8 @@ def circuit(): UserWarning, ) return "builtin.module{}" + + +def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: + """Generic decomposition wrapper.""" + return str(python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data)) diff --git a/frontend/catalyst/utils/precompile_decomposition_rules.py b/frontend/catalyst/utils/precompile_decomposition_rules.py index 2ca103944a..43cfd52901 100644 --- a/frontend/catalyst/utils/precompile_decomposition_rules.py +++ b/frontend/catalyst/utils/precompile_decomposition_rules.py @@ -14,22 +14,17 @@ """Utilities for AOT compiling PennyLane's decomposition rules to MLIR Bytecode.""" -import warnings from pathlib import Path -import jax import pennylane as qp from jax._src.lib.mlir import ir -from pennylane.operation import Operator +from pennylane.operation import Operator, Operator2 from catalyst.compiler import _quantum_opt -from catalyst.jax_primitives import decomposition_rule -from catalyst.utils.exceptions import CompileError +from catalyst.device.python_decompositions import get_graph_op_id, python_decomposition from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH # TODO: Uncomment dynamic size wires ops once they are supported -# FIXME: Use the Gate class instead of this list of compiler ops -# https://github.com/PennyLaneAI/pennylane/pull/8767 COMPILER_OPS_FOR_DECOMPOSITION = { qp.CNOT, qp.ControlledPhaseShift, @@ -96,113 +91,52 @@ def get_abstract_args(op_class: type[Operator]) -> list[type]: return [float for _ in range(op_class.num_params)] -def get_func_from_circuit(module) -> str | None: +def get_rules_from_module(module) -> str: """ - Get the string representation of `rule_wrapper` from module, if it exists. + Parse and modify decomposition rules from a ModuleOp. Args: module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted Returns: - str: string representation of FuncOp named `rule_wrapper` from module - None: if no such FuncOp can be found + str: The string representation of any decomposition rules from `module`, pre-pending the + `__builtin_` prefix to their names. """ - decomp_func_op = None + funcOps = [] def find_condition(op): - nonlocal decomp_func_op if op.name == "func.func": - if ir.StringAttr(op.attributes["sym_name"]).value == "rule_wrapper": - decomp_func_op = op - return ir.WalkResult.INTERRUPT + if "target_gate" in op.attributes: + op.attributes["sym_name"] = ir.StringAttr.get( + "__builtin_" + op.attributes["sym_name"].value.strip('"') + ) + funcOps.append(op) + return ir.WalkResult.SKIP return ir.WalkResult.ADVANCE module.operation.walk(find_condition) - return str(decomp_func_op) + "\n" if decomp_func_op else None - - -def compile_rule( - op_class, - abstract_args, - op_num_wires, - rule, - dev, -) -> str | None: - """ - Get the string representation of a compiled rule from a python decomposition rule, if possible. - - NOTE: rules with string params are not currently supported. - - Args: - op_class: A PennyLane class subclassing Operation - op_num_wires: the number of wires used by op_class - rule (DecompositionRule): the decomposition rule to be compiled - dev (Device): a device for qjit - - Returns: - str: string representation of the mlir of the decomposition rule. - """ - qp.decomposition.enable_graph() - - # WARNING: do not rename this function, we use it to extract the rule from the compiled - # circuit - @decomposition_rule(is_qreg=True, op_type=op_class.__name__) - def rule_wrapper(*args, wires, **_): - return rule(*args, wires=wires, **_) - - @qp.qjit(capture=True, target="mlir") - @qp.qnode(dev) - def circuit(): - rule_wrapper(*abstract_args, wires=jax.core.ShapedArray((op_num_wires,), int)) - return qp.probs() - - return get_func_from_circuit(circuit.mlir_module) - - -def compile_op_decomp_rules( - op_class: type[Operator], -) -> dict[str, str | None]: - """ - Compile all decomposition rules for op_class. - - Note: the modules include the full circuit IR. - - Args: - op_class (type[Operator]): the op class to compile decomposition rules for. - - Returns: - dict[str, str | None]: decomposition rule names to compiled mlir modules. - """ - op_decomp_rules = qp.decomposition.decomposition_graph.list_decomps(op_class) - - mlir_modules: dict[str, str | None] = {} - - if not hasattr(op_class, "num_wires") or not op_class.num_wires: - warnings.warn( - f"Cannot compile decomposition rules for op {op_class.__name__} with an unknown number " - + "of wires." + return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else "" + + +def parse_operator_data(op): + """Parse operator data from an Operator/Operator2 instance.""" + if isinstance(op, Operator2): + # TODO: use real getters here + dynamic_shape = op.getDynamicShape() + wire_lens = op.getWireLens() + static_data = op.getStaticData() + return dynamic_shape, wire_lens, static_data + if issubclass(op, Operator): + # NOTE: handling this the old-fashioned way, remove once Operator2 migration is complete + dynamic_shape = get_abstract_args(op) + num_wires = op.num_wires if op.num_wires else 0 + return dynamic_shape, [num_wires], {} + else: + raise ValueError( + "Only AbstractOperator and CompressedResourceOp types are supported for generating a " + f"graph ID, got {op} of type {type(op)}" ) - return mlir_modules - - dev = qp.device("null.qubit", wires=op_class.num_wires) - - abstract_args = get_abstract_args(op_class) # pylint: disable=protected-access - - for rule in op_decomp_rules: - try: - rule_name = rule._impl.__name__ # pylint: disable=protected-access - mlir_modules[rule_name] = compile_rule( - op_class, abstract_args, op_class.num_wires, rule, dev - ) - except CompileError as e: - warnings.warn(f"Failed to compile {rule_name}: {e}") - except Exception as e: # pylint: disable=broad-exception-caught - warnings.warn(f"Unexpected error while trying to compile {rule_name}: {e}") - finally: - qp.decomposition.disable_graph() - - return mlir_modules def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): @@ -216,11 +150,21 @@ def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): """ Path(decomp_file_path).parent.mkdir(parents=True, exist_ok=True) - mlir_rules = "".join( - str(mlir).replace("@rule_wrapper", f"@__builtin_{name}") - for func in COMPILER_OPS_FOR_DECOMPOSITION - for name, mlir in compile_op_decomp_rules(func).items() - ) + bytecode_lib = "" + + with ir.Context(): + # TODO: update this for Operator2, PL will implement a precompilation registry + for op in COMPILER_OPS_FOR_DECOMPOSITION: + dynamic_data, wire_lens, static_data = parse_operator_data(op) + if static_data: + # we cannot precompile if the rule takes static data + continue + + mlir_rules = python_decomposition( + op.__name__, get_graph_op_id(op), dynamic_data, wire_lens, {} + ) + + bytecode_lib += get_rules_from_module(mlir_rules) + "\n" bytecode = _quantum_opt( "--emit-bytecode", @@ -228,7 +172,7 @@ def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): "--convert-to-value-semantics", "--canonicalize", "--register-decomp-rule-resource", - stdin=mlir_rules.encode("utf-8"), + stdin=bytecode_lib.encode("utf-8"), text=None, ) diff --git a/frontend/test/pytest/test_QPD.py b/frontend/test/pytest/test_QPD.py deleted file mode 100644 index a4986d9d7d..0000000000 --- a/frontend/test/pytest/test_QPD.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright 2026 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. - -"""Unit tests for the python decompositions module.""" - -import pennylane as qp -import pytest - -from catalyst.device.python_decompositions import python_decomposition_wrapper - - -class TestQPD: - """Test the python wrapper functions used for compile-time decomposition rule lowering.""" - - def test_paulirot_wrapper(self): - """Test that the paulirot QPD wrapper correctly returns the IR as a string.""" - result = python_decomposition_wrapper( - "PauliRot", "PauliRot[f64][3]{pauli_word:XZZ}", [0.4], [3], {"pauli_word": "XZZ"} - ) - assert isinstance(result, str) - assert "PauliRot[f64][3]{pauli_word:XZZ}__pauli_rot_decomposition" in result - assert "Hadamard" in result - assert "multirz" in result - - def test_multiple_rules(self): - """Test that the python decomposition wrapper supports multiple rules.""" - with qp.decomposition.local_decomps(): - - def test_resources(): - return {qp.X: 1} - - @qp.register_resources(test_resources) - def test_decomp(angle, wires, pauli_word): # pylint: disable=unused-argument - qp.RX(angle, wires[0]) - - qp.add_decomps(qp.PauliRot, test_decomp) - - result = python_decomposition_wrapper( - "PauliRot", "PauliRot[f64][3]{pauli_word:XYX}", [float], [3], {"pauli_word": "XYX"} - ) - - assert "PauliRot[f64][3]{pauli_word:XYX}_test_decomp" in result - assert "test_decomp" in result - - -if __name__ == "__main__": - pytest.main(["-x", __file__]) diff --git a/frontend/test/pytest/test_precompile_decomp_rules.py b/frontend/test/pytest/test_precompile_decomp_rules.py deleted file mode 100644 index f2e78327d6..0000000000 --- a/frontend/test/pytest/test_precompile_decomp_rules.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright 2026 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. - -"""Tests for the decomposition rule precompilation utilities.""" - -from pathlib import Path - -import pennylane as qp -import pytest - -from catalyst.compiler import _quantum_opt -from catalyst.utils.precompile_decomposition_rules import ( - compile_op_decomp_rules, - get_abstract_args, - precompile_decomp_rules, -) -from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH - - -class TestGetAbstractArgs: - """Tests for get_abstract_args.""" - - def test_ignore_wires(self): - """Test that get_abstract_args correctly ignores WiresLike params.""" - assert not get_abstract_args(qp.X) - - def test_missing_ndim_params(self): - """Test that get_abstract_args correctly handles missing ndim_params properties.""" - assert not get_abstract_args(qp.H) - - def test_0_in_ndim_params(self): - """Test that get_abstract_args correctly handles 0 values in ndim_params.""" - assert get_abstract_args(qp.RX) == [float] - - def test_multiple_values_in_ndim_params(self): - """Test that get_abstract_args correctly handles length > 1 ndim_params.""" - assert get_abstract_args(qp.U3) == [float, float, float] - - def test_dimension_failure(self): - """ - Test that get_abstract_args correctly raises an exception given an op with - multi-dimensional params. - """ - with pytest.raises(ValueError, match="Cannot generate arguments"): - get_abstract_args(qp.ControlledQubitUnitary) - - -class TestCompileOpDecompRules: - """Tests for compile_op_decomp_rules.""" - - def test_compile_hadamard_rules(self): - """Test that compile_op_decomp_rules successfully compiles each decomp rule for Hadamard.""" - rules = compile_op_decomp_rules(qp.H) - - assert "_hadamard_to_rz_rx" in rules - assert "_hadamard_to_rz_ry" in rules - - def test_compile_rx_rules(self): - """Test that compile_op_decomp_rules successfully compiles each decomp rule for RX gates.""" - rules = compile_op_decomp_rules(qp.RX) - - assert "_rx_to_rot" in rules - assert "_rx_to_rz_ry" in rules - assert "_rx_to_ry_cliff" in rules - assert "_rx_to_rz_cliff" in rules - assert "_rx_to_ppr" in rules - - def test_fails_with_unknown_wires(self): - """ - Test that compile_op_decomp_rules warns when the number of wires is unknown. - """ - with pytest.warns(): - compile_op_decomp_rules(qp.Identity) - - def test_compile_error(self): - """Test that compile_op_decomp_rules warns when compilation of a rule fails.""" - - with pytest.warns(match="Failed to compile"): - with qp.decomposition.local_decomps(): - - class FakeOp(qp.operation.Operator): - """Test class with incompatible decomp rule.""" - - num_wires = 3 - num_params = 1 - ndim_params = (0,) - - @qp.register_resources({}) - def fake_op_decomp(param, wires): - _quantum_opt(stdin="module {") - return param, wires - - qp.add_decomps(FakeOp, fake_op_decomp) - - compile_op_decomp_rules(FakeOp) - - def test_unexpected_error(self): - """Test that compile_op_decomp_rules warns when an unexpected exception is thrown.""" - - with pytest.warns(match="Unexpected error"): - - with qp.decomposition.local_decomps(): - - class NewFakeOp(qp.operation.Operator): - """Test class without ndim_params.""" - - num_wires = 1 - num_params = 1 - - @qp.register_resources({}) - def fake_op_decomp(string): - qp.PauliRot(2, string) - - qp.add_decomps(NewFakeOp, fake_op_decomp) - - compile_op_decomp_rules(NewFakeOp) - - -def test_bytecode_file(): - """Test that the bytecode file is generated correctly.""" - orig_bcfile = Path(BYTECODE_FILE_PATH) - tmp_bcfile = None - - if orig_bcfile.exists(): - tmp_bcfile = orig_bcfile.replace(BYTECODE_FILE_PATH + ".tmpbackup") - - try: - precompile_decomp_rules() - assert orig_bcfile.exists() - - finally: - if tmp_bcfile: - tmp_bcfile = tmp_bcfile.replace(orig_bcfile) - else: - orig_bcfile.unlink(missing_ok=True) - - # NOTE: empty pass is needed to prevent running default pipeline - rules = _quantum_opt("--empty", BYTECODE_FILE_PATH) - - assert "_isingxy_to_h_cy" in rules - assert "_doublexcit" in rules - assert "_pauliz_to_ps" in rules - assert "_cphase_to_ppr" in rules - assert "_crot" in rules - - -if __name__ == "__main__": - pytest.main(["-x", __file__]) diff --git a/frontend/test/pytest/test_rule_lowering.py b/frontend/test/pytest/test_rule_lowering.py new file mode 100644 index 0000000000..6cd00b5fab --- /dev/null +++ b/frontend/test/pytest/test_rule_lowering.py @@ -0,0 +1,133 @@ +# Copyright 2026 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. + +"""Unit tests for the python decompositions module.""" + +from pathlib import Path + +import pennylane as qp +import pytest + +from catalyst.compiler import _quantum_opt +from catalyst.device.python_decompositions import python_decomposition_wrapper +from catalyst.utils.precompile_decomposition_rules import ( + get_abstract_args, + precompile_decomp_rules, +) +from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH + + +class TestPrecompiled: + """Tests for precompiled decomposition rules.""" + + # Tests for get_abstract_args helper + + def test_ignore_wires(self): + """Test that get_abstract_args correctly ignores WiresLike params.""" + assert not get_abstract_args(qp.X) + + def test_missing_ndim_params(self): + """Test that get_abstract_args correctly handles missing ndim_params properties.""" + assert not get_abstract_args(qp.H) + + def test_0_in_ndim_params(self): + """Test that get_abstract_args correctly handles 0 values in ndim_params.""" + assert get_abstract_args(qp.RX) == [float] + + def test_multiple_values_in_ndim_params(self): + """Test that get_abstract_args correctly handles length > 1 ndim_params.""" + assert get_abstract_args(qp.U3) == [float, float, float] + + def test_dimension_failure(self): + """ + Test that get_abstract_args correctly raises an exception given an op with + multi-dimensional params. + """ + with pytest.raises(ValueError, match="Cannot generate arguments"): + get_abstract_args(qp.ControlledQubitUnitary) + + # Test for bytecode file + + def test_bytecode_file(self): + """Test that the bytecode file is generated correctly.""" + orig_bcfile = Path(BYTECODE_FILE_PATH) + tmp_bcfile = None + + if orig_bcfile.exists(): + tmp_bcfile = orig_bcfile.replace(BYTECODE_FILE_PATH + ".tmpbackup") + + try: + precompile_decomp_rules() + assert orig_bcfile.exists() + + finally: + if tmp_bcfile: + tmp_bcfile = tmp_bcfile.replace(orig_bcfile) + else: + orig_bcfile.unlink(missing_ok=True) + + # NOTE: empty pass is needed to prevent running default pipeline + rules = _quantum_opt("--empty", BYTECODE_FILE_PATH) + + assert "__builtin__isingxy_to_h_cy" in rules + assert "__builtin__doublexcit" in rules + assert "__builtin__pauliz_to_ps" in rules + assert "__builtin__cphase_to_ppr" in rules + assert "__builtin__crot" in rules + + +class TestTraceTime: + """Placeholder for future tests of trace-time decomposition rule lowering.""" + + +class TestOnDemand: + """ + Test the python wrapper functions used for on-demand, compile-time decomposition rule lowering. + """ + + def test_paulirot(self): + """Test that the QPD wrapper correctly returns the IR as a string.""" + result = python_decomposition_wrapper( + "PauliRot", "PauliRot[f64][3]{pauli_word:XZZ}", ["i32"], [3], {"pauli_word": "XZZ"} + ) + assert isinstance(result, str) + assert "_pauli_rot_decomposition" in result + assert 'target_gate = "PauliRot[f64][3]{pauli_word:XZZ}"' in result + assert "Hadamard" in result + assert "multirz" in result + + def test_multiple_rules(self): + """Test that the python decomposition wrapper supports multiple rules.""" + with qp.decomposition.local_decomps(): + + def test_resources(pauli_word): # pylint: disable=unused-argument + return {qp.X: 1} + + @qp.register_resources(test_resources) + def test_decomp(angle, wires, pauli_word): # pylint: disable=unused-argument + qp.RX(angle, wires[0]) + + qp.add_decomps(qp.PauliRot, test_decomp) + + result = python_decomposition_wrapper( + "PauliRot", "PauliRot[f64][3]{pauli_word:XYX}", ["f64"], [3], {"pauli_word": "XYX"} + ) + + assert "test_decomp" in result + assert "_pauli_rot_decomp" in result + assert 'target_gate = "PauliRot[f64][3]{pauli_word:XYX}"' in result + + +if __name__ == "__main__": + pytest.main(["-x", __file__]) diff --git a/mlir/lib/Quantum/IR/QuantumInterfaces.cpp b/mlir/lib/Quantum/IR/QuantumInterfaces.cpp index 7e4eca210d..ce8c358584 100644 --- a/mlir/lib/Quantum/IR/QuantumInterfaces.cpp +++ b/mlir/lib/Quantum/IR/QuantumInterfaces.cpp @@ -22,6 +22,10 @@ #include "llvm/Support/raw_ostream.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/TypeRange.h" +#include "mlir/IR/Types.h" #include "mlir/Support/LLVM.h" using namespace mlir; diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp index 48ce6e108e..d18512ebdf 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp @@ -399,12 +399,11 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBasewalk([&](mlir::func::FuncOp func) { - if (func.getName().starts_with(opId)) { + if (func->hasAttr("target_gate")) { mlir::OwningOpRef outOp; func->remove(); outOp = mlir::OwningOpRef(func); mlir::func::FuncOp funcOp = outOp.get(); - funcOp->setAttr("target_gate", mlir::StringAttr::get(context, opId)); // if we fail to add one of the decomps, we still want to try for the rest std::ignore = addRuleNode(funcOp, ruleNodes); diff --git a/mlir/unittests/Quantum/Interfaces/TestDecomposableGateInterface.cpp b/mlir/unittests/Quantum/Interfaces/TestDecomposableGateInterface.cpp index f179501284..2d31be2071 100644 --- a/mlir/unittests/Quantum/Interfaces/TestDecomposableGateInterface.cpp +++ b/mlir/unittests/Quantum/Interfaces/TestDecomposableGateInterface.cpp @@ -14,12 +14,15 @@ // limitations under the License. #include +#include #include #include #include "gtest/gtest.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Format.h" // for gtest printing on failure +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/raw_ostream.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/AsmState.h" @@ -34,6 +37,7 @@ #include "mlir/IR/TypeRange.h" #include "mlir/IR/Types.h" #include "mlir/Parser/Parser.h" +#include "mlir/Support/LLVM.h" #include "Quantum/IR/QuantumDialect.h" #include "Quantum/IR/QuantumInterfaces.h" @@ -474,3 +478,46 @@ func.func @testfunc(%first : tensor<1xi64>, %secondthird : tensor<2xi64>) { ASSERT_EQ(op.getGraphOpId(), "testOperatorUID[i1,f64,i64][1,2]{}[248]"); } + +TEST(DecomposableGateInterfaceTests, OperatorOpNestedArgs) +{ + std::string moduleStr = R"mlir( +func.func @testfunc(%arg: tensor<3xi64>, %arg2: tensor<2xf64>) { + + %reg = quantum.alloc(4) : !quantum.reg + %q0 = quantum.extract %reg[0] : !quantum.reg -> !quantum.bit + + %0 = quantum.operator "testNestedArgs"(%arg : tensor<3xi64>, %arg2: tensor<2xf64>) + qubits(%q0) + return +} + )mlir"; + + // Parsing boilerplate + DialectRegistry registry; + registry.insert(); + MLIRContext context(registry); + ParserConfig config(&context, /*verifyAfterParse=*/false); + OwningOpRef module = parseSourceString(moduleStr, config); + + DecomposableGate op; + module->walk([&](OperatorOp walkOp) { op = walkOp; }); + + ASSERT_EQ(op.getOperatorName(), "testNestedArgs"); + + // This is needed to keep the backing array from being deleted + mlir::Type int64 = mlir::IntegerType::get(&context, 64); + llvm::SmallVector backing( + {mlir::RankedTensorType::get(ArrayRef({3}), int64), + mlir::RankedTensorType::get(ArrayRef({2}), mlir::Float64Type::get(&context))}); + mlir::TypeRange expectedDynamicShape(backing); + mlir::TypeRange actual(op.getDynamicShape()); + ASSERT_EQ(llvm::SmallVector(actual), + llvm::SmallVector(expectedDynamicShape)); + + ASSERT_EQ(op.getWireLens(), std::vector({1})); + + ASSERT_EQ(op.getStaticData(), mlir::DictionaryAttr::get(&context, {})); + + ASSERT_EQ(op.getGraphOpId(), "testNestedArgs[[i64,i64,i64],[f64,f64]][1]{}"); +}