From 034e5003c2084f46ed2e9699f1d6072a683c3e09 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 19 Jun 2026 16:10:15 -0400 Subject: [PATCH 01/48] add translation and lowering of OperatorOp --- .../catalyst/from_plxpr/qfunc_interpreter.py | 22 ++++++++++++++++ .../from_plxpr/qref_jax_primitives.py | 26 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index a1f8a8c098..fca616acd2 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -25,6 +25,7 @@ import pennylane as qp from jax._src.sharding_impls import UNSPECIFIED from pennylane.capture import PlxprInterpreter, pause +from pennylane.core.operator.operator2 import operator_p from pennylane.capture.primitives import cond_prim as pl_cond_prim from pennylane.capture.primitives import ctrl_transform_prim as plxpr_ctrl_transform_prim from pennylane.capture.primitives import measure_prim as plxpr_measure_prim @@ -51,6 +52,7 @@ qref_set_basis_state_p, qref_set_state_p, qref_unitary_p, + qref_operator_op, ) from catalyst.jax_primitives import ( counts_p, @@ -290,6 +292,26 @@ def __call__(self, jaxpr, *args): return self.eval(jaxpr.jaxpr, jaxpr.consts, *args) +@PLxPRToQuantumJaxprInterpreter.register_primitive(operator_p) +def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, **kwargs): + + if hybrid_lens or hybrid_trees: + raise NotImplementedError + + wire_inputs = args[len(op_cls.dynamic_argnames):] + new_wires = [w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in wire_inputs] + + qref_operator_op.bind( + *args[:len(op_cls.dynamic_argnames)], + *new_wires, + op_cls=op_cls, + hybrid_lens=hybrid_lens, + hybrid_trees=hybrid_trees, + **kwargs + ) + return [] + + # pylint: disable=unused-argument, too-many-arguments def _qubit_unitary_bind_call( *invals, op, qubits_len, params_len, ctrl_len, adjoint, hyperparameters diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 4e3e94d9a4..c31768a786 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -72,6 +72,7 @@ MeasureOp, MultiRZOp, NamedObsOp, + OperatorOp, PauliRotOp, PCPhaseOp, QubitUnitaryOp, @@ -164,6 +165,7 @@ class MeasurementPlane(Enum): qref_compbasis_p = Primitive("qref_compbasis") qref_namedobs_p = Primitive("qref_namedobs") qref_hermitian_p = Primitive("qref_hermitian") +qref_operator_op = Primitive("qref_operator") # @@ -797,6 +799,29 @@ def _qref_named_obs_lowering(jax_ctx: mlir.LoweringRuleContext, qubit: ir.Value, return NamedObsOp(result_type, qubit, obsId).results +qref_operator_op.multiple_results = True + +@qref_operator_op.def_abstract_eval +def _qref_operator_op_abstract_eval(*args, **kwargs): + return [] + +def _operator_op_lowering(jax_ctx: mlir.LoweringRuleContext, *args, op_cls, hybrid_lens, hybrid_trees, wire_lens, **static_args): + params = args[:len(op_cls.dynamic_argnames)] + qubits = args[len(op_cls.dynamic_argnames):] + + name_attr = ir.StringAttr.get(op_cls.__name__) + + OperatorOp(op_name=name_attr, + params = params, + qubits = qubits, + forward_args=[], + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + UID=None, + arr_qubit_indices=[] + ) + return [] # # hermitian observable @@ -821,6 +846,7 @@ def _qref_hermitian_lowering(jax_ctx: mlir.LoweringRuleContext, matrix: ir.Value CUSTOM_LOWERING_RULES = ( + (qref_operator_op, _operator_op_lowering), (qref_alloc_p, _qref_alloc_lowering), (qref_dealloc_p, _qref_dealloc_lowering), (qref_get_p, _qref_get_lowering), From a9308fc8df8ab63ca57dad127d79bbd2808e299b Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 19 Jun 2026 17:29:03 -0400 Subject: [PATCH 02/48] sonnet's AI recommendation to fix python bindings issue --- mlir/python/dialects/qref.py | 45 ++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/mlir/python/dialects/qref.py b/mlir/python/dialects/qref.py index 5522c08da8..951de35d86 100644 --- a/mlir/python/dialects/qref.py +++ b/mlir/python/dialects/qref.py @@ -16,3 +16,48 @@ # pylint: disable=relative-beyond-top-level from ._qref_ops_gen import * # noqa: F401 +from ._qref_ops_gen import OperatorOp as _OperatorOpGen +from ._ods_common import ( + get_default_loc_context as _ods_get_default_loc_context, + get_op_results_or_values as _get_op_results_or_values, +) +from ._ods_common import _cext as _ods_cext +_ods_ir = _ods_cext.ir + + +class OperatorOp(_OperatorOpGen): + def __init__(self, op_name, params, forward_args, qubits, ctrl_qubits, + ctrl_values, arr_qubit_indices=None, adjoint=False, UID=None, *, + qreg=None, arr_ctrl_indices=None, arr_ctrl_values=None, + static_data=None, param_map=None, qubit_map=None, loc=None, ip=None): + operands = [] + attributes = {} + # op_name is a StringProp — must go in attributes + attributes["op_name"] = (op_name if isinstance(op_name, _ods_ir.Attribute) + else _ods_ir.StringAttr.get(op_name)) + operands.append(_get_op_results_or_values(params)) + operands.append(_get_op_results_or_values(forward_args)) + operands.append(_get_op_results_or_values(qubits)) + operands.append(_get_op_results_or_values(ctrl_qubits)) + operands.append(ctrl_values if ctrl_values is not None else []) + operands.append(_get_op_results_or_values(qreg)) + operands.append(arr_qubit_indices if arr_qubit_indices is not None else []) + operands.append(arr_ctrl_indices) + operands.append(arr_ctrl_values) + _ods_context = _ods_get_default_loc_context(loc) + if bool(adjoint): + attributes["adjoint"] = _ods_ir.UnitAttr.get(_ods_context) + if UID is not None: + attributes["UID"] = _ods_ir.IntegerAttr.get( + _ods_ir.IntegerType.get_signless(64), UID) + if static_data is not None: + attributes["static_data"] = static_data + if param_map is not None: + attributes["param_map"] = param_map + if qubit_map is not None: + attributes["qubit_map"] = qubit_map + # bypass _OperatorOpGen.__init__ and go directly to OpView + super(_OperatorOpGen, self).__init__( + self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, + self._ODS_RESULT_SEGMENTS, attributes=attributes, results=[], + operands=operands, successors=None, regions=None, loc=loc, ip=ip) \ No newline at end of file From 8599ec6219b754d73edf9eb78805b1860203a7c9 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 10:02:36 -0400 Subject: [PATCH 03/48] switch to StrAttr --- mlir/include/QRef/IR/QRefOps.td | 2 +- mlir/include/Quantum/IR/QuantumOps.td | 2 +- mlir/lib/QRef/IR/QRefOps.cpp | 3 ++- mlir/lib/Quantum/IR/QuantumOps.cpp | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/mlir/include/QRef/IR/QRefOps.td b/mlir/include/QRef/IR/QRefOps.td index 075f871e2c..028916927f 100644 --- a/mlir/include/QRef/IR/QRefOps.td +++ b/mlir/include/QRef/IR/QRefOps.td @@ -412,7 +412,7 @@ def OperatorOp : UnitaryGate_Op<"operator", [ParametrizedGate, AttrSizedOperandS }]; let arguments = (ins - StringProp:$op_name, + StrAttr:$op_name, Variadic:$params, Variadic:$forward_args, diff --git a/mlir/include/Quantum/IR/QuantumOps.td b/mlir/include/Quantum/IR/QuantumOps.td index 5b7db9bebf..8623922262 100644 --- a/mlir/include/Quantum/IR/QuantumOps.td +++ b/mlir/include/Quantum/IR/QuantumOps.td @@ -613,7 +613,7 @@ def OperatorOp : UnitaryGate_Op<"operator", [ParametrizedGate, NoMemoryEffect, }]; let arguments = (ins - StringProp:$op_name, + StrAttr:$op_name, Variadic:$params, Variadic:$forward_args, diff --git a/mlir/lib/QRef/IR/QRefOps.cpp b/mlir/lib/QRef/IR/QRefOps.cpp index f53e3be8cc..8f4e24cd9d 100644 --- a/mlir/lib/QRef/IR/QRefOps.cpp +++ b/mlir/lib/QRef/IR/QRefOps.cpp @@ -594,8 +594,9 @@ ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &result) if (parser.parseString(&opName)) { return failure(); } + result.addAttribute("op_name", builder.getStringAttr(opName)); auto &opProperties = result.getOrAddProperties(); - opProperties.setOpName(opName); + // 2. Parse variadic params: (%arg0: type, ...) SmallVector params; diff --git a/mlir/lib/Quantum/IR/QuantumOps.cpp b/mlir/lib/Quantum/IR/QuantumOps.cpp index 736d9f2d4b..821003d195 100644 --- a/mlir/lib/Quantum/IR/QuantumOps.cpp +++ b/mlir/lib/Quantum/IR/QuantumOps.cpp @@ -899,7 +899,7 @@ ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &result) return failure(); } auto &opProperties = result.getOrAddProperties(); - opProperties.setOpName(opName); + result.addAttribute("op_name", builder.getStringAttr(opName)); // 2. Parse variadic params: (%arg0: type, ...) SmallVector params; From 5c4504e4f19b04edae3a539444cd0499dcb2fa87 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 10:23:51 -0400 Subject: [PATCH 04/48] bump pl version --- .dep-versions | 2 +- doc/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.dep-versions b/.dep-versions index 4ca08f1d5c..eff2c832ec 100644 --- a/.dep-versions +++ b/.dep-versions @@ -8,7 +8,7 @@ enzyme=v0.0.238 # For a custom PL version, update the package version here and at # 'doc/requirements.txt' -pennylane=0.46.0.dev24 +pennylane=0.46.0.dev38 # For a custom LQ/LK version, update the package version here and at # 'doc/requirements.txt' diff --git a/doc/requirements.txt b/doc/requirements.txt index c6721a731a..3459767b3b 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -34,4 +34,4 @@ lxml_html_clean --extra-index-url https://test.pypi.org/simple/ pennylane-lightning-kokkos==0.46.0-dev10 pennylane-lightning==0.46.0-dev10 -pennylane==0.46.0.dev24 +pennylane==0.46.0.dev38 From b23289ad514afd768691ae0b61233cee122b6858 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 10:58:05 -0400 Subject: [PATCH 05/48] remove qref patch --- mlir/python/dialects/qref.py | 45 ------------------------------------ 1 file changed, 45 deletions(-) diff --git a/mlir/python/dialects/qref.py b/mlir/python/dialects/qref.py index 951de35d86..5522c08da8 100644 --- a/mlir/python/dialects/qref.py +++ b/mlir/python/dialects/qref.py @@ -16,48 +16,3 @@ # pylint: disable=relative-beyond-top-level from ._qref_ops_gen import * # noqa: F401 -from ._qref_ops_gen import OperatorOp as _OperatorOpGen -from ._ods_common import ( - get_default_loc_context as _ods_get_default_loc_context, - get_op_results_or_values as _get_op_results_or_values, -) -from ._ods_common import _cext as _ods_cext -_ods_ir = _ods_cext.ir - - -class OperatorOp(_OperatorOpGen): - def __init__(self, op_name, params, forward_args, qubits, ctrl_qubits, - ctrl_values, arr_qubit_indices=None, adjoint=False, UID=None, *, - qreg=None, arr_ctrl_indices=None, arr_ctrl_values=None, - static_data=None, param_map=None, qubit_map=None, loc=None, ip=None): - operands = [] - attributes = {} - # op_name is a StringProp — must go in attributes - attributes["op_name"] = (op_name if isinstance(op_name, _ods_ir.Attribute) - else _ods_ir.StringAttr.get(op_name)) - operands.append(_get_op_results_or_values(params)) - operands.append(_get_op_results_or_values(forward_args)) - operands.append(_get_op_results_or_values(qubits)) - operands.append(_get_op_results_or_values(ctrl_qubits)) - operands.append(ctrl_values if ctrl_values is not None else []) - operands.append(_get_op_results_or_values(qreg)) - operands.append(arr_qubit_indices if arr_qubit_indices is not None else []) - operands.append(arr_ctrl_indices) - operands.append(arr_ctrl_values) - _ods_context = _ods_get_default_loc_context(loc) - if bool(adjoint): - attributes["adjoint"] = _ods_ir.UnitAttr.get(_ods_context) - if UID is not None: - attributes["UID"] = _ods_ir.IntegerAttr.get( - _ods_ir.IntegerType.get_signless(64), UID) - if static_data is not None: - attributes["static_data"] = static_data - if param_map is not None: - attributes["param_map"] = param_map - if qubit_map is not None: - attributes["qubit_map"] = qubit_map - # bypass _OperatorOpGen.__init__ and go directly to OpView - super(_OperatorOpGen, self).__init__( - self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, - self._ODS_RESULT_SEGMENTS, attributes=attributes, results=[], - operands=operands, successors=None, regions=None, loc=loc, ip=ip) \ No newline at end of file From 6d4de2b446732617a011be7283b0326df4a6abc6 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 12:54:03 -0400 Subject: [PATCH 06/48] formatting --- mlir/lib/QRef/IR/QRefOps.cpp | 3 +-- mlir/lib/Quantum/IR/QuantumOps.cpp | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/mlir/lib/QRef/IR/QRefOps.cpp b/mlir/lib/QRef/IR/QRefOps.cpp index 8f4e24cd9d..3e754fc03b 100644 --- a/mlir/lib/QRef/IR/QRefOps.cpp +++ b/mlir/lib/QRef/IR/QRefOps.cpp @@ -482,7 +482,7 @@ void OperatorOp::print(OpAsmPrinter &p) // 5. Attribute Dictionary SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", - "operandSegmentSizes"}; + "operandSegmentSizes", "op_name"}; p.printOptionalAttrDict(getOperation()->getAttrs(), elidedAttrs); p.increaseIndent(); @@ -596,7 +596,6 @@ ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &result) } result.addAttribute("op_name", builder.getStringAttr(opName)); auto &opProperties = result.getOrAddProperties(); - // 2. Parse variadic params: (%arg0: type, ...) SmallVector params; diff --git a/mlir/lib/Quantum/IR/QuantumOps.cpp b/mlir/lib/Quantum/IR/QuantumOps.cpp index 821003d195..62977102c4 100644 --- a/mlir/lib/Quantum/IR/QuantumOps.cpp +++ b/mlir/lib/Quantum/IR/QuantumOps.cpp @@ -785,8 +785,8 @@ void OperatorOp::print(OpAsmPrinter &p) } // 5. Attribute Dictionary - SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", - "operandSegmentSizes", "resultSegmentSizes"}; + SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", + "operandSegmentSizes", "resultSegmentSizes", "op_name"}; p.printOptionalAttrDict(getOperation()->getAttrs(), elidedAttrs); p.increaseIndent(); From daca7ae39dc2b802878a8b71fa12784f49899db9 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 12:54:36 -0400 Subject: [PATCH 07/48] black --- .../catalyst/from_plxpr/qfunc_interpreter.py | 10 +++-- .../from_plxpr/qref_jax_primitives.py | 43 ++++++++++++------- 2 files changed, 34 insertions(+), 19 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index fca616acd2..c7e3cb45a2 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -298,16 +298,18 @@ def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, **kwargs): if hybrid_lens or hybrid_trees: raise NotImplementedError - wire_inputs = args[len(op_cls.dynamic_argnames):] - new_wires = [w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in wire_inputs] + wire_inputs = args[len(op_cls.dynamic_argnames) :] + new_wires = [ + w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in wire_inputs + ] qref_operator_op.bind( - *args[:len(op_cls.dynamic_argnames)], + *args[: len(op_cls.dynamic_argnames)], *new_wires, op_cls=op_cls, hybrid_lens=hybrid_lens, hybrid_trees=hybrid_trees, - **kwargs + **kwargs, ) return [] diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index c31768a786..4dd957288a 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -799,30 +799,43 @@ def _qref_named_obs_lowering(jax_ctx: mlir.LoweringRuleContext, qubit: ir.Value, return NamedObsOp(result_type, qubit, obsId).results + qref_operator_op.multiple_results = True + @qref_operator_op.def_abstract_eval def _qref_operator_op_abstract_eval(*args, **kwargs): return [] -def _operator_op_lowering(jax_ctx: mlir.LoweringRuleContext, *args, op_cls, hybrid_lens, hybrid_trees, wire_lens, **static_args): - params = args[:len(op_cls.dynamic_argnames)] - qubits = args[len(op_cls.dynamic_argnames):] - + +def _operator_op_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, + **static_args, +): + params = args[: len(op_cls.dynamic_argnames)] + qubits = args[len(op_cls.dynamic_argnames) :] + name_attr = ir.StringAttr.get(op_cls.__name__) - - OperatorOp(op_name=name_attr, - params = params, - qubits = qubits, - forward_args=[], - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - UID=None, - arr_qubit_indices=[] - ) + + OperatorOp( + op_name=name_attr, + params=params, + qubits=qubits, + forward_args=[], + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + UID=None, + arr_qubit_indices=[], + ) return [] + # # hermitian observable # From 79e060e3ac936908de17847ff8b530d9914d5f77 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Wed, 24 Jun 2026 16:58:46 -0400 Subject: [PATCH 08/48] starting to add some tests --- frontend/catalyst/from_plxpr/from_plxpr.py | 4 + .../catalyst/from_plxpr/qfunc_interpreter.py | 3 +- .../from_plxpr/qref_jax_primitives.py | 23 ++- frontend/catalyst/jax_extras/lowering.py | 3 + frontend/test/lit/test_operator.py | 146 ++++++++++++++++++ 5 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 frontend/test/lit/test_operator.py diff --git a/frontend/catalyst/from_plxpr/from_plxpr.py b/frontend/catalyst/from_plxpr/from_plxpr.py index f3b0366f57..bafceb02b0 100644 --- a/frontend/catalyst/from_plxpr/from_plxpr.py +++ b/frontend/catalyst/from_plxpr/from_plxpr.py @@ -25,6 +25,7 @@ import jax import pennylane as qp from jax.extend.core import ClosedJaxpr, Jaxpr +from pennylane.core.operator.operator2 import operator_p from pennylane.capture import PlxprInterpreter, qnode_prim from pennylane.capture.primitives import transform_prim from pennylane.transforms import decompose as pl_decompose @@ -407,6 +408,9 @@ def _handle_decompose_transform(self, inner_jaxpr, consts, non_const_args, tkwar # return self.eval(final_jaxpr.jaxpr, consts, *non_const_args) return next_eval.eval(inner_jaxpr, consts, *non_const_args) +@WorkflowInterpreter.register_primitive(operator_p) +def _error_on_operator(self, *args, op_cls, **kwargs): + raise ValueError(f"Operator {op_cls} must occur inside a qnode.") # pylint: disable=too-many-arguments @WorkflowInterpreter.register_primitive(transform_prim) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index c7e3cb45a2..918139323c 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -295,7 +295,8 @@ def __call__(self, jaxpr, *args): @PLxPRToQuantumJaxprInterpreter.register_primitive(operator_p) def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, **kwargs): - if hybrid_lens or hybrid_trees: + if hybrid_lens or hybrid_trees or op_cls.static_argnames: + # only support compilable_argnames for the moment raise NotImplementedError wire_inputs = args[len(op_cls.dynamic_argnames) :] diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 4dd957288a..315003a187 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -30,6 +30,7 @@ from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp from pennylane.capture.primitives import adjoint_transform_prim as plxpr_adjoint_transform_prim from pennylane.wires import AbstractQubit +from pennylane.pytrees import unflatten # TODO: remove after jax v0.7.2 upgrade # Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between @@ -38,6 +39,7 @@ # once JAX updates to a compatible MLIR version # pylint: disable=ungrouped-imports from catalyst.jax_extras.patches import mock_attributes +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.jax_primitives import ( AbstractObs, _named_obs_attribute, @@ -815,12 +817,26 @@ def _operator_op_lowering( hybrid_lens, hybrid_trees, wire_lens, - **static_args, + **static_data, ): params = args[: len(op_cls.dynamic_argnames)] qubits = args[len(op_cls.dynamic_argnames) :] - name_attr = ir.StringAttr.get(op_cls.__name__) + name_attr = get_mlir_attribute_from_pyval(op_cls.__name__) + + repack_static_data = {k: unflatten(*v) for k, v in static_data.items()} + processed_static_data = get_mlir_attribute_from_pyval(repack_static_data) + + param_map = {name: ir.DenseI64ArrayAttr.get([ind]) for ind, name in enumerate(op_cls.dynamic_argnames)} + processed_param_map = get_mlir_attribute_from_pyval(param_map) + + qubit_map = {} + ind = 0 + for name, size in zip(op_cls.wire_argnames, wire_lens): + qubit_map[name] = ir.DenseI64ArrayAttr.get(list(range(ind, ind+size))) + ind += size + + processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) OperatorOp( op_name=name_attr, @@ -832,6 +848,9 @@ def _operator_op_lowering( adjoint=False, UID=None, arr_qubit_indices=[], + param_map=processed_param_map, + static_data=processed_static_data, + qubit_map=processed_qubit_map, ) return [] diff --git a/frontend/catalyst/jax_extras/lowering.py b/frontend/catalyst/jax_extras/lowering.py index d372890229..3812dc97a0 100644 --- a/frontend/catalyst/jax_extras/lowering.py +++ b/frontend/catalyst/jax_extras/lowering.py @@ -202,6 +202,9 @@ def get_mlir_attribute_from_pyval(value): attr = None match value: + case ir.Attribute(): + attr = value + case bool(): attr = ir.BoolAttr.get(value) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py new file mode 100644 index 0000000000..9260ddcb21 --- /dev/null +++ b/frontend/test/lit/test_operator.py @@ -0,0 +1,146 @@ +# Copyright 2022-2023 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 operator in Catalyst.""" + + +# RUN: %PYTHON %s | FileCheck %s + +import numpy as np +import pennylane as qp + +class NoParams(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +@qp.qjit(target="mlir", capture=True) +@qp.qnode(qp.device('null.qubit', wires=2)) +def c_no_params(): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: qref.operator "NoParams"() qubits([[q0]]) + # CHECK static_data = {} + # CHECK param_map = {} qubit_map = {wires = [0]} + NoParams(wires=0) + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.operator "NoParams"() qubits([[q1]], [[q2]]) + # CHECK: static_data = {} + # CHECK: param_map = {} qubit_map = {wires = [0, 1]} + NoParams(wires=(0,1)) + return qp.state() + +print(c_no_params.mlir) + +class SingleParam(qp.core.Operator2): + + dynamic_argnames = ("x", ) + + def __init__(self, x, wires): + super().__init__(x, wires=wires) + +@qp.qjit(target="mlir", capture=True) +@qp.qnode(qp.device('null.qubit', wires=3)) +def c_single_param(x : float): + + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + + # CHECK: qref.operator "SingleParam"({{%.+}}: tensor) qubits([[q0]]) + # CHECK: static_data = {} + # CHECK: param_map = {x = [0]} qubit_map = {wires = [0]} + SingleParam(x, 0) + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.operator "SingleParam"({{%.+}}: tensor<4x4xf64>) qubits([[q1]], [[q2]]) + # CHECK: static_data = {} + # CHECK: param_map = {x = [0]} qubit_map = {wires = [0, 1]} + SingleParam(np.eye(4), (1,2)) + + return qp.state() + +print(c_single_param.mlir) + + +class CompilableData(qp.core.Operator2): + + compilable_argnames = ("a", "b", "thing") + + def __init__(self, a, b, thing, wires): + super().__init__(a=a, b=b, thing=thing, wires=wires) + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device('null.qubit', wires=3)) +def c_compilable(): + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.operator "CompilableData"() qubits([[q1]], [[q2]]) + # CHECK: static_data = {a = true, b = "some string", thing = [1, true, "string"]} + # CHECK: param_map = {} qubit_map = {wires = [0, 1]} + + CompilableData(True, "some string", (1, True, "string"), wires=(0,1)) + + return qp.state() + + +print(c_compilable.mlir) + + +class MultipleRegisters(qp.core.Operator2): + + wire_argnames = ("reg1", "reg2") + + def __init__(self, reg1, reg2): + super().__init__(reg1=reg1, reg2=reg2) + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device('null.qubit', wires=5)) +def c_multiple_registers(): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + + # CHECK: qref.operator "MultipleRegisters"() qubits([[q0]], [[q2]], [[q3]], [[q4]]) + # CHECK: static_data = {} + # CHECK: param_map = {} qubit_map = {reg1 = [0], reg2 = [1, 2, 3]} + MultipleRegisters(0, (2,3,4)) + return qp.state() + + +print(c_multiple_registers.mlir) + + +class MultiParams(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + + # note also having non-standard order with dynamic inputs after wires + def __init__(self, wires, a, b, c): + super().__init__(wires, a, b, c) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device('null.qubit', wires=1)) +def c_multi_params(): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + + # CHECK: qref.operator "MultiParams"({{%.+}}: tensor, {{%.+}}: tensor<4x2x1xf64>, {{%.+}}: tensor<3xi64>) qubits([[q0]]) + # CHECK: static_data = {} + # CHECK: param_map = {a = [0], b = [1], c = [2]} qubit_map = {wires = [0]} + MultiParams(0, 0.5, c=np.array([1, 2 ,3]), b=np.zeros((4,2,1))) + return qp.state() + +print(c_multi_params.mlir) \ No newline at end of file From 98c88c4216a24dfddfc56bb9e340c1cd1713b50e Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 25 Jun 2026 10:28:02 -0400 Subject: [PATCH 09/48] tests, changelog, formatting --- doc/releases/changelog-dev.md | 4 ++ frontend/catalyst/from_plxpr/from_plxpr.py | 4 +- .../catalyst/from_plxpr/qfunc_interpreter.py | 4 +- .../from_plxpr/qref_jax_primitives.py | 11 +-- frontend/test/lit/test_operator.py | 39 +++++----- frontend/test/pytest/test_operator.py | 72 +++++++++++++++++++ .../Transforms/value_semantics_conversion.cpp | 1 - .../reference_semantics_conversion.cpp | 1 - 8 files changed, 111 insertions(+), 25 deletions(-) create mode 100644 frontend/test/pytest/test_operator.py diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index c990c2ad88..f5b0045e82 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -5,6 +5,10 @@

Improvements 🛠

+* The new `pennylane.core.Operator2` can now be lowered to MLIR with program capture for operators + without non-lowerable arguments. + [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969/) + * The `ResourceAnalysis` pass now reports each loop body and each subroutine as its own entry instead of folding their gate counts into the caller. Loops with constant bounds appear as `for_loop_` with their trip count. Loops with dynamic bounds appear as `dyn_for_loop_` with a stable diff --git a/frontend/catalyst/from_plxpr/from_plxpr.py b/frontend/catalyst/from_plxpr/from_plxpr.py index bafceb02b0..86bdc8ac88 100644 --- a/frontend/catalyst/from_plxpr/from_plxpr.py +++ b/frontend/catalyst/from_plxpr/from_plxpr.py @@ -25,9 +25,9 @@ import jax import pennylane as qp from jax.extend.core import ClosedJaxpr, Jaxpr -from pennylane.core.operator.operator2 import operator_p from pennylane.capture import PlxprInterpreter, qnode_prim from pennylane.capture.primitives import transform_prim +from pennylane.core.operator.operator2 import operator_p from pennylane.transforms import decompose as pl_decompose from catalyst.device import extract_backend_info @@ -408,10 +408,12 @@ def _handle_decompose_transform(self, inner_jaxpr, consts, non_const_args, tkwar # return self.eval(final_jaxpr.jaxpr, consts, *non_const_args) return next_eval.eval(inner_jaxpr, consts, *non_const_args) + @WorkflowInterpreter.register_primitive(operator_p) def _error_on_operator(self, *args, op_cls, **kwargs): raise ValueError(f"Operator {op_cls} must occur inside a qnode.") + # pylint: disable=too-many-arguments @WorkflowInterpreter.register_primitive(transform_prim) def handle_transform( diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index 918139323c..1f6b4af815 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -25,12 +25,12 @@ import pennylane as qp from jax._src.sharding_impls import UNSPECIFIED from pennylane.capture import PlxprInterpreter, pause -from pennylane.core.operator.operator2 import operator_p from pennylane.capture.primitives import cond_prim as pl_cond_prim from pennylane.capture.primitives import ctrl_transform_prim as plxpr_ctrl_transform_prim from pennylane.capture.primitives import measure_prim as plxpr_measure_prim from pennylane.capture.primitives import pauli_measure_prim as plxpr_pauli_measure_prim from pennylane.capture.primitives import quantum_subroutine_prim, transform_prim +from pennylane.core.operator.operator2 import operator_p from pennylane.ftqc.primitives import measure_in_basis_prim as plxpr_measure_in_basis_prim from pennylane.measurements import CountsMP from pennylane.wires import AbstractQubit, is_abstract_qubit @@ -46,13 +46,13 @@ qref_measure_in_basis_p, qref_measure_p, qref_namedobs_p, + qref_operator_op, qref_pauli_measure_p, qref_pauli_rot_p, qref_qinst_p, qref_set_basis_state_p, qref_set_state_p, qref_unitary_p, - qref_operator_op, ) from catalyst.jax_primitives import ( counts_p, diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 315003a187..6535a85cfa 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -29,8 +29,10 @@ ) from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp from pennylane.capture.primitives import adjoint_transform_prim as plxpr_adjoint_transform_prim -from pennylane.wires import AbstractQubit from pennylane.pytrees import unflatten +from pennylane.wires import AbstractQubit + +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval # TODO: remove after jax v0.7.2 upgrade # Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between @@ -39,7 +41,6 @@ # once JAX updates to a compatible MLIR version # pylint: disable=ungrouped-imports from catalyst.jax_extras.patches import mock_attributes -from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.jax_primitives import ( AbstractObs, _named_obs_attribute, @@ -827,13 +828,15 @@ def _operator_op_lowering( repack_static_data = {k: unflatten(*v) for k, v in static_data.items()} processed_static_data = get_mlir_attribute_from_pyval(repack_static_data) - param_map = {name: ir.DenseI64ArrayAttr.get([ind]) for ind, name in enumerate(op_cls.dynamic_argnames)} + param_map = { + name: ir.DenseI64ArrayAttr.get([ind]) for ind, name in enumerate(op_cls.dynamic_argnames) + } processed_param_map = get_mlir_attribute_from_pyval(param_map) qubit_map = {} ind = 0 for name, size in zip(op_cls.wire_argnames, wire_lens): - qubit_map[name] = ir.DenseI64ArrayAttr.get(list(range(ind, ind+size))) + qubit_map[name] = ir.DenseI64ArrayAttr.get(list(range(ind, ind + size))) ind += size processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 9260ddcb21..0ae5fc2d25 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -13,12 +13,12 @@ # limitations under the License. """Tests for operator in Catalyst.""" - # RUN: %PYTHON %s | FileCheck %s import numpy as np import pennylane as qp + class NoParams(qp.core.Operator2): def __init__(self, wires): @@ -26,10 +26,10 @@ def __init__(self, wires): @qp.qjit(target="mlir", capture=True) -@qp.qnode(qp.device('null.qubit', wires=2)) +@qp.qnode(qp.device("null.qubit", wires=2)) def c_no_params(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: qref.operator "NoParams"() qubits([[q0]]) + # CHECK: qref.operator "NoParams"() qubits([[q0]]) # CHECK static_data = {} # CHECK param_map = {} qubit_map = {wires = [0]} NoParams(wires=0) @@ -39,21 +39,24 @@ def c_no_params(): # CHECK: qref.operator "NoParams"() qubits([[q1]], [[q2]]) # CHECK: static_data = {} # CHECK: param_map = {} qubit_map = {wires = [0, 1]} - NoParams(wires=(0,1)) + NoParams(wires=(0, 1)) return qp.state() + print(c_no_params.mlir) + class SingleParam(qp.core.Operator2): - dynamic_argnames = ("x", ) + dynamic_argnames = ("x",) def __init__(self, x, wires): super().__init__(x, wires=wires) + @qp.qjit(target="mlir", capture=True) -@qp.qnode(qp.device('null.qubit', wires=3)) -def c_single_param(x : float): +@qp.qnode(qp.device("null.qubit", wires=3)) +def c_single_param(x: float): # CHECK: [[q0:%.+]] = qref.get {{%.+}} @@ -67,10 +70,11 @@ def c_single_param(x : float): # CHECK: qref.operator "SingleParam"({{%.+}}: tensor<4x4xf64>) qubits([[q1]], [[q2]]) # CHECK: static_data = {} # CHECK: param_map = {x = [0]} qubit_map = {wires = [0, 1]} - SingleParam(np.eye(4), (1,2)) + SingleParam(np.eye(4), (1, 2)) return qp.state() + print(c_single_param.mlir) @@ -81,8 +85,9 @@ class CompilableData(qp.core.Operator2): def __init__(self, a, b, thing, wires): super().__init__(a=a, b=b, thing=thing, wires=wires) + @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device('null.qubit', wires=3)) +@qp.qnode(qp.device("null.qubit", wires=3)) def c_compilable(): # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} @@ -90,7 +95,7 @@ def c_compilable(): # CHECK: static_data = {a = true, b = "some string", thing = [1, true, "string"]} # CHECK: param_map = {} qubit_map = {wires = [0, 1]} - CompilableData(True, "some string", (1, True, "string"), wires=(0,1)) + CompilableData(True, "some string", (1, True, "string"), wires=(0, 1)) return qp.state() @@ -105,8 +110,9 @@ class MultipleRegisters(qp.core.Operator2): def __init__(self, reg1, reg2): super().__init__(reg1=reg1, reg2=reg2) + @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device('null.qubit', wires=5)) +@qp.qnode(qp.device("null.qubit", wires=5)) def c_multiple_registers(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} @@ -116,7 +122,7 @@ def c_multiple_registers(): # CHECK: qref.operator "MultipleRegisters"() qubits([[q0]], [[q2]], [[q3]], [[q4]]) # CHECK: static_data = {} # CHECK: param_map = {} qubit_map = {reg1 = [0], reg2 = [1, 2, 3]} - MultipleRegisters(0, (2,3,4)) + MultipleRegisters(0, (2, 3, 4)) return qp.state() @@ -130,17 +136,18 @@ class MultiParams(qp.core.Operator2): # note also having non-standard order with dynamic inputs after wires def __init__(self, wires, a, b, c): super().__init__(wires, a, b, c) - + @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device('null.qubit', wires=1)) +@qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_params(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "MultiParams"({{%.+}}: tensor, {{%.+}}: tensor<4x2x1xf64>, {{%.+}}: tensor<3xi64>) qubits([[q0]]) # CHECK: static_data = {} # CHECK: param_map = {a = [0], b = [1], c = [2]} qubit_map = {wires = [0]} - MultiParams(0, 0.5, c=np.array([1, 2 ,3]), b=np.zeros((4,2,1))) + MultiParams(0, 0.5, c=np.array([1, 2, 3]), b=np.zeros((4, 2, 1))) return qp.state() -print(c_multi_params.mlir) \ No newline at end of file + +print(c_multi_params.mlir) diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py new file mode 100644 index 0000000000..0ab4ac9668 --- /dev/null +++ b/frontend/test/pytest/test_operator.py @@ -0,0 +1,72 @@ +# Copyright 2022-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. + +import pennylane as qp +import pytest + + +class DummyOp(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +def test_error_on_operator_outside_qnode(): + """Test that an error is raised for an Operator2 outside the qnode.""" + + with pytest.raises(ValueError, match="must occur inside a qnode."): + + @qp.qjit(capture=True, target="jaxpr") + def f(): + DummyOp(2) + return 2 + + +def test_hybrid_not_supported_yet(): + """Test that hybrid arguments are not yet supported.""" + + class OperatorArgument(qp.core.Operator2): + + hybrid_argnames = ("op",) + wire_argnames = () + + def __init__(self, op): + super().__init__(op) + + with pytest.raises(NotImplementedError): + + @qp.qjit(capture=True) + @qp.qnode(qp.device("null.qubit", wires=3)) + def c(): + OperatorArgument(DummyOp(0)) + return qp.state() + + +def test_static_argnames(): + """Test that static arguments are not yet supported.""" + + class StaticArgsOp(qp.core.Operator2): + + static_argnames = ("thing",) + + def __init__(self, thing, wires): + super().__init__(thing, wires) + + with pytest.raises(NotImplementedError): + + @qp.qjit(capture=True) + @qp.qnode(qp.device("null.qubit", wires=2)) + def c(): + StaticArgsOp("hello", 0) + return qp.state() diff --git a/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp b/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp index d3475f46c2..689af5bb9c 100644 --- a/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp +++ b/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp @@ -1159,7 +1159,6 @@ void handleGate(IRRewriter &builder, qref::QuantumOperation rGateOp, QubitValueT builder.getDenseI32ArrayAttr({nTargets, nCtrls, nQreg})); // Properties are not handled via the generic attribute fields, so we set them separately. - vGateOp.setOpName(rOperatorOp.getOpName()); vGateOp.setAdjoint(rOperatorOp.getAdjoint()); vGateOp.setUID(rOperatorOp.getUID()); } diff --git a/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp b/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp index 3dbbf0baf1..5c35f70782 100644 --- a/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp +++ b/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp @@ -368,7 +368,6 @@ void handleGate(IRRewriter &builder, quantum::QuantumOperation vGateOp, QubitVal rGateOp->removeAttr("resultSegmentSizes"); // Properties are not handled via the generic attribute fields, so we set them separately. - rGateOp.setOpName(vOperatorOp.getOpName()); rGateOp.setAdjoint(vOperatorOp.getAdjoint()); rGateOp.setUID(vOperatorOp.getUID()); } From 862375eddba2fd28fb2eaec3baa8cf7e7a3c2f0c Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 25 Jun 2026 10:29:48 -0400 Subject: [PATCH 10/48] pylint --- frontend/test/pytest/test_operator.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py index 0ab4ac9668..c637bb0f94 100644 --- a/frontend/test/pytest/test_operator.py +++ b/frontend/test/pytest/test_operator.py @@ -11,7 +11,10 @@ # 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 new Operator2 class. +""" +# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring import pennylane as qp import pytest From 6e4eaf761c9dd6b5d92dc3d0e30e93039b68c13d Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 25 Jun 2026 10:35:01 -0400 Subject: [PATCH 11/48] why didn't black work --- frontend/test/pytest/test_operator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py index c637bb0f94..557df7658d 100644 --- a/frontend/test/pytest/test_operator.py +++ b/frontend/test/pytest/test_operator.py @@ -14,7 +14,8 @@ """ Tests for the new Operator2 class. """ -# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring + +# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring import pennylane as qp import pytest From 78b4f9ee59c4316ea9923dfe23d7017175e8f535 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 25 Jun 2026 13:52:05 -0400 Subject: [PATCH 12/48] more pylint --- frontend/test/lit/test_operator.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 0ae5fc2d25..abd2beb55d 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -13,6 +13,8 @@ # limitations under the License. """Tests for operator in Catalyst.""" +# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring + # RUN: %PYTHON %s | FileCheck %s import numpy as np @@ -143,6 +145,7 @@ def __init__(self, wires, a, b, c): def c_multi_params(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # pylint: disable=line-too-long # CHECK: qref.operator "MultiParams"({{%.+}}: tensor, {{%.+}}: tensor<4x2x1xf64>, {{%.+}}: tensor<3xi64>) qubits([[q0]]) # CHECK: static_data = {} # CHECK: param_map = {a = [0], b = [1], c = [2]} qubit_map = {wires = [0]} From 51639cdcece335f9c0550d5780280fda35b733b4 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 26 Jun 2026 16:47:17 -0400 Subject: [PATCH 13/48] add specialized lowerings --- .../from_plxpr/qref_jax_primitives.py | 219 +++++++++++++++++- 1 file changed, 210 insertions(+), 9 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 6535a85cfa..8701b1646e 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -811,6 +811,14 @@ def _qref_operator_op_abstract_eval(*args, **kwargs): return [] +def _is_custom_op(op_cls, params): + if op_cls.static_argnames or op_cls.hybrid_argnames or op_cls.compilable_argnames: + return False + if op_cls.wire_argnames != ("wires",): + return False + return all(p.shape == () and "float" in p.dtype.name for p in params) + + def _operator_op_lowering( jax_ctx: mlir.LoweringRuleContext, *args, @@ -820,6 +828,10 @@ def _operator_op_lowering( wire_lens, **static_data, ): + ctx = jax_ctx.module_context.context + ctx.allow_unregistered_dialects = True + if op_cls.__name__ in _SPECIAL_LOWERINGS: + return _SPECIAL_LOWERINGS[op_cls.__name__](jax_ctx, *args, op_cls=op_cls, hybrid_lens=hybrid_lens, hybrid_trees=hybrid_trees, wire_lens=wire_lens, **static_data) params = args[: len(op_cls.dynamic_argnames)] qubits = args[len(op_cls.dynamic_argnames) :] @@ -841,23 +853,212 @@ def _operator_op_lowering( processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) - OperatorOp( - op_name=name_attr, - params=params, + ctrl_qubits = [] + ctrl_values = [] + adjoint = False + + if op_cls.__name__ == "PCPhase": + assert len(params) == 2, "PCPhase takes two float parameters" + PCPhaseOp( + theta=extract_scalar(safe_cast_to_f64(params[0], op_cls), op_cls), + dim=extract_scalar(safe_cast_to_f64(params[1], op_cls), op_cls), + qubits=qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, + ) + return () + elif _is_custom_op(op_cls, params): + params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] + CustomOp( + params=params, + qubits=qubits, + gate_name=name_attr, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, + ) + else: + OperatorOp( + op_name=name_attr, + params=params, + qubits=qubits, + forward_args=[], + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, + UID=None, + arr_qubit_indices=[], + param_map=processed_param_map, + static_data=processed_static_data, + qubit_map=processed_qubit_map, + ) + return [] + + +_SPECIAL_LOWERINGS = {} + +def _register_special_lowering(op_name): + def decorator(f): + _SPECIAL_LOWERINGS[op_name] = f + return f + return decorator + +@_register_special_lowering("MultiRZ") +def _multirz_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + theta=extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"), + qubits=args[1:] + MultiRZOp( + theta=theta, qubits=qubits, - forward_args=[], ctrl_qubits=[], ctrl_values=[], adjoint=False, - UID=None, - arr_qubit_indices=[], - param_map=processed_param_map, - static_data=processed_static_data, - qubit_map=processed_qubit_map, ) return [] +@_register_special_lowering("PCPhase") +def _pcphase_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + qubits = args[2:] + PCPhaseOp( + theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), + dim=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), + qubits=qubits, + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + ) + return () + +@_register_special_lowering("GlobalPhase") +def _special_gphase_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + GlobalPhaseOp( + angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"), + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + ) + return () + + +@_register_special_lowering("QubitUnitary") +def _special_unitary_lowering( + jax_ctx: mlir.LoweringRuleContext, + matrix, + *qubits, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + ctrl_qubits = [] + ctrl_values = [] + + for q in qubits: + assert ir.OpaqueType.isinstance(q.type) + assert ir.OpaqueType(q.type).dialect_namespace == "qref" + assert ir.OpaqueType(q.type).data == "bit" + + matrix_type = matrix.type + is_tensor = ir.RankedTensorType.isinstance(matrix_type) + shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None + is_2d_tensor = len(shape) == 2 if is_tensor else False + if not is_2d_tensor: + raise TypeError("QubitUnitary must be a 2 dimensional tensor.") + + possibly_complex_type = ir.RankedTensorType(matrix_type).element_type + is_complex = ir.ComplexType.isinstance(possibly_complex_type) + is_f64_type = False + + if is_complex: + complex_type = ir.ComplexType(possibly_complex_type) + possibly_f64_type = complex_type.element_type + is_f64_type = ir.F64Type.isinstance(possibly_f64_type) + + is_complex_f64_type = is_complex and is_f64_type + if not is_complex_f64_type: + f64_type = ir.F64Type.get() + complex_f64_type = ir.ComplexType.get(f64_type) + tensor_complex_f64_type = ir.RankedTensorType.get(shape, complex_f64_type) + matrix = StableHLOConvertOp(tensor_complex_f64_type, matrix).result + + ctrl_values_i1 = [ + TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values + ] + + QubitUnitaryOp( + matrix=matrix, + qubits=qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values_i1, + adjoint=False, + ) + + return () + +@_register_special_lowering("PauliRot") +def _special_paulirot_lowering( + jax_ctx: mlir.LoweringRuleContext, + angle, + *qubits, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, + pauli_word, +): + pauli_word = unflatten(*pauli_word) + ctrl_qubits = [] + ctrl_values = [] + + for q in qubits: + assert ir.OpaqueType.isinstance(q.type) + assert ir.OpaqueType(q.type).dialect_namespace == "qref" + assert ir.OpaqueType(q.type).data == "bit" + + angle = safe_cast_to_f64(angle, "PauliRot") + angle = extract_scalar(angle, "PauliRot") + assert ir.F64Type.isinstance(angle.type) + + pauli_word = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in pauli_word]) + + ctrl_values_i1 = [ + TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values + ] + + PauliRotOp( + angle=angle, + pauli_product=pauli_word, + qubits=qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values_i1, + adjoint=False, + ) + + return () + # # hermitian observable # From 8513e91058f40b0f67da1a1541c9328b458261df Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Tue, 30 Jun 2026 11:24:35 -0400 Subject: [PATCH 14/48] Apply suggestions from code review Co-authored-by: Christina Lee --- frontend/catalyst/from_plxpr/from_plxpr.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/frontend/catalyst/from_plxpr/from_plxpr.py b/frontend/catalyst/from_plxpr/from_plxpr.py index 86bdc8ac88..f3b0366f57 100644 --- a/frontend/catalyst/from_plxpr/from_plxpr.py +++ b/frontend/catalyst/from_plxpr/from_plxpr.py @@ -27,7 +27,6 @@ from jax.extend.core import ClosedJaxpr, Jaxpr from pennylane.capture import PlxprInterpreter, qnode_prim from pennylane.capture.primitives import transform_prim -from pennylane.core.operator.operator2 import operator_p from pennylane.transforms import decompose as pl_decompose from catalyst.device import extract_backend_info @@ -409,11 +408,6 @@ def _handle_decompose_transform(self, inner_jaxpr, consts, non_const_args, tkwar return next_eval.eval(inner_jaxpr, consts, *non_const_args) -@WorkflowInterpreter.register_primitive(operator_p) -def _error_on_operator(self, *args, op_cls, **kwargs): - raise ValueError(f"Operator {op_cls} must occur inside a qnode.") - - # pylint: disable=too-many-arguments @WorkflowInterpreter.register_primitive(transform_prim) def handle_transform( From ca108c89003daff3f174dc16dad57f3b289dd9a5 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Tue, 30 Jun 2026 11:28:56 -0400 Subject: [PATCH 15/48] styling --- .../from_plxpr/qref_jax_primitives.py | 60 +++++++++++++------ 1 file changed, 41 insertions(+), 19 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index a3863a068d..64e76eea74 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -811,6 +811,14 @@ def _qref_operator_p_abstract_eval(*args, **kwargs): return [] +def _is_custom_op(op_cls, params): + if op_cls.static_argnames or op_cls.hybrid_argnames or op_cls.compilable_argnames: + return False + if op_cls.wire_argnames != ("wires",): + return False + return all(p.shape == () and "float" in p.dtype.name for p in params) + + def _qref_operator_p_lowering( jax_ctx: mlir.LoweringRuleContext, *args, @@ -823,7 +831,15 @@ def _qref_operator_p_lowering( ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True if op_cls.__name__ in _SPECIAL_LOWERINGS: - return _SPECIAL_LOWERINGS[op_cls.__name__](jax_ctx, *args, op_cls=op_cls, hybrid_lens=hybrid_lens, hybrid_trees=hybrid_trees, wire_lens=wire_lens, **static_data) + return _SPECIAL_LOWERINGS[op_cls.__name__]( + jax_ctx, + *args, + op_cls=op_cls, + hybrid_lens=hybrid_lens, + hybrid_trees=hybrid_trees, + wire_lens=wire_lens, + **static_data, + ) params = args[: len(op_cls.dynamic_argnames)] qubits = args[len(op_cls.dynamic_argnames) :] @@ -860,32 +876,35 @@ def _qref_operator_p_lowering( adjoint=adjoint, ) else: - OperatorOp( - op_name=name_attr, - params=params, - qubits=qubits, - qreg=None, - forward_args=[], - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - UID=None, - arr_qubit_indices=[], - param_map=processed_param_map, - static_data=processed_static_data, - qubit_map=processed_qubit_map, - ) + OperatorOp( + op_name=name_attr, + params=params, + qubits=qubits, + qreg=None, + forward_args=[], + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + UID=None, + arr_qubit_indices=[], + param_map=processed_param_map, + static_data=processed_static_data, + qubit_map=processed_qubit_map, + ) return [] _SPECIAL_LOWERINGS = {} + def _register_special_lowering(op_name): def decorator(f): _SPECIAL_LOWERINGS[op_name] = f return f + return decorator + @_register_special_lowering("MultiRZ") def _multirz_lowering( jax_ctx: mlir.LoweringRuleContext, @@ -893,10 +912,10 @@ def _multirz_lowering( op_cls, hybrid_lens, hybrid_trees, - wire_lens, + wire_lens ): - theta=extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"), - qubits=args[1:] + theta = (extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"),) + qubits = args[1:] MultiRZOp( theta=theta, qubits=qubits, @@ -927,6 +946,7 @@ def _pcphase_lowering( ) return () + @_register_special_lowering("GlobalPhase") def _special_gphase_lowering( jax_ctx: mlir.LoweringRuleContext, @@ -1000,6 +1020,7 @@ def _special_unitary_lowering( return () + @_register_special_lowering("PauliRot") def _special_paulirot_lowering( jax_ctx: mlir.LoweringRuleContext, @@ -1041,6 +1062,7 @@ def _special_paulirot_lowering( return () + # # hermitian observable # From 65ec5fbc7d02c84d5853574ad787da5a652584aa Mon Sep 17 00:00:00 2001 From: albi3ro Date: Tue, 30 Jun 2026 13:54:02 -0400 Subject: [PATCH 16/48] move things to their own file --- .../from_plxpr/qref_jax_primitives.py | 265 +-------------- .../from_plxpr/qref_operator2_primitives.py | 319 ++++++++++++++++++ frontend/test/lit/test_operator.py | 81 ++++- 3 files changed, 390 insertions(+), 275 deletions(-) create mode 100644 frontend/catalyst/from_plxpr/qref_operator2_primitives.py diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 64e76eea74..085c4d9bdb 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -29,7 +29,6 @@ ) from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp from pennylane.capture.primitives import adjoint_transform_prim as plxpr_adjoint_transform_prim -from pennylane.pytrees import unflatten from pennylane.wires import AbstractQubit from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval @@ -50,6 +49,8 @@ from catalyst.utils.extra_bindings import FromElementsOp, TensorExtractOp from catalyst.utils.patching import Patcher +from .qref_operator2_primitives import qref_operator_p, _qref_operator_p_lowering + with Patcher( ( _ods_cext, @@ -75,7 +76,6 @@ MeasureOp, MultiRZOp, NamedObsOp, - OperatorOp, PauliRotOp, PCPhaseOp, QubitUnitaryOp, @@ -168,7 +168,7 @@ class MeasurementPlane(Enum): qref_compbasis_p = Primitive("qref_compbasis") qref_namedobs_p = Primitive("qref_namedobs") qref_hermitian_p = Primitive("qref_hermitian") -qref_operator_p = Primitive("qref_operator") + # @@ -803,265 +803,6 @@ def _qref_named_obs_lowering(jax_ctx: mlir.LoweringRuleContext, qubit: ir.Value, return NamedObsOp(result_type, qubit, obsId).results -qref_operator_p.multiple_results = True - - -@qref_operator_p.def_abstract_eval -def _qref_operator_p_abstract_eval(*args, **kwargs): - return [] - - -def _is_custom_op(op_cls, params): - if op_cls.static_argnames or op_cls.hybrid_argnames or op_cls.compilable_argnames: - return False - if op_cls.wire_argnames != ("wires",): - return False - return all(p.shape == () and "float" in p.dtype.name for p in params) - - -def _qref_operator_p_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, - **static_data, -): - ctx = jax_ctx.module_context.context - ctx.allow_unregistered_dialects = True - if op_cls.__name__ in _SPECIAL_LOWERINGS: - return _SPECIAL_LOWERINGS[op_cls.__name__]( - jax_ctx, - *args, - op_cls=op_cls, - hybrid_lens=hybrid_lens, - hybrid_trees=hybrid_trees, - wire_lens=wire_lens, - **static_data, - ) - params = args[: len(op_cls.dynamic_argnames)] - qubits = args[len(op_cls.dynamic_argnames) :] - - name_attr = get_mlir_attribute_from_pyval(op_cls.__name__) - - repack_static_data = {k: unflatten(*v) for k, v in static_data.items()} - processed_static_data = get_mlir_attribute_from_pyval(repack_static_data) - - param_map = { - name: ir.DenseI64ArrayAttr.get([ind]) for ind, name in enumerate(op_cls.dynamic_argnames) - } - processed_param_map = get_mlir_attribute_from_pyval(param_map) - - qubit_map = {} - ind = 0 - for name, size in zip(op_cls.wire_argnames, wire_lens): - qubit_map[name] = ir.DenseI64ArrayAttr.get(list(range(ind, ind + size))) - ind += size - - processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) - - ctrl_qubits = [] - ctrl_values = [] - adjoint = False - - if _is_custom_op(op_cls, params): - params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] - CustomOp( - params=params, - qubits=qubits, - gate_name=name_attr, - ctrl_qubits=ctrl_qubits, - ctrl_values=ctrl_values, - adjoint=adjoint, - ) - else: - OperatorOp( - op_name=name_attr, - params=params, - qubits=qubits, - qreg=None, - forward_args=[], - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - UID=None, - arr_qubit_indices=[], - param_map=processed_param_map, - static_data=processed_static_data, - qubit_map=processed_qubit_map, - ) - return [] - - -_SPECIAL_LOWERINGS = {} - - -def _register_special_lowering(op_name): - def decorator(f): - _SPECIAL_LOWERINGS[op_name] = f - return f - - return decorator - - -@_register_special_lowering("MultiRZ") -def _multirz_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens -): - theta = (extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"),) - qubits = args[1:] - MultiRZOp( - theta=theta, - qubits=qubits, - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - ) - return [] - - -@_register_special_lowering("PCPhase") -def _pcphase_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, -): - qubits = args[2:] - PCPhaseOp( - theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), - dim=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), - qubits=qubits, - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - ) - return () - - -@_register_special_lowering("GlobalPhase") -def _special_gphase_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, -): - GlobalPhaseOp( - angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"), - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, - ) - return () - - -@_register_special_lowering("QubitUnitary") -def _special_unitary_lowering( - jax_ctx: mlir.LoweringRuleContext, - matrix, - *qubits, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, -): - ctrl_qubits = [] - ctrl_values = [] - - for q in qubits: - assert ir.OpaqueType.isinstance(q.type) - assert ir.OpaqueType(q.type).dialect_namespace == "qref" - assert ir.OpaqueType(q.type).data == "bit" - - matrix_type = matrix.type - is_tensor = ir.RankedTensorType.isinstance(matrix_type) - shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None - is_2d_tensor = len(shape) == 2 if is_tensor else False - if not is_2d_tensor: - raise TypeError("QubitUnitary must be a 2 dimensional tensor.") - - possibly_complex_type = ir.RankedTensorType(matrix_type).element_type - is_complex = ir.ComplexType.isinstance(possibly_complex_type) - is_f64_type = False - - if is_complex: - complex_type = ir.ComplexType(possibly_complex_type) - possibly_f64_type = complex_type.element_type - is_f64_type = ir.F64Type.isinstance(possibly_f64_type) - - is_complex_f64_type = is_complex and is_f64_type - if not is_complex_f64_type: - f64_type = ir.F64Type.get() - complex_f64_type = ir.ComplexType.get(f64_type) - tensor_complex_f64_type = ir.RankedTensorType.get(shape, complex_f64_type) - matrix = StableHLOConvertOp(tensor_complex_f64_type, matrix).result - - ctrl_values_i1 = [ - TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values - ] - - QubitUnitaryOp( - matrix=matrix, - qubits=qubits, - ctrl_qubits=ctrl_qubits, - ctrl_values=ctrl_values_i1, - adjoint=False, - ) - - return () - - -@_register_special_lowering("PauliRot") -def _special_paulirot_lowering( - jax_ctx: mlir.LoweringRuleContext, - angle, - *qubits, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, - pauli_word, -): - pauli_word = unflatten(*pauli_word) - ctrl_qubits = [] - ctrl_values = [] - - for q in qubits: - assert ir.OpaqueType.isinstance(q.type) - assert ir.OpaqueType(q.type).dialect_namespace == "qref" - assert ir.OpaqueType(q.type).data == "bit" - - angle = safe_cast_to_f64(angle, "PauliRot") - angle = extract_scalar(angle, "PauliRot") - assert ir.F64Type.isinstance(angle.type) - - pauli_word = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in pauli_word]) - - ctrl_values_i1 = [ - TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values - ] - - PauliRotOp( - angle=angle, - pauli_product=pauli_word, - qubits=qubits, - ctrl_qubits=ctrl_qubits, - ctrl_values=ctrl_values_i1, - adjoint=False, - ) - - return () - # # hermitian observable diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py new file mode 100644 index 0000000000..7f554f2ed4 --- /dev/null +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -0,0 +1,319 @@ +# 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. +"""This module contains JAX-compatible quantum primitives to support the lowering +of quantum operations, measurements, and observables to reference semantics JAXPR. +""" +# pylint: disable=unused-argument +from jax._src.lib.mlir import ir +from jax.extend.core import Primitive +from jax.interpreters import mlir +from jaxlib.mlir._mlir_libs import _mlir as _ods_cext +from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp +from pennylane.pytrees import unflatten + +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval + +# TODO: remove after jax v0.7.2 upgrade +# Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between +# Catalyst's MLIR version and the MLIR version used by JAX. The current JAX version has not +# yet updated to the latest MLIR, causing compatibility issues. This workaround will be removed +# once JAX updates to a compatible MLIR version +# pylint: disable=ungrouped-imports +from catalyst.jax_extras.patches import mock_attributes +from catalyst.jax_primitives import ( + extract_scalar, + safe_cast_to_f64, +) +from catalyst.utils.extra_bindings import TensorExtractOp +from catalyst.utils.patching import Patcher + +with Patcher( + ( + _ods_cext, + "globals", + mock_attributes( + # pylint: disable=c-extension-no-member + _ods_cext.globals, + {"register_traceback_file_exclusion": lambda x: None}, + ), + ), +): + from mlir_quantum.dialects.qref import ( + CustomOp, + GlobalPhaseOp, + MultiRZOp, + OperatorOp, + PauliRotOp, + PCPhaseOp, + QubitUnitaryOp, + ) + + +_SPECIAL_LOWERINGS = {} + + +def _register_special_lowering(op_name): + def decorator(f): + _SPECIAL_LOWERINGS[op_name] = f + return f + + return decorator + +qref_operator_p = Primitive("qref_operator") +qref_operator_p.multiple_results = True + + +@qref_operator_p.def_abstract_eval +def _qref_operator_p_abstract_eval(*args, **kwargs): + return [] + + +def _is_custom_op(op_cls, params): + if op_cls.static_argnames or op_cls.hybrid_argnames or op_cls.compilable_argnames: + return False + if op_cls.wire_argnames != ("wires",): + return False + return all(p.shape == () and "float" in p.dtype.name for p in params) + + +def _qref_operator_p_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + **kwargs, +): + ctx = jax_ctx.module_context.context + ctx.allow_unregistered_dialects = True + if op_cls.__name__ in _SPECIAL_LOWERINGS: + return _SPECIAL_LOWERINGS[op_cls.__name__]( + jax_ctx, + *args, + op_cls=op_cls, + **kwargs + ) + hybrid_lens = kwargs.pop("hybrid_lens") + hybrid_trees = kwargs.pop("hybrid_trees") + wire_lens = kwargs.pop("wire_lens") + params = args[: len(op_cls.dynamic_argnames)] + qubits = args[len(op_cls.dynamic_argnames) :] + + name_attr = get_mlir_attribute_from_pyval(op_cls.__name__) + + repack_static_data = {k: unflatten(*v) for k, v in kwargs.items()} + processed_static_data = get_mlir_attribute_from_pyval(repack_static_data) + + param_map = { + name: ir.DenseI64ArrayAttr.get([ind]) for ind, name in enumerate(op_cls.dynamic_argnames) + } + processed_param_map = get_mlir_attribute_from_pyval(param_map) + + qubit_map = {} + ind = 0 + for name, size in zip(op_cls.wire_argnames, wire_lens): + qubit_map[name] = ir.DenseI64ArrayAttr.get(list(range(ind, ind + size))) + ind += size + + processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) + + ctrl_qubits = [] + ctrl_values = [] + adjoint = False + + if _is_custom_op(op_cls, params): + params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] + CustomOp( + params=params, + qubits=qubits, + gate_name=name_attr, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, + ) + else: + OperatorOp( + op_name=name_attr, + params=params, + qubits=qubits, + qreg=None, + forward_args=[], + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + UID=None, + arr_qubit_indices=[], + param_map=processed_param_map, + static_data=processed_static_data, + qubit_map=processed_qubit_map, + ) + return [] + + + + +@_register_special_lowering("MultiRZ") +def _multirz_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens +): + theta = (extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"),) + qubits = args[1:] + MultiRZOp( + theta=theta, + qubits=qubits, + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + ) + return [] + + +@_register_special_lowering("PCPhase") +def _pcphase_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + qubits = args[2:] + PCPhaseOp( + theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), + dim=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), + qubits=qubits, + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + ) + return () + + +@_register_special_lowering("GlobalPhase") +def _special_gphase_lowering( + jax_ctx: mlir.LoweringRuleContext, + *args, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + GlobalPhaseOp( + angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"), + ctrl_qubits=[], + ctrl_values=[], + adjoint=False, + ) + return () + + +@_register_special_lowering("QubitUnitary") +def _special_unitary_lowering( + jax_ctx: mlir.LoweringRuleContext, + matrix, + *qubits, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, +): + ctrl_qubits = [] + ctrl_values = [] + + for q in qubits: + assert ir.OpaqueType.isinstance(q.type) + assert ir.OpaqueType(q.type).dialect_namespace == "qref" + assert ir.OpaqueType(q.type).data == "bit" + + matrix_type = matrix.type + is_tensor = ir.RankedTensorType.isinstance(matrix_type) + shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None + is_2d_tensor = len(shape) == 2 if is_tensor else False + if not is_2d_tensor: + raise TypeError("QubitUnitary must be a 2 dimensional tensor.") + + possibly_complex_type = ir.RankedTensorType(matrix_type).element_type + is_complex = ir.ComplexType.isinstance(possibly_complex_type) + is_f64_type = False + + if is_complex: + complex_type = ir.ComplexType(possibly_complex_type) + possibly_f64_type = complex_type.element_type + is_f64_type = ir.F64Type.isinstance(possibly_f64_type) + + is_complex_f64_type = is_complex and is_f64_type + if not is_complex_f64_type: + f64_type = ir.F64Type.get() + complex_f64_type = ir.ComplexType.get(f64_type) + tensor_complex_f64_type = ir.RankedTensorType.get(shape, complex_f64_type) + matrix = StableHLOConvertOp(tensor_complex_f64_type, matrix).result + + ctrl_values_i1 = [ + TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values + ] + + QubitUnitaryOp( + matrix=matrix, + qubits=qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values_i1, + adjoint=False, + ) + + return () + + +@_register_special_lowering("PauliRot") +def _special_paulirot_lowering( + jax_ctx: mlir.LoweringRuleContext, + angle, + *qubits, + op_cls, + hybrid_lens, + hybrid_trees, + wire_lens, + pauli_word, +): + pauli_word = unflatten(*pauli_word) + ctrl_qubits = [] + ctrl_values = [] + + for q in qubits: + assert ir.OpaqueType.isinstance(q.type) + assert ir.OpaqueType(q.type).dialect_namespace == "qref" + assert ir.OpaqueType(q.type).data == "bit" + + angle = safe_cast_to_f64(angle, "PauliRot") + angle = extract_scalar(angle, "PauliRot") + assert ir.F64Type.isinstance(angle.type) + + pauli_word = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in pauli_word]) + + ctrl_values_i1 = [ + TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values + ] + + PauliRotOp( + angle=angle, + pauli_product=pauli_word, + qubits=qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values_i1, + adjoint=False, + ) + + return () diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 12bafa1f1a..b2ce026243 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -23,8 +23,11 @@ class NoParams(qp.core.Operator2): - def __init__(self, wires): - super().__init__(wires=wires) + # have to use different wire argnames or will in up CustomOp + wire_argnames = ("reg",) + + def __init__(self, reg): + super().__init__(reg=reg) @qp.qjit(target="mlir", capture=True) @@ -33,27 +36,51 @@ def c_no_params(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "NoParams"() qubits([[q0]]) # CHECK: static_data = {} - # CHECK: param_map = {} qubit_map = {wires = [0]} - NoParams(wires=0) + # CHECK: param_map = {} qubit_map = {reg = [0]} + NoParams(reg=0) # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "NoParams"() qubits([[q1]], [[q2]]) # CHECK: static_data = {} - # CHECK: param_map = {} qubit_map = {wires = [0, 1]} - NoParams(wires=(0, 1)) + # CHECK: param_map = {} qubit_map = {reg = [0, 1]} + NoParams(reg=(0, 1)) return qp.state() print(c_no_params.mlir) +class NoParamsCustomOp(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +@qp.qjit(target="mlir", capture=True) +@qp.qnode(qp.device("null.qubit", wires=2)) +def c_no_params_custom(): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: qref.custom "NoParamsCustomOp"() [[q0]] : !qref.bit + NoParamsCustomOp(wires=0) + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.custom "NoParamsCustomOp"() [[q1]], [[q2]] : !qref.bit, !qref.bit + NoParamsCustomOp(wires=(0, 1)) + return qp.state() + + +print(c_no_params_custom.mlir) + + class SingleParam(qp.core.Operator2): dynamic_argnames = ("x",) + wire_argnames = ("reg",) - def __init__(self, x, wires): - super().__init__(x, wires=wires) + def __init__(self, x, reg): + super().__init__(x, reg=reg) @qp.qjit(target="mlir", capture=True) @@ -64,14 +91,14 @@ def c_single_param(x: float): # CHECK: qref.operator "SingleParam"({{%.+}}: tensor) qubits([[q0]]) # CHECK: static_data = {} - # CHECK: param_map = {x = [0]} qubit_map = {wires = [0]} + # CHECK: param_map = {x = [0]} qubit_map = {reg = [0]} SingleParam(x, 0) # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "SingleParam"({{%.+}}: tensor<4x4xf64>) qubits([[q1]], [[q2]]) # CHECK: static_data = {} - # CHECK: param_map = {x = [0]} qubit_map = {wires = [0, 1]} + # CHECK: param_map = {x = [0]} qubit_map = {reg = [0, 1]} SingleParam(np.eye(4), (1, 2)) return qp.state() @@ -80,6 +107,33 @@ def c_single_param(x: float): print(c_single_param.mlir) +class SingleParamCustomOp(qp.core.Operator2): + + dynamic_argnames = ("x",) + + def __init__(self, x, wires): + super().__init__(x, wires=wires) + + +@qp.qjit(target="mlir", capture=True) +@qp.qnode(qp.device("null.qubit", wires=3)) +def c_single_param_custom(x: float): + + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}: tensor) [[q0]] -> !qref.bit + SingleParamCustomOp(x, 0) + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}: tensor) [[q1]], [[q2]]) -> !qref.bit !qref.bit + SingleParamCustomOp(0.5, (1, 2)) + + return qp.state() + + +print(c_single_param_custom.mlir) + + class CompilableData(qp.core.Operator2): compilable_argnames = ("a", "b", "thing") @@ -134,10 +188,11 @@ def c_multiple_registers(): class MultiParams(qp.core.Operator2): dynamic_argnames = ("a", "b", "c") + wire_argnames = ("reg", ) # note also having non-standard order with dynamic inputs after wires - def __init__(self, wires, a, b, c): - super().__init__(wires, a, b, c) + def __init__(self, reg, a, b, c): + super().__init__(reg, a, b, c) @qp.qjit(capture=True, target="mlir") @@ -148,7 +203,7 @@ def c_multi_params(): # pylint: disable=line-too-long # CHECK: qref.operator "MultiParams"({{%.+}}: tensor, {{%.+}}: tensor<4x2x1xf64>, {{%.+}}: tensor<3xi64>) qubits([[q0]]) # CHECK: static_data = {} - # CHECK: param_map = {a = [0], b = [1], c = [2]} qubit_map = {wires = [0]} + # CHECK: param_map = {a = [0], b = [1], c = [2]} qubit_map = {reg = [0]} MultiParams(0, 0.5, c=np.array([1, 2, 3]), b=np.zeros((4, 2, 1))) return qp.state() From ab0c04dd776f2665f25d64de560103fda6635090 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Tue, 30 Jun 2026 16:26:33 -0400 Subject: [PATCH 17/48] more testing --- .../from_plxpr/qref_operator2_primitives.py | 10 +- frontend/test/lit/test_operator.py | 26 ++- frontend/test/pytest/test_operator.py | 165 +++++++++++++++++- 3 files changed, 193 insertions(+), 8 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 7f554f2ed4..33e39159b8 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -79,12 +79,12 @@ def _qref_operator_p_abstract_eval(*args, **kwargs): return [] -def _is_custom_op(op_cls, params): +def _is_custom_op(op_cls, avals_in): if op_cls.static_argnames or op_cls.hybrid_argnames or op_cls.compilable_argnames: return False if op_cls.wire_argnames != ("wires",): return False - return all(p.shape == () and "float" in p.dtype.name for p in params) + return all(p.shape == () and "float" in p.dtype.name for p in avals_in) def _qref_operator_p_lowering( @@ -130,7 +130,7 @@ def _qref_operator_p_lowering( ctrl_values = [] adjoint = False - if _is_custom_op(op_cls, params): + if _is_custom_op(op_cls, jax_ctx.avals_in[:len(op_cls.dynamic_argnames)]): params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] CustomOp( params=params, @@ -170,7 +170,7 @@ def _multirz_lowering( hybrid_trees, wire_lens ): - theta = (extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ"),) + theta = extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ") qubits = args[1:] MultiRZOp( theta=theta, @@ -194,7 +194,7 @@ def _pcphase_lowering( qubits = args[2:] PCPhaseOp( theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), - dim=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), + dim=extract_scalar(safe_cast_to_f64(args[1], "PCPhase"), "PCPhase"), qubits=qubits, ctrl_qubits=[], ctrl_values=[], diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index b2ce026243..818a059a23 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -120,12 +120,12 @@ def __init__(self, x, wires): def c_single_param_custom(x: float): # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}: tensor) [[q0]] -> !qref.bit + # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}) [[q0]] : !qref.bit SingleParamCustomOp(x, 0) # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} - # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}: tensor) [[q1]], [[q2]]) -> !qref.bit !qref.bit + # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}) [[q1]], [[q2]] : !qref.bit, !qref.bit SingleParamCustomOp(0.5, (1, 2)) return qp.state() @@ -209,3 +209,25 @@ def c_multi_params(): print(c_multi_params.mlir) + +class MultiParamsCustom(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + + # note also having non-standard order with dynamic inputs after wires + def __init__(self, wires, a, b, c): + super().__init__(wires, a, b, c) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=1)) +def c_multi_param_custom(): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + + # pylint: disable=line-too-long + # CHECK: qref.custom "MultiParamsCustom"({{%.+}}, {{%.+}}, {{%.+}}) [[q0]] : !qref.bit + MultiParamsCustom(0, 0.5, c=0.7, b=2.4) + return qp.state() + + +print(c_multi_param_custom.mlir) diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py index 6558f5972d..f1d0da0577 100644 --- a/frontend/test/pytest/test_operator.py +++ b/frontend/test/pytest/test_operator.py @@ -18,13 +18,176 @@ # pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring import pennylane as qp import pytest - +from jax import numpy as jnp class DummyOp(qp.core.Operator2): def __init__(self, wires): super().__init__(wires=wires) +class PauliX(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + +class RX(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + +class CRX(qp.core.Operator2): + + dynamic_argnames = ("phi", ) + + def __init__(self, phi, wires): + super().__init__(phi, wires=wires) + +class MultiRZ(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + + +class Hadamard(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +class PauliRot(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("pauli_word", ) + + def __init__(self, phi, pauli_word, wires): + super().__init__(phi, pauli_word, wires) + + + +class GlobalPhase(qp.core.Operator2): + + dynamic_argnames = ("phi", ) + wire_argnames = () + + def __init__(self, phi): + super().__init__(phi=phi) + +class QubitUnitary(qp.core.Operator2): + + dynamic_argnames = ("matrix", ) + + def __init__(self, matrix, wires): + super().__init__(matrix, wires) + +class PCPhase(qp.core.Operator2): + + dynamic_argnames = ("phi", "dim") + + def __init__(self, phi, dim, wires): + super().__init__(phi, dim, wires) + + +class TestOperator2Execution: + + def test_custom_op_supported(self): + """Test that Operator2 versions of core ops are supported and can be executed.""" + + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=3)) + def c(x): + PauliX(0) + RX(x, 1) + CRX(2*x, (0,2)) + return qp.expval(qp.Z(0)), qp.expval(qp.Z(1)), qp.expval(qp.Z(2)) + + res1, res2, res3 = c(0.5) + + assert qp.math.allclose(res1, -1) + assert qp.math.allclose(res2, jnp.cos(0.5)) + assert qp.math.allclose(res3, jnp.cos(1.0)) + + def test_MultiRZ(self): + """Test that MultiRZ can be executed.""" + + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=3)) + def c(x): + Hadamard(0) + Hadamard(1) + # skip on 2 for comparison + MultiRZ(x, (0,1,2)) + return qp.expval(qp.X(0)), qp.expval(qp.X(1)), qp.expval(qp.X(2)) + + r1, r2, r3 = c(0.5) + assert qp.math.allclose(r1, jnp.cos(0.5)) + assert qp.math.allclose(r2, jnp.cos(0.5)) + assert qp.math.allclose(r3, 0) + + def test_paulirot(self): + """Test that PauliRot can be executed.""" + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=3)) + def c(x): + Hadamard(2) + PauliRot(x, "XYZ", (0,1,2)) + return qp.expval(qp.Z(0)), qp.expval(qp.Z(1)), qp.expval(qp.X(2)) + + r1, r2, r3 = c(1.2) + assert qp.math.allclose(r1, jnp.cos(1.2)) + assert qp.math.allclose(r2, jnp.cos(1.2)) + assert qp.math.allclose(r3, jnp.cos(1.2)) + + def test_globalphase(self): + """Test that global phase can be executed.""" + + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=1)) + def c(x): + GlobalPhase(x) + return qp.state() + + state = c(0.5) + assert qp.math.allclose(state, jnp.exp(0.5 * -1j) * jnp.array([1, 0])) + + def test_QubitUnitary(self): + """Test that QubitUnitary can be executed.""" + + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=3)) + def c(): + QubitUnitary(jnp.array([[0,1],[1,0]]), 0) + QubitUnitary(qp.CNOT.compute_matrix(), (0,1)) + return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) + + r1, r2 = c() + assert qp.math.allclose(r1, -1) + assert qp.math.allclose(r2, -1) + + def test_PCPhase(self): + """Test that PCPhase can be executed.""" + + @qp.qjit(capture=True) + @qp.qnode(qp.device('lightning.qubit', wires=3)) + def c(x, dim): + Hadamard(0) + Hadamard(1) + PCPhase(x, dim, (0,1)) + return qp.state() + + state2 = c(0.5, 2) + plus = jnp.exp(0.5j)/jnp.sqrt(2) + minus = jnp.exp(-0.5j)/jnp.sqrt(2) + expected2 = jnp.array([plus, plus, minus, minus]) + assert qp.math.allclose(state2, expected2) + + state3 = c(0.5, 3) + expected3 = jnp.array([plus, plus, plus, minus]) + assert qp.math.allclose(state3, expected3) + def test_hybrid_not_supported_yet(): """Test that hybrid arguments are not yet supported.""" From 277c4852eccef79bfcf958903ef958007653a07b Mon Sep 17 00:00:00 2001 From: albi3ro Date: Tue, 30 Jun 2026 17:24:34 -0400 Subject: [PATCH 18/48] more tests --- frontend/test/lit/test_operator.py | 20 ++++++++++++++++++++ frontend/test/pytest/test_operator.py | 2 +- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 818a059a23..27d5ab8b44 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -231,3 +231,23 @@ def c_multi_param_custom(): print(c_multi_param_custom.mlir) + +class MultiRZ(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device('null.qubit', wires=2)) +def c(x: float): + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + + # CHECK: qref.multirz({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit + MultiRZ(x, (0,1,2)) + return qp.state() + +print(c.mlir) \ No newline at end of file diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py index f1d0da0577..69e0cf0472 100644 --- a/frontend/test/pytest/test_operator.py +++ b/frontend/test/pytest/test_operator.py @@ -157,7 +157,7 @@ def test_QubitUnitary(self): """Test that QubitUnitary can be executed.""" @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=3)) + @qp.qjitqnode(qp.device('lightning.qubit', wires=3)) def c(): QubitUnitary(jnp.array([[0,1],[1,0]]), 0) QubitUnitary(qp.CNOT.compute_matrix(), (0,1)) From a42a5c63b4cd2fb1f777907423bdbccfb9bbf13a Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 11:29:57 -0400 Subject: [PATCH 19/48] more testing, changelog, formatting --- doc/releases/changelog-dev.md | 7 +- .../from_plxpr/qref_jax_primitives.py | 6 +- .../from_plxpr/qref_operator2_primitives.py | 51 ++------ frontend/test/lit/test_operator.py | 121 +++++++++++++++++- frontend/test/pytest/test_operator.py | 59 +++++---- 5 files changed, 167 insertions(+), 77 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index f5b0045e82..d1959e5e1e 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -6,7 +6,12 @@

Improvements 🛠

* The new `pennylane.core.Operator2` can now be lowered to MLIR with program capture for operators - without non-lowerable arguments. + without non-lowerable arguments. `Operator2` classes with scalar float dynamic inputs, a single + register of wires called `wires`, and no static/ compilable metadata are lowered to `CustomOp` and + can be compiled and executed using the existing infrastructure. `Operator2` versions of + `GlobalPhase`, `MultiRZ`,`PauliRot`, `PCPhase`, and `QubitUnitary` can also be lowered to their + specialized representations, compiled, and executed. + [(#2979)](https://github.com/PennyLaneAI/catalyst/pull/2979) [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969/) * The `ResourceAnalysis` pass now reports each loop body and each subroutine as its own entry diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index 085c4d9bdb..903942dfe3 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -31,8 +31,6 @@ from pennylane.capture.primitives import adjoint_transform_prim as plxpr_adjoint_transform_prim from pennylane.wires import AbstractQubit -from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval - # TODO: remove after jax v0.7.2 upgrade # Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between # Catalyst's MLIR version and the MLIR version used by JAX. The current JAX version has not @@ -49,7 +47,7 @@ from catalyst.utils.extra_bindings import FromElementsOp, TensorExtractOp from catalyst.utils.patching import Patcher -from .qref_operator2_primitives import qref_operator_p, _qref_operator_p_lowering +from .qref_operator2_primitives import _qref_operator_p_lowering, qref_operator_p with Patcher( ( @@ -170,7 +168,6 @@ class MeasurementPlane(Enum): qref_hermitian_p = Primitive("qref_hermitian") - # # qref_alloc_p # @@ -803,7 +800,6 @@ def _qref_named_obs_lowering(jax_ctx: mlir.LoweringRuleContext, qubit: ir.Value, return NamedObsOp(result_type, qubit, obsId).results - # # hermitian observable # diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 33e39159b8..10c362e0d1 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -14,6 +14,7 @@ """This module contains JAX-compatible quantum primitives to support the lowering of quantum operations, measurements, and observables to reference semantics JAXPR. """ + # pylint: disable=unused-argument from jax._src.lib.mlir import ir from jax.extend.core import Primitive @@ -70,6 +71,7 @@ def decorator(f): return decorator + qref_operator_p = Primitive("qref_operator") qref_operator_p.multiple_results = True @@ -96,14 +98,10 @@ def _qref_operator_p_lowering( ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True if op_cls.__name__ in _SPECIAL_LOWERINGS: - return _SPECIAL_LOWERINGS[op_cls.__name__]( - jax_ctx, - *args, - op_cls=op_cls, - **kwargs - ) - hybrid_lens = kwargs.pop("hybrid_lens") - hybrid_trees = kwargs.pop("hybrid_trees") + return _SPECIAL_LOWERINGS[op_cls.__name__](jax_ctx, *args, op_cls=op_cls, **kwargs) + # will be used in future improvements + hybrid_lens = kwargs.pop("hybrid_lens") # pylint: disable=unused-variable + hybrid_trees = kwargs.pop("hybrid_trees") # pylint: disable=unused-variable wire_lens = kwargs.pop("wire_lens") params = args[: len(op_cls.dynamic_argnames)] qubits = args[len(op_cls.dynamic_argnames) :] @@ -130,7 +128,7 @@ def _qref_operator_p_lowering( ctrl_values = [] adjoint = False - if _is_custom_op(op_cls, jax_ctx.avals_in[:len(op_cls.dynamic_argnames)]): + if _is_custom_op(op_cls, jax_ctx.avals_in[: len(op_cls.dynamic_argnames)]): params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] CustomOp( params=params, @@ -159,17 +157,8 @@ def _qref_operator_p_lowering( return [] - - @_register_special_lowering("MultiRZ") -def _multirz_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens -): +def _multirz_lowering(jax_ctx: mlir.LoweringRuleContext, *args, **_): theta = extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ") qubits = args[1:] MultiRZOp( @@ -186,10 +175,7 @@ def _multirz_lowering( def _pcphase_lowering( jax_ctx: mlir.LoweringRuleContext, *args, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, + **_, ): qubits = args[2:] PCPhaseOp( @@ -208,9 +194,7 @@ def _special_gphase_lowering( jax_ctx: mlir.LoweringRuleContext, *args, op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, + **_, ): GlobalPhaseOp( angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"), @@ -222,15 +206,7 @@ def _special_gphase_lowering( @_register_special_lowering("QubitUnitary") -def _special_unitary_lowering( - jax_ctx: mlir.LoweringRuleContext, - matrix, - *qubits, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, -): +def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits, **_): ctrl_qubits = [] ctrl_values = [] @@ -282,11 +258,8 @@ def _special_paulirot_lowering( jax_ctx: mlir.LoweringRuleContext, angle, *qubits, - op_cls, - hybrid_lens, - hybrid_trees, - wire_lens, pauli_word, + **_, ): pauli_word = unflatten(*pauli_word) ctrl_qubits = [] diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 27d5ab8b44..d81fc9caa9 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -50,6 +50,7 @@ def c_no_params(): print(c_no_params.mlir) + class NoParamsCustomOp(qp.core.Operator2): def __init__(self, wires): @@ -73,7 +74,6 @@ def c_no_params_custom(): print(c_no_params_custom.mlir) - class SingleParam(qp.core.Operator2): dynamic_argnames = ("x",) @@ -188,7 +188,7 @@ def c_multiple_registers(): class MultiParams(qp.core.Operator2): dynamic_argnames = ("a", "b", "c") - wire_argnames = ("reg", ) + wire_argnames = ("reg",) # note also having non-standard order with dynamic inputs after wires def __init__(self, reg, a, b, c): @@ -210,6 +210,7 @@ def c_multi_params(): print(c_multi_params.mlir) + class MultiParamsCustom(qp.core.Operator2): dynamic_argnames = ("a", "b", "c") @@ -232,6 +233,7 @@ def c_multi_param_custom(): print(c_multi_param_custom.mlir) + class MultiRZ(qp.core.Operator2): dynamic_argnames = ("phi",) @@ -239,15 +241,122 @@ class MultiRZ(qp.core.Operator2): def __init__(self, phi, wires): super().__init__(phi, wires) + @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device('null.qubit', wires=2)) -def c(x: float): +@qp.qnode(qp.device("null.qubit", wires=2)) +def circuit(x: float): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: qref.multirz({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit - MultiRZ(x, (0,1,2)) + MultiRZ(x, (0, 1, 2)) + return qp.state() + + +print(circuit.mlir) + + +class PauliRot(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("pauli_word",) + + def __init__(self, phi, pauli_word, wires): + super().__init__(phi, pauli_word, wires) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=3)) +def circuit(x: float): + + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + + # CHECK: qref.paulirot ["X", "Y", "Z"]({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit + PauliRot(x, "XYZ", (0, 1, 2)) + + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q5:%.+]] = qref.get {{%.+}} + + # CHECK: qref.paulirot ["Y", "Z", "X"]({{%.+}}) [[q3]], [[q4]], [[q5]] : !qref.bit, !qref.bit, !qref.bit + PauliRot(x, "YZX", (0, 1, 2)) + + return qp.probs(wires=(0, 1, 2)) + + +print(circuit.mlir) + + +class GlobalPhase(qp.core.Operator2): + + dynamic_argnames = ("phi",) + wire_argnames = () + + def __init__(self, phi): + super().__init__(phi=phi) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=3)) +def circuit(x: float): + + # CHECK: qref.gphase({{%.+}}) + GlobalPhase(x) return qp.state() -print(c.mlir) \ No newline at end of file + +print(circuit.mlir) + + +class QubitUnitary(qp.core.Operator2): + + dynamic_argnames = ("matrix",) + + def __init__(self, matrix, wires): + super().__init__(matrix, wires) + + +@qp.qjit(capture=True) +@qp.qnode(qp.device("lightning.qubit", wires=3)) +def c(): + + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: qref.unitary({{%.+}}. : tensor<2x2xcomplex>) [[q0]] : !qref.bit + + QubitUnitary(np.array([[0, 1], [1, 0]]), 0) + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.unitary({{%.+}}. : tensor<4x4xcomplex>) [[q1]], [[q2]] : !qref.bit, !qref.bit + + QubitUnitary(qp.CNOT.compute_matrix(), (0, 1)) + return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) + + +print(circuit.mlir) + + +class PCPhase(qp.core.Operator2): + + dynamic_argnames = ("phi", "dim") + + def __init__(self, phi, dim, wires): + super().__init__(phi, dim, wires) + + +@qp.qjit(capture=True) +@qp.qnode(qp.device("lightning.qubit", wires=2)) +def c(x: float, dim: int): + + # CHECK: [[q1:%.+]] = qref.get {{%.+}} + # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: qref.pcphase({{%.+}}, {{%.+}}) [[q1]], [[q2]] : !qref.bit, !qref.bit + + PCPhase(x, dim, (0, 1)) + return qp.state() + + +print(circuit.mlir) diff --git a/frontend/test/pytest/test_operator.py b/frontend/test/pytest/test_operator.py index 69e0cf0472..4af23fc69c 100644 --- a/frontend/test/pytest/test_operator.py +++ b/frontend/test/pytest/test_operator.py @@ -20,16 +20,19 @@ import pytest from jax import numpy as jnp + class DummyOp(qp.core.Operator2): def __init__(self, wires): super().__init__(wires=wires) + class PauliX(qp.core.Operator2): def __init__(self, wires): super().__init__(wires=wires) + class RX(qp.core.Operator2): dynamic_argnames = ("phi",) @@ -37,19 +40,13 @@ class RX(qp.core.Operator2): def __init__(self, phi, wires): super().__init__(phi, wires) -class CRX(qp.core.Operator2): - - dynamic_argnames = ("phi", ) - - def __init__(self, phi, wires): - super().__init__(phi, wires=wires) -class MultiRZ(qp.core.Operator2): +class CRX(qp.core.Operator2): dynamic_argnames = ("phi",) def __init__(self, phi, wires): - super().__init__(phi, wires) + super().__init__(phi, wires=wires) class Hadamard(qp.core.Operator2): @@ -58,31 +55,40 @@ def __init__(self, wires): super().__init__(wires=wires) +class MultiRZ(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + + class PauliRot(qp.core.Operator2): dynamic_argnames = ("phi",) - compilable_argnames = ("pauli_word", ) + compilable_argnames = ("pauli_word",) def __init__(self, phi, pauli_word, wires): super().__init__(phi, pauli_word, wires) - class GlobalPhase(qp.core.Operator2): - dynamic_argnames = ("phi", ) + dynamic_argnames = ("phi",) wire_argnames = () def __init__(self, phi): super().__init__(phi=phi) + class QubitUnitary(qp.core.Operator2): - dynamic_argnames = ("matrix", ) + dynamic_argnames = ("matrix",) def __init__(self, matrix, wires): super().__init__(matrix, wires) + class PCPhase(qp.core.Operator2): dynamic_argnames = ("phi", "dim") @@ -97,11 +103,11 @@ def test_custom_op_supported(self): """Test that Operator2 versions of core ops are supported and can be executed.""" @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=3)) + @qp.qnode(qp.device("lightning.qubit", wires=3)) def c(x): PauliX(0) RX(x, 1) - CRX(2*x, (0,2)) + CRX(2 * x, (0, 2)) return qp.expval(qp.Z(0)), qp.expval(qp.Z(1)), qp.expval(qp.Z(2)) res1, res2, res3 = c(0.5) @@ -114,12 +120,12 @@ def test_MultiRZ(self): """Test that MultiRZ can be executed.""" @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=3)) + @qp.qnode(qp.device("lightning.qubit", wires=3)) def c(x): Hadamard(0) Hadamard(1) # skip on 2 for comparison - MultiRZ(x, (0,1,2)) + MultiRZ(x, (0, 1, 2)) return qp.expval(qp.X(0)), qp.expval(qp.X(1)), qp.expval(qp.X(2)) r1, r2, r3 = c(0.5) @@ -129,11 +135,12 @@ def c(x): def test_paulirot(self): """Test that PauliRot can be executed.""" + @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=3)) + @qp.qnode(qp.device("lightning.qubit", wires=3)) def c(x): Hadamard(2) - PauliRot(x, "XYZ", (0,1,2)) + PauliRot(x, "XYZ", (0, 1, 2)) return qp.expval(qp.Z(0)), qp.expval(qp.Z(1)), qp.expval(qp.X(2)) r1, r2, r3 = c(1.2) @@ -145,7 +152,7 @@ def test_globalphase(self): """Test that global phase can be executed.""" @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=1)) + @qp.qnode(qp.device("lightning.qubit", wires=1)) def c(x): GlobalPhase(x) return qp.state() @@ -157,10 +164,10 @@ def test_QubitUnitary(self): """Test that QubitUnitary can be executed.""" @qp.qjit(capture=True) - @qp.qjitqnode(qp.device('lightning.qubit', wires=3)) + @qp.qnode(qp.device("lightning.qubit", wires=3)) def c(): - QubitUnitary(jnp.array([[0,1],[1,0]]), 0) - QubitUnitary(qp.CNOT.compute_matrix(), (0,1)) + QubitUnitary(jnp.array([[0, 1], [1, 0]]), 0) + QubitUnitary(qp.CNOT.compute_matrix(), (0, 1)) return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) r1, r2 = c() @@ -171,16 +178,16 @@ def test_PCPhase(self): """Test that PCPhase can be executed.""" @qp.qjit(capture=True) - @qp.qnode(qp.device('lightning.qubit', wires=3)) + @qp.qnode(qp.device("lightning.qubit", wires=2)) def c(x, dim): Hadamard(0) Hadamard(1) - PCPhase(x, dim, (0,1)) + PCPhase(x, dim, (0, 1)) return qp.state() state2 = c(0.5, 2) - plus = jnp.exp(0.5j)/jnp.sqrt(2) - minus = jnp.exp(-0.5j)/jnp.sqrt(2) + plus = jnp.exp(0.5j) / 2 + minus = jnp.exp(-0.5j) / 2 expected2 = jnp.array([plus, plus, minus, minus]) assert qp.math.allclose(state2, expected2) From bb0728a75519bf7dec62ca400f5b53598dcd8ed1 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 11:37:10 -0400 Subject: [PATCH 20/48] codefactor --- frontend/test/lit/test_operator.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index d81fc9caa9..5d8a919746 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -13,7 +13,7 @@ # limitations under the License. """Tests for operator in Catalyst.""" -# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring +# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring, line-too-long # RUN: %PYTHON %s | FileCheck %s @@ -268,7 +268,7 @@ def __init__(self, phi, pauli_word, wires): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) -def circuit(x: float): +def circuit_paulirot(x: float): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} @@ -287,7 +287,7 @@ def circuit(x: float): return qp.probs(wires=(0, 1, 2)) -print(circuit.mlir) +print(circuit_paulirot.mlir) class GlobalPhase(qp.core.Operator2): @@ -301,14 +301,14 @@ def __init__(self, phi): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) -def circuit(x: float): +def circuit_gphase(x: float): # CHECK: qref.gphase({{%.+}}) GlobalPhase(x) return qp.state() -print(circuit.mlir) +print(circuit_gphase.mlir) class QubitUnitary(qp.core.Operator2): @@ -321,7 +321,7 @@ def __init__(self, matrix, wires): @qp.qjit(capture=True) @qp.qnode(qp.device("lightning.qubit", wires=3)) -def c(): +def circuit_qubitunitary(): # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.unitary({{%.+}}. : tensor<2x2xcomplex>) [[q0]] : !qref.bit @@ -336,7 +336,7 @@ def c(): return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) -print(circuit.mlir) +print(circuit_qubitunitary.mlir) class PCPhase(qp.core.Operator2): @@ -349,7 +349,7 @@ def __init__(self, phi, dim, wires): @qp.qjit(capture=True) @qp.qnode(qp.device("lightning.qubit", wires=2)) -def c(x: float, dim: int): +def c_pcphase(x: float, dim: int): # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} @@ -359,4 +359,4 @@ def c(x: float, dim: int): return qp.state() -print(circuit.mlir) +print(c_pcphase.mlir) From 6a40341f013bb71fb10908af93ca95392b8a5efd Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 14:14:09 -0400 Subject: [PATCH 21/48] I HATE LITgit add *git add *git add *git add *git add *git add *git add *git add * --- frontend/test/lit/test_operator.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 5d8a919746..f01236e002 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -33,6 +33,8 @@ def __init__(self, reg): @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=2)) def c_no_params(): + # CHECK-LABEL: func.func public @c_no_params + # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "NoParams"() qubits([[q0]]) # CHECK: static_data = {} @@ -60,6 +62,8 @@ def __init__(self, wires): @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=2)) def c_no_params_custom(): + # CHECK-LABEL: func.func public @c_no_params_custom + # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.custom "NoParamsCustomOp"() [[q0]] : !qref.bit NoParamsCustomOp(wires=0) @@ -86,6 +90,7 @@ def __init__(self, x, reg): @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=3)) def c_single_param(x: float): + # CHECK-LABEL: func.func public @c_single_param # CHECK: [[q0:%.+]] = qref.get {{%.+}} @@ -118,6 +123,7 @@ def __init__(self, x, wires): @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=3)) def c_single_param_custom(x: float): + # CHECK-LABEL: func.func public @c_single_param_custom # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}) [[q0]] : !qref.bit @@ -145,6 +151,8 @@ def __init__(self, a, b, thing, wires): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def c_compilable(): + # CHECK-LABEL: func.func public @c_compilable + # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: qref.operator "CompilableData"() qubits([[q1]], [[q2]]) @@ -170,6 +178,8 @@ def __init__(self, reg1, reg2): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=5)) def c_multiple_registers(): + # CHECK-LABEL: func.func public @c_multiple_registers + # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: [[q3:%.+]] = qref.get {{%.+}} @@ -198,6 +208,8 @@ def __init__(self, reg, a, b, c): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_params(): + # CHECK-LABEL: func.func public @c_multi_params + # CHECK: [[q0:%.+]] = qref.get {{%.+}} # pylint: disable=line-too-long @@ -223,6 +235,7 @@ def __init__(self, wires, a, b, c): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_param_custom(): + # CHECK-LABEL: func.func public @circuit_multi_param_custom # CHECK: [[q0:%.+]] = qref.get {{%.+}} # pylint: disable=line-too-long @@ -269,6 +282,7 @@ def __init__(self, phi, pauli_word, wires): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_paulirot(x: float): + # CHECK-LABEL: func.func public @circuit_paulirot # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} @@ -302,6 +316,7 @@ def __init__(self, phi): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_gphase(x: float): + # CHECK-LABEL: func.func public @circuit_gphase # CHECK: qref.gphase({{%.+}}) GlobalPhase(x) @@ -322,15 +337,16 @@ def __init__(self, matrix, wires): @qp.qjit(capture=True) @qp.qnode(qp.device("lightning.qubit", wires=3)) def circuit_qubitunitary(): + # CHECK-LABEL: func.func public @circuit_qubitunitary # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: qref.unitary({{%.+}}. : tensor<2x2xcomplex>) [[q0]] : !qref.bit + # CHECK: qref.unitary({{%.+}} : tensor<2x2xcomplex>) [[q0]] : !qref.bit QubitUnitary(np.array([[0, 1], [1, 0]]), 0) # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} - # CHECK: qref.unitary({{%.+}}. : tensor<4x4xcomplex>) [[q1]], [[q2]] : !qref.bit, !qref.bit + # CHECK: qref.unitary({{%.+}} : tensor<4x4xcomplex>) [[q1]], [[q2]] : !qref.bit, !qref.bit QubitUnitary(qp.CNOT.compute_matrix(), (0, 1)) return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) @@ -350,10 +366,11 @@ def __init__(self, phi, dim, wires): @qp.qjit(capture=True) @qp.qnode(qp.device("lightning.qubit", wires=2)) def c_pcphase(x: float, dim: int): + # CHECK-LABEL: func.func public @c_pcphase - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} - # CHECK: qref.pcphase({{%.+}}, {{%.+}}) [[q1]], [[q2]] : !qref.bit, !qref.bit + # CHECK: [[q11:%.+]] = qref.get {{%.+}} + # CHECK: [[q12:%.+]] = qref.get {{%.+}} + # CHECK: qref.pcphase({{%.+}}, {{%.+}}) [[q11]], [[q12]] : !qref.bit, !qref.bit PCPhase(x, dim, (0, 1)) return qp.state() From 32c3a8115b1d9e0053911b17131a75c60d8beab5 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Thu, 2 Jul 2026 14:38:11 -0400 Subject: [PATCH 22/48] Apply suggestions from code review Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- frontend/catalyst/from_plxpr/qref_operator2_primitives.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 10c362e0d1..d41cfad153 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """This module contains JAX-compatible quantum primitives to support the lowering -of quantum operations, measurements, and observables to reference semantics JAXPR. +of quantum operations to reference semantics JAXPR. """ # pylint: disable=unused-argument From 2903e03e45500cf1ba4be67ccef1ce2ee7e2f984 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 14:51:25 -0400 Subject: [PATCH 23/48] add general validation that wires are actual wires in lowering --- .../from_plxpr/qref_operator2_primitives.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index d41cfad153..76f937b8c4 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -89,12 +89,22 @@ def _is_custom_op(op_cls, avals_in): return all(p.shape == () and "float" in p.dtype.name for p in avals_in) +def _general_validation(*args, op_cls, wire_lens, **kwargs): + num_normal_wires = sum(wire_lens) + wires = args[len(op_cls.dynamic_argnames): (len(op_cls.dynamic_argnames) + num_normal_wires)] + for w in wires: + assert ir.OpaqueType.isinstance(w.type) + assert ir.OpaqueType(w.type).dialect_namespace == "qref" + assert ir.OpaqueType(w.type).data == "bit" + + def _qref_operator_p_lowering( jax_ctx: mlir.LoweringRuleContext, *args, op_cls, **kwargs, ): + _general_validation(*args, op_cls, **kwargs) ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True if op_cls.__name__ in _SPECIAL_LOWERINGS: @@ -210,11 +220,6 @@ def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits ctrl_qubits = [] ctrl_values = [] - for q in qubits: - assert ir.OpaqueType.isinstance(q.type) - assert ir.OpaqueType(q.type).dialect_namespace == "qref" - assert ir.OpaqueType(q.type).data == "bit" - matrix_type = matrix.type is_tensor = ir.RankedTensorType.isinstance(matrix_type) shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None @@ -265,11 +270,6 @@ def _special_paulirot_lowering( ctrl_qubits = [] ctrl_values = [] - for q in qubits: - assert ir.OpaqueType.isinstance(q.type) - assert ir.OpaqueType(q.type).dialect_namespace == "qref" - assert ir.OpaqueType(q.type).data == "bit" - angle = safe_cast_to_f64(angle, "PauliRot") angle = extract_scalar(angle, "PauliRot") assert ir.F64Type.isinstance(angle.type) From 7fb66dc8493fd69b5167c73f36063410fd34b63f Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 15:11:10 -0400 Subject: [PATCH 24/48] stupid mistake --- frontend/catalyst/from_plxpr/qref_operator2_primitives.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 76f937b8c4..2b5e0697db 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -91,7 +91,7 @@ def _is_custom_op(op_cls, avals_in): def _general_validation(*args, op_cls, wire_lens, **kwargs): num_normal_wires = sum(wire_lens) - wires = args[len(op_cls.dynamic_argnames): (len(op_cls.dynamic_argnames) + num_normal_wires)] + wires = args[len(op_cls.dynamic_argnames) : (len(op_cls.dynamic_argnames) + num_normal_wires)] for w in wires: assert ir.OpaqueType.isinstance(w.type) assert ir.OpaqueType(w.type).dialect_namespace == "qref" @@ -104,7 +104,7 @@ def _qref_operator_p_lowering( op_cls, **kwargs, ): - _general_validation(*args, op_cls, **kwargs) + _general_validation(*args, op_cls=op_cls, **kwargs) ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True if op_cls.__name__ in _SPECIAL_LOWERINGS: From fdd290a7789c337a847830898e1dd362b46efc96 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Thu, 2 Jul 2026 15:15:49 -0400 Subject: [PATCH 25/48] Apply suggestion from @albi3ro --- doc/releases/changelog-dev.md | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index d1959e5e1e..ae615044ad 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -6,11 +6,8 @@

Improvements 🛠

* The new `pennylane.core.Operator2` can now be lowered to MLIR with program capture for operators - without non-lowerable arguments. `Operator2` classes with scalar float dynamic inputs, a single - register of wires called `wires`, and no static/ compilable metadata are lowered to `CustomOp` and - can be compiled and executed using the existing infrastructure. `Operator2` versions of - `GlobalPhase`, `MultiRZ`,`PauliRot`, `PCPhase`, and `QubitUnitary` can also be lowered to their - specialized representations, compiled, and executed. + without non-lowerable arguments. `Operator2` classes are now lowered to specialized operations + where applicable, unlocking compilation and execution for these cases. [(#2979)](https://github.com/PennyLaneAI/catalyst/pull/2979) [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969/) From ab5484de66127a89f2571eac49b68b7e658a9bed Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 26 Jun 2026 17:44:51 -0400 Subject: [PATCH 26/48] Add support for lowering adjoint and control with Operator2 --- .../catalyst/from_plxpr/qfunc_interpreter.py | 35 +++++++++++++++---- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index 0d51ee2705..b7ed5c968c 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -210,16 +210,24 @@ def _check_measurement_with_dynamic_allocation(self, measurement): """Check some constraints regarding dynamic allocation.""" if self.has_dynamic_allocation: if len(measurement.wires) == 0 and not isinstance(measurement, qp.measurements.StateMP): - raise CompileError(textwrap.dedent(""" + raise CompileError( + textwrap.dedent( + """ Terminal measurements must take in an explicit list of wires when dynamically allocated wires are present in the program. - """)) + """ + ) + ) if any(is_abstract_qubit(w) for w in measurement.wires): - raise CompileError(textwrap.dedent(""" + raise CompileError( + textwrap.dedent( + """ Terminal measurements cannot take in dynamically allocated wires since they must be temporary. - """)) + """ + ) + ) # pylint: disable=too-many-branches def interpret_measurement(self, measurement): @@ -293,23 +301,38 @@ def __call__(self, jaxpr, *args): @PLxPRToQuantumJaxprInterpreter.register_primitive(operator_p) -def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, **kwargs): +def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, n_ctrls, adjoint, **kwargs): if hybrid_lens or hybrid_trees or op_cls.static_argnames: # only support compilable_argnames for the moment raise NotImplementedError - wire_inputs = args[len(op_cls.dynamic_argnames) :] + if n_ctrls: + wire_inputs = args[len(op_cls.dynamic_argnames) : -2 * n_ctrls] + control_wire_inputs = args[-2 * n_ctrls : -n_ctrls] + control_values = args[-n_ctrls:] + else: + wire_inputs = args[len(op_cls.dynamic_argnames) :] + control_wire_inputs = control_values = () + new_wires = [ w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in wire_inputs ] + new_control_wires = [ + w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) + for w in control_wire_inputs + ] qref_operator_p.bind( *args[: len(op_cls.dynamic_argnames)], *new_wires, + *new_control_wires, + *control_values, op_cls=op_cls, hybrid_lens=hybrid_lens, hybrid_trees=hybrid_trees, + n_ctrls=n_ctrls, + adjoint=adjoint, **kwargs, ) return [] From 95f62c307128f1faf1a3293a36d1d946aca62a75 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 29 Jun 2026 11:21:03 -0400 Subject: [PATCH 27/48] Change adjoint from property to attribute --- mlir/include/QRef/IR/QRefOps.td | 2 +- mlir/include/Quantum/IR/QuantumOps.td | 2 +- mlir/lib/QRef/IR/QRefOps.cpp | 8 +++++--- mlir/lib/QRef/Transforms/value_semantics_conversion.cpp | 1 - mlir/lib/Quantum/IR/QuantumOps.cpp | 9 ++++++--- .../Transforms/reference_semantics_conversion.cpp | 1 - 6 files changed, 13 insertions(+), 10 deletions(-) diff --git a/mlir/include/QRef/IR/QRefOps.td b/mlir/include/QRef/IR/QRefOps.td index 028916927f..5d8a6278ec 100644 --- a/mlir/include/QRef/IR/QRefOps.td +++ b/mlir/include/QRef/IR/QRefOps.td @@ -429,7 +429,7 @@ def OperatorOp : UnitaryGate_Op<"operator", [ParametrizedGate, AttrSizedOperandS // additional functionality provided by quantum.operator, being last allows the builder // to use default values in the signature, so that these args are truly optional - UnitProp:$adjoint, + UnitAttr:$adjoint, OptionalProp:$UID, DefaultValuedOptionalAttr:$static_data, DefaultValuedOptionalAttr:$param_map, diff --git a/mlir/include/Quantum/IR/QuantumOps.td b/mlir/include/Quantum/IR/QuantumOps.td index 8623922262..9f1cdbf975 100644 --- a/mlir/include/Quantum/IR/QuantumOps.td +++ b/mlir/include/Quantum/IR/QuantumOps.td @@ -630,7 +630,7 @@ def OperatorOp : UnitaryGate_Op<"operator", [ParametrizedGate, NoMemoryEffect, // additional functionality provided by quantum.operator, being last allows the builder // to use default values in the signature, so that these args are truly optional - UnitProp:$adjoint, + UnitAttr:$adjoint, OptionalProp:$UID, DefaultValuedOptionalAttr:$static_data, DefaultValuedOptionalAttr:$param_map, diff --git a/mlir/lib/QRef/IR/QRefOps.cpp b/mlir/lib/QRef/IR/QRefOps.cpp index 3e754fc03b..fa53be7400 100644 --- a/mlir/lib/QRef/IR/QRefOps.cpp +++ b/mlir/lib/QRef/IR/QRefOps.cpp @@ -23,6 +23,8 @@ #include "QRef/IR/QRefDialect.h" #include "Quantum/IR/QuantumInterfaces.h" +#include + using namespace mlir; using namespace catalyst::qref; @@ -481,8 +483,8 @@ void OperatorOp::print(OpAsmPrinter &p) } // 5. Attribute Dictionary - SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", - "operandSegmentSizes", "op_name"}; + SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", + "operandSegmentSizes", "op_name", "adjoint"}; p.printOptionalAttrDict(getOperation()->getAttrs(), elidedAttrs); p.increaseIndent(); @@ -616,7 +618,7 @@ ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &result) // 3. Optional adjoint marker. if (succeeded(parser.parseOptionalKeyword("adj"))) { - opProperties.setAdjoint(true); + result.addAttribute("adjoint", builder.getUnitAttr()); } SmallVector qubits; diff --git a/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp b/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp index 689af5bb9c..68cb97ea67 100644 --- a/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp +++ b/mlir/lib/QRef/Transforms/value_semantics_conversion.cpp @@ -1159,7 +1159,6 @@ void handleGate(IRRewriter &builder, qref::QuantumOperation rGateOp, QubitValueT builder.getDenseI32ArrayAttr({nTargets, nCtrls, nQreg})); // Properties are not handled via the generic attribute fields, so we set them separately. - vGateOp.setAdjoint(rOperatorOp.getAdjoint()); vGateOp.setUID(rOperatorOp.getUID()); } else { diff --git a/mlir/lib/Quantum/IR/QuantumOps.cpp b/mlir/lib/Quantum/IR/QuantumOps.cpp index 62977102c4..0f7e36e978 100644 --- a/mlir/lib/Quantum/IR/QuantumOps.cpp +++ b/mlir/lib/Quantum/IR/QuantumOps.cpp @@ -30,6 +30,8 @@ #include "Quantum/IR/QuantumAttrDefs.h" #include "Quantum/IR/QuantumDialect.h" +#include + using namespace mlir; using namespace catalyst::quantum; @@ -785,8 +787,9 @@ void OperatorOp::print(OpAsmPrinter &p) } // 5. Attribute Dictionary - SmallVector elidedAttrs = {"static_data", "param_map", "qubit_map", - "operandSegmentSizes", "resultSegmentSizes", "op_name"}; + SmallVector elidedAttrs = { + "static_data", "param_map", "qubit_map", "operandSegmentSizes", + "resultSegmentSizes", "op_name", "adjoint"}; p.printOptionalAttrDict(getOperation()->getAttrs(), elidedAttrs); p.increaseIndent(); @@ -920,7 +923,7 @@ ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &result) // 3. Optional adjoint marker. if (succeeded(parser.parseOptionalKeyword("adj"))) { - opProperties.setAdjoint(true); + result.addAttribute("adjoint", builder.getUnitAttr()); } SmallVector inQubits; diff --git a/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp b/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp index 5c35f70782..d89082f4ed 100644 --- a/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp +++ b/mlir/lib/Quantum/Transforms/reference_semantics_conversion.cpp @@ -368,7 +368,6 @@ void handleGate(IRRewriter &builder, quantum::QuantumOperation vGateOp, QubitVal rGateOp->removeAttr("resultSegmentSizes"); // Properties are not handled via the generic attribute fields, so we set them separately. - rGateOp.setAdjoint(vOperatorOp.getAdjoint()); rGateOp.setUID(vOperatorOp.getUID()); } else { From 801caaa2783c62dc10e9f471987548d953188715 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 29 Jun 2026 13:04:21 -0400 Subject: [PATCH 28/48] Revert formatting --- .../catalyst/from_plxpr/qfunc_interpreter.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index b7ed5c968c..8c894b7b3d 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -210,24 +210,16 @@ def _check_measurement_with_dynamic_allocation(self, measurement): """Check some constraints regarding dynamic allocation.""" if self.has_dynamic_allocation: if len(measurement.wires) == 0 and not isinstance(measurement, qp.measurements.StateMP): - raise CompileError( - textwrap.dedent( - """ + raise CompileError(textwrap.dedent(""" Terminal measurements must take in an explicit list of wires when dynamically allocated wires are present in the program. - """ - ) - ) + """)) if any(is_abstract_qubit(w) for w in measurement.wires): - raise CompileError( - textwrap.dedent( - """ + raise CompileError(textwrap.dedent(""" Terminal measurements cannot take in dynamically allocated wires since they must be temporary. - """ - ) - ) + """)) # pylint: disable=too-many-branches def interpret_measurement(self, measurement): From a561f536bfb17c6b570ffa5dc1bf51cc4ed8ec63 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 29 Jun 2026 13:45:38 -0400 Subject: [PATCH 29/48] Add lit tests --- doc/releases/changelog-dev.md | 3 +- frontend/test/lit/test_operator.py | 75 ++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index ae615044ad..327762a341 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -9,7 +9,8 @@ without non-lowerable arguments. `Operator2` classes are now lowered to specialized operations where applicable, unlocking compilation and execution for these cases. [(#2979)](https://github.com/PennyLaneAI/catalyst/pull/2979) - [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969/) + [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969) + [(#2980)](https://github.com/PennyLaneAI/catalyst/pull/2980) * The `ResourceAnalysis` pass now reports each loop body and each subroutine as its own entry instead of folding their gate counts into the caller. Loops with constant bounds appear as `for_loop_` diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index f01236e002..cd4659ed14 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -377,3 +377,78 @@ def c_pcphase(x: float, dim: int): print(c_pcphase.mlir) + + +@qp.qjit(capture=True) +@qp.qnode(qp.device("lightning.qubit", wires=1)) +def c_adjoint(): + # CHECK-LABEL: func.func public @c_adjoint + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + op0 = NoParams(wires=0) + qp.adjoint(op0) + + # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) + qp.adjoint(qp.adjoint(op0)) + return qp.state() + + +print(c_adjoint.mlir) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=4)) +def c_controlled(): + # CHECK-LABEL: func.func public @c_controlled + + # CHECK: [[true:%.+]] = arith.constant true + # CHECK: [[false:%.+]] = arith.constant false + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + + # CHECK: qref.operator "NoParams"() qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + op0 = NoParams(wires=0) + qp.ops.ControlledOp2(op0, [1], control_values=[False]) + + # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1_1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2_1:%.+]] = qref.get {{%.+}}[ 2] + # CHECK: [[q3_1:%.+]] = qref.get {{%.+}}[ 3] + + # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) + # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[true]], [[false]], [[true]]) + qp.ops.ControlledOp2( + qp.ops.ControlledOp2(op0, [1], control_values=[True]), [2, 3], control_values=[False, True] + ) + return qp.state() + + +print(c_controlled.mlir) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=4)) +def c_adjoint_and_controlled(): + # CHECK-LABEL: func.func public @c_adjoint_and_controlled + + # CHECK: [[false:%.+]] = arith.constant false + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + op0 = NoParams(wires=0) + qp.ops.ControlledOp2(qp.adjoint(op0), [1], [False]) + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + qp.adjoint(qp.ops.ControlledOp2(op0, [1], [False])) + return qp.state() + + +print(c_adjoint_and_controlled.mlir) From 1ae38755fb836831e750cd636c4355253d1dfdeb Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 29 Jun 2026 16:24:52 -0400 Subject: [PATCH 30/48] Update operator_p import path --- frontend/catalyst/from_plxpr/qfunc_interpreter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index 8c894b7b3d..4daad35beb 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -28,9 +28,9 @@ from pennylane.capture.primitives import cond_prim as pl_cond_prim from pennylane.capture.primitives import ctrl_transform_prim as plxpr_ctrl_transform_prim from pennylane.capture.primitives import measure_prim as plxpr_measure_prim +from pennylane.capture.primitives import operator_p from pennylane.capture.primitives import pauli_measure_prim as plxpr_pauli_measure_prim from pennylane.capture.primitives import quantum_subroutine_prim, transform_prim -from pennylane.core.operator.operator2 import operator_p from pennylane.ftqc.primitives import measure_in_basis_prim as plxpr_measure_in_basis_prim from pennylane.measurements import CountsMP from pennylane.wires import AbstractQubit, is_abstract_qubit From 0f8efa82251b46c559e7cc9448c6fd97159d2e14 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 29 Jun 2026 16:27:17 -0400 Subject: [PATCH 31/48] [skip ci] Skip CI From 6a2b58133edaa7243ef3ef25dafbb8b71c50c357 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Tue, 30 Jun 2026 16:30:51 -0400 Subject: [PATCH 32/48] Transmit ctrl/adjoint information to spacial lowerings --- frontend/test/lit/test_operator.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index cd4659ed14..415f3141ca 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -410,7 +410,9 @@ def c_controlled(): # CHECK: qref.operator "NoParams"() qubits([[q0]]) # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) - op0 = NoParams(wires=0) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + op0 = NoParams(reg=0) qp.ops.ControlledOp2(op0, [1], control_values=[False]) # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] @@ -420,6 +422,8 @@ def c_controlled(): # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[true]], [[false]], [[true]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} qp.ops.ControlledOp2( qp.ops.ControlledOp2(op0, [1], control_values=[True]), [2, 3], control_values=[False, True] ) @@ -440,13 +444,17 @@ def c_adjoint_and_controlled(): # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) - op0 = NoParams(wires=0) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + op0 = NoParams(reg=0) qp.ops.ControlledOp2(qp.adjoint(op0), [1], [False]) # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} qp.adjoint(qp.ops.ControlledOp2(op0, [1], [False])) return qp.state() From 360851edab4bf0e0d33bb10a7f4839cf811d1a99 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Tue, 30 Jun 2026 16:30:51 -0400 Subject: [PATCH 33/48] [skip ci] Skip CI From e235a1f7aecb34ade56cc650324886ac20229dd4 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Thu, 2 Jul 2026 15:50:05 -0400 Subject: [PATCH 34/48] LIT IS THE WORST THING EVER INVENTED --- frontend/test/lit/test_operator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index f01236e002..31ea141b44 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -235,7 +235,7 @@ def __init__(self, wires, a, b, c): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_param_custom(): - # CHECK-LABEL: func.func public @circuit_multi_param_custom + # CHECK-LABEL: func.func public @c_multi_param_custom # CHECK: [[q0:%.+]] = qref.get {{%.+}} # pylint: disable=line-too-long From 01f2b2361e947cea76305e14ea40cc365b6257a5 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Thu, 2 Jul 2026 15:50:23 -0400 Subject: [PATCH 35/48] Update dep-versions --- .dep-versions | 2 +- doc/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.dep-versions b/.dep-versions index eff2c832ec..2b4af52c04 100644 --- a/.dep-versions +++ b/.dep-versions @@ -8,7 +8,7 @@ enzyme=v0.0.238 # For a custom PL version, update the package version here and at # 'doc/requirements.txt' -pennylane=0.46.0.dev38 +pennylane=0.46.0.dev45 # For a custom LQ/LK version, update the package version here and at # 'doc/requirements.txt' diff --git a/doc/requirements.txt b/doc/requirements.txt index 3459767b3b..067df129b4 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -34,4 +34,4 @@ lxml_html_clean --extra-index-url https://test.pypi.org/simple/ pennylane-lightning-kokkos==0.46.0-dev10 pennylane-lightning==0.46.0-dev10 -pennylane==0.46.0.dev38 +pennylane==0.46.0.dev45 From f65c9675e2759da230625c3651f35df0e29c3881 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Thu, 2 Jul 2026 15:50:24 -0400 Subject: [PATCH 36/48] [skip ci] Skip CI From e6788d3313c9391378c7450311a75d08703a278a Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Thu, 2 Jul 2026 16:14:33 -0400 Subject: [PATCH 37/48] Update based on upstream changes --- .../from_plxpr/qref_operator2_primitives.py | 125 ++++++------ frontend/test/lit/test_operator.py | 178 +++++++++--------- 2 files changed, 144 insertions(+), 159 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 2b5e0697db..72bcfbb50a 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -104,17 +104,40 @@ def _qref_operator_p_lowering( op_cls, **kwargs, ): - _general_validation(*args, op_cls=op_cls, **kwargs) ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True - if op_cls.__name__ in _SPECIAL_LOWERINGS: - return _SPECIAL_LOWERINGS[op_cls.__name__](jax_ctx, *args, op_cls=op_cls, **kwargs) - # will be used in future improvements + _general_validation(*args, op_cls=op_cls, **kwargs) + hybrid_lens = kwargs.pop("hybrid_lens") # pylint: disable=unused-variable hybrid_trees = kwargs.pop("hybrid_trees") # pylint: disable=unused-variable wire_lens = kwargs.pop("wire_lens") + adjoint = kwargs.pop("adjoint") + n_ctrls = kwargs.pop("n_ctrls") + + if n_ctrls: + ctrl_qubits = args[-2 * n_ctrls : -n_ctrls] + ctrl_values = [ + TensorExtractOp(ir.IntegerType.get_signless(1), val, []).result + for val in args[-n_ctrls:] + ] + qubits_slice = slice(len(op_cls.dynamic_argnames), -2 * n_ctrls) + else: + ctrl_qubits = ctrl_values = () + qubits_slice = slice(len(op_cls.dynamic_argnames), None) + params = args[: len(op_cls.dynamic_argnames)] - qubits = args[len(op_cls.dynamic_argnames) :] + qubits = args[qubits_slice] + + if op_cls.__name__ in _SPECIAL_LOWERINGS: + return _SPECIAL_LOWERINGS[op_cls.__name__]( + *params, + *qubits, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, + **kwargs, + ) + # will be used in future improvements name_attr = get_mlir_attribute_from_pyval(op_cls.__name__) @@ -134,10 +157,6 @@ def _qref_operator_p_lowering( processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map) - ctrl_qubits = [] - ctrl_values = [] - adjoint = False - if _is_custom_op(op_cls, jax_ctx.avals_in[: len(op_cls.dynamic_argnames)]): params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params] CustomOp( @@ -155,9 +174,9 @@ def _qref_operator_p_lowering( qubits=qubits, qreg=None, forward_args=[], - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, UID=None, arr_qubit_indices=[], param_map=processed_param_map, @@ -168,58 +187,43 @@ def _qref_operator_p_lowering( @_register_special_lowering("MultiRZ") -def _multirz_lowering(jax_ctx: mlir.LoweringRuleContext, *args, **_): - theta = extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ") - qubits = args[1:] +def _multirz_lowering(theta, *qubits, ctrl_qubits, ctrl_values, adjoint): MultiRZOp( - theta=theta, + theta=extract_scalar(safe_cast_to_f64(theta, "MultiRZ"), "MultiRZ"), qubits=qubits, - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, ) return [] @_register_special_lowering("PCPhase") -def _pcphase_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - **_, -): - qubits = args[2:] +def _pcphase_lowering(theta, dim, *qubits, ctrl_qubits, ctrl_values, adjoint): PCPhaseOp( - theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"), - dim=extract_scalar(safe_cast_to_f64(args[1], "PCPhase"), "PCPhase"), + theta=extract_scalar(safe_cast_to_f64(theta, "PCPhase"), "PCPhase"), + dim=extract_scalar(safe_cast_to_f64(dim, "PCPhase"), "PCPhase"), qubits=qubits, - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, ) return () @_register_special_lowering("GlobalPhase") -def _special_gphase_lowering( - jax_ctx: mlir.LoweringRuleContext, - *args, - op_cls, - **_, -): +def _special_gphase_lowering(angle, *_, ctrl_qubits, ctrl_values, adjoint): GlobalPhaseOp( - angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"), - ctrl_qubits=[], - ctrl_values=[], - adjoint=False, + angle=extract_scalar(safe_cast_to_f64(angle, "GlobalPhase"), "GlobalPhase"), + ctrl_qubits=ctrl_qubits, + ctrl_values=ctrl_values, + adjoint=adjoint, ) return () @_register_special_lowering("QubitUnitary") -def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits, **_): - ctrl_qubits = [] - ctrl_values = [] - +def _special_unitary_lowering(matrix, *qubits, ctrl_qubits, ctrl_values, adjoint): matrix_type = matrix.type is_tensor = ir.RankedTensorType.isinstance(matrix_type) shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None @@ -243,50 +247,29 @@ def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits tensor_complex_f64_type = ir.RankedTensorType.get(shape, complex_f64_type) matrix = StableHLOConvertOp(tensor_complex_f64_type, matrix).result - ctrl_values_i1 = [ - TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values - ] - QubitUnitaryOp( matrix=matrix, qubits=qubits, ctrl_qubits=ctrl_qubits, - ctrl_values=ctrl_values_i1, - adjoint=False, + ctrl_values=ctrl_values, + adjoint=adjoint, ) return () @_register_special_lowering("PauliRot") -def _special_paulirot_lowering( - jax_ctx: mlir.LoweringRuleContext, - angle, - *qubits, - pauli_word, - **_, -): +def _special_paulirot_lowering(angle, *qubits, ctrl_qubits, ctrl_values, adjoint, pauli_word): pauli_word = unflatten(*pauli_word) - ctrl_qubits = [] - ctrl_values = [] - - angle = safe_cast_to_f64(angle, "PauliRot") - angle = extract_scalar(angle, "PauliRot") - assert ir.F64Type.isinstance(angle.type) - pauli_word = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in pauli_word]) - ctrl_values_i1 = [ - TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values - ] - PauliRotOp( - angle=angle, + angle=extract_scalar(safe_cast_to_f64(angle, "PauliRot"), "PauliRot"), pauli_product=pauli_word, qubits=qubits, ctrl_qubits=ctrl_qubits, - ctrl_values=ctrl_values_i1, - adjoint=False, + ctrl_values=ctrl_values, + adjoint=adjoint, ) return () diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 415f3141ca..0a3770f868 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -53,6 +53,89 @@ def c_no_params(): print(c_no_params.mlir) +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("lightning.qubit", wires=1)) +def c_adjoint(): + # CHECK-LABEL: func.func public @c_adjoint + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + op0 = NoParams(reg=0) + qp.adjoint(op0) + + # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) + qp.adjoint(qp.adjoint(op0)) + return qp.state() + + +print(c_adjoint.mlir) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=4)) +def c_controlled(): + # CHECK-LABEL: func.func public @c_controlled + + # CHECK: [[true:%.+]] = arith.constant true + # CHECK: [[false:%.+]] = arith.constant false + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + + # CHECK: qref.operator "NoParams"() qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + op0 = NoParams(reg=0) + qp.ops.ControlledOp2(op0, [1], control_values=[False]) + + # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1_1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2_1:%.+]] = qref.get {{%.+}}[ 2] + # CHECK: [[q3_1:%.+]] = qref.get {{%.+}}[ 3] + + # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) + # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[false]], [[true]], [[true]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + qp.ops.ControlledOp2( + qp.ops.ControlledOp2(op0, [3], control_values=[True]), [1, 2], control_values=[False, True] + ) + return qp.state() + + +print(c_controlled.mlir) + + +@qp.qjit(capture=True, target="mlir") +@qp.qnode(qp.device("null.qubit", wires=4)) +def c_adjoint_and_controlled(): + # CHECK-LABEL: func.func public @c_adjoint_and_controlled + + # CHECK: [[false:%.+]] = arith.constant false + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + op0 = NoParams(reg=0) + qp.ops.ControlledOp2(qp.adjoint(op0), [1], [False]) + + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) + # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + qp.adjoint(qp.ops.ControlledOp2(op0, [1], [False])) + return qp.state() + + +print(c_adjoint_and_controlled.mlir) + + class NoParamsCustomOp(qp.core.Operator2): def __init__(self, wires): @@ -235,7 +318,7 @@ def __init__(self, wires, a, b, c): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_param_custom(): - # CHECK-LABEL: func.func public @circuit_multi_param_custom + # CHECK-LABEL: func.func public @c_multi_param_custom # CHECK: [[q0:%.+]] = qref.get {{%.+}} # pylint: disable=line-too-long @@ -257,7 +340,9 @@ def __init__(self, phi, wires): @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=2)) -def circuit(x: float): +def circuit_multirz(x: float): + # CHECK-LABEL: func.func public @circuit_multirz + # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} # CHECK: [[q2:%.+]] = qref.get {{%.+}} @@ -267,7 +352,7 @@ def circuit(x: float): return qp.state() -print(circuit.mlir) +print(circuit_multirz.mlir) class PauliRot(qp.core.Operator2): @@ -334,7 +419,7 @@ def __init__(self, matrix, wires): super().__init__(matrix, wires) -@qp.qjit(capture=True) +@qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("lightning.qubit", wires=3)) def circuit_qubitunitary(): # CHECK-LABEL: func.func public @circuit_qubitunitary @@ -363,7 +448,7 @@ def __init__(self, phi, dim, wires): super().__init__(phi, dim, wires) -@qp.qjit(capture=True) +@qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("lightning.qubit", wires=2)) def c_pcphase(x: float, dim: int): # CHECK-LABEL: func.func public @c_pcphase @@ -377,86 +462,3 @@ def c_pcphase(x: float, dim: int): print(c_pcphase.mlir) - - -@qp.qjit(capture=True) -@qp.qnode(qp.device("lightning.qubit", wires=1)) -def c_adjoint(): - # CHECK-LABEL: func.func public @c_adjoint - - # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) - op0 = NoParams(wires=0) - qp.adjoint(op0) - - # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) - qp.adjoint(qp.adjoint(op0)) - return qp.state() - - -print(c_adjoint.mlir) - - -@qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device("null.qubit", wires=4)) -def c_controlled(): - # CHECK-LABEL: func.func public @c_controlled - - # CHECK: [[true:%.+]] = arith.constant true - # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] - - # CHECK: qref.operator "NoParams"() qubits([[q0]]) - # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) - # CHECK-NEXT: static_data = {} - # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - op0 = NoParams(reg=0) - qp.ops.ControlledOp2(op0, [1], control_values=[False]) - - # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: [[q1_1:%.+]] = qref.get {{%.+}}[ 1] - # CHECK: [[q2_1:%.+]] = qref.get {{%.+}}[ 2] - # CHECK: [[q3_1:%.+]] = qref.get {{%.+}}[ 3] - - # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) - # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[true]], [[false]], [[true]]) - # CHECK-NEXT: static_data = {} - # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - qp.ops.ControlledOp2( - qp.ops.ControlledOp2(op0, [1], control_values=[True]), [2, 3], control_values=[False, True] - ) - return qp.state() - - -print(c_controlled.mlir) - - -@qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device("null.qubit", wires=4)) -def c_adjoint_and_controlled(): - # CHECK-LABEL: func.func public @c_adjoint_and_controlled - - # CHECK: [[false:%.+]] = arith.constant false - - # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] - # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) - # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) - # CHECK-NEXT: static_data = {} - # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - op0 = NoParams(reg=0) - qp.ops.ControlledOp2(qp.adjoint(op0), [1], [False]) - - # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] - # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] - # CHECK: qref.operator "NoParams"() adj qubits([[q0]]) - # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) - # CHECK-NEXT: static_data = {} - # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - qp.adjoint(qp.ops.ControlledOp2(op0, [1], [False])) - return qp.state() - - -print(c_adjoint_and_controlled.mlir) From c23610248cc0d832b48360b75ef53640e30b17db Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 11:56:40 -0400 Subject: [PATCH 38/48] Finish adding tests. Won't pass until PL-side fix is uploaded to TestPyPI --- frontend/test/lit/test_operator.py | 59 ++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 3 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 0a3770f868..45b20310f6 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -54,7 +54,7 @@ def c_no_params(): @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device("lightning.qubit", wires=1)) +@qp.qnode(qp.device("null.qubit", wires=1)) def c_adjoint(): # CHECK-LABEL: func.func public @c_adjoint @@ -146,6 +146,7 @@ def __init__(self, wires): @qp.qnode(qp.device("null.qubit", wires=2)) def c_no_params_custom(): # CHECK-LABEL: func.func public @c_no_params_custom + # CHECK: [[false:%.+]] = arith.constant false # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.custom "NoParamsCustomOp"() [[q0]] : !qref.bit @@ -155,6 +156,12 @@ def c_no_params_custom(): # CHECK: [[q2:%.+]] = qref.get {{%.+}} # CHECK: qref.custom "NoParamsCustomOp"() [[q1]], [[q2]] : !qref.bit, !qref.bit NoParamsCustomOp(wires=(0, 1)) + + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: qref.custom "NoParamsCustomOp"() [[q3]] adj ctrls([[q4]]) + # CHECK-SAME: ctrlvals([[false]]) : !qref.bit ctrls !qref.bit + qp.ctrl(qp.adjoint(NoParamsCustomOp(wires=(0,))), [2], [False]) return qp.state() @@ -339,9 +346,10 @@ def __init__(self, phi, wires): @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device("null.qubit", wires=2)) +@qp.qnode(qp.device("null.qubit", wires=4)) def circuit_multirz(x: float): # CHECK-LABEL: func.func public @circuit_multirz + # CHECK: [[false:%.+]] = arith.constant false # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} @@ -349,6 +357,16 @@ def circuit_multirz(x: float): # CHECK: qref.multirz({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit MultiRZ(x, (0, 1, 2)) + + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q5:%.+]] = qref.get {{%.+}} + + # MultiRZ gets automatically canonicalized, so it will never have the `adj` attribute + # CHECK: [[theta:%.+]] = arith.negf {{%.+}} : f64 + # CHECK: qref.multirz([[theta]]) [[q3]], [[q4]] ctrls([[q5]]) ctrlvals([[false]]) : + # CHECK-SAME: !qref.bit, !qref.bit ctrls !qref.bit + qp.ctrl(qp.adjoint(MultiRZ(x, wires=(0, 1))), [3], [False]) return qp.state() @@ -368,6 +386,7 @@ def __init__(self, phi, pauli_word, wires): @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_paulirot(x: float): # CHECK-LABEL: func.func public @circuit_paulirot + # CHECK: [[false:%.+]] = arith.constant false # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: [[q1:%.+]] = qref.get {{%.+}} @@ -383,6 +402,14 @@ def circuit_paulirot(x: float): # CHECK: qref.paulirot ["Y", "Z", "X"]({{%.+}}) [[q3]], [[q4]], [[q5]] : !qref.bit, !qref.bit, !qref.bit PauliRot(x, "YZX", (0, 1, 2)) + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q5:%.+]] = qref.get {{%.+}} + + # CHECK: qref.paulirot ["Y", "Z"]({{%.+}}) [[q3]], [[q4]] adj ctrls([[q5]]) + # CHECK-SAME: ctrlvals([[false]]) : !qref.bit, !qref.bit ctrls !qref.bit + qp.ctrl(qp.adjoint(PauliRot(x, "YZ", (0, 1))), [2], [False]) + return qp.probs(wires=(0, 1, 2)) @@ -402,9 +429,14 @@ def __init__(self, phi): @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_gphase(x: float): # CHECK-LABEL: func.func public @circuit_gphase + # CHECK: [[false:%.+]] = arith.constant false # CHECK: qref.gphase({{%.+}}) GlobalPhase(x) + + # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: qref.gphase({{%.+}}) adj ctrls([[q0]]) ctrlvals([[false]]) : ctrls !qref.bit + qp.ctrl(qp.adjoint(GlobalPhase(x)), [0], [False]) return qp.state() @@ -423,6 +455,7 @@ def __init__(self, matrix, wires): @qp.qnode(qp.device("lightning.qubit", wires=3)) def circuit_qubitunitary(): # CHECK-LABEL: func.func public @circuit_qubitunitary + # CHECK: [[false:%.+]] = arith.constant false # CHECK: [[q0:%.+]] = qref.get {{%.+}} # CHECK: qref.unitary({{%.+}} : tensor<2x2xcomplex>) [[q0]] : !qref.bit @@ -434,6 +467,14 @@ def circuit_qubitunitary(): # CHECK: qref.unitary({{%.+}} : tensor<4x4xcomplex>) [[q1]], [[q2]] : !qref.bit, !qref.bit QubitUnitary(qp.CNOT.compute_matrix(), (0, 1)) + + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: qref.unitary({{%.+}} : tensor<4x4xcomplex>) [[q3]], [[q4]] adj ctrls([[q5]]) + # CHECK-SAME: ctrlvals([[false]]) : !qref.bit, !qref.bit ctrls !qref.bit + + qp.ctrl(qp.adjoint(QubitUnitary(qp.CNOT.compute_matrix(), (0, 1))), [2], [False]) return qp.expval(qp.Z(0)), qp.expval(qp.Z(0)) @@ -449,15 +490,27 @@ def __init__(self, phi, dim, wires): @qp.qjit(capture=True, target="mlir") -@qp.qnode(qp.device("lightning.qubit", wires=2)) +@qp.qnode(qp.device("lightning.qubit", wires=3)) def c_pcphase(x: float, dim: int): # CHECK-LABEL: func.func public @c_pcphase + # CHECK: [[false:%.+]] = arith.constant false # CHECK: [[q11:%.+]] = qref.get {{%.+}} # CHECK: [[q12:%.+]] = qref.get {{%.+}} # CHECK: qref.pcphase({{%.+}}, {{%.+}}) [[q11]], [[q12]] : !qref.bit, !qref.bit PCPhase(x, dim, (0, 1)) + + # CHECK: [[q3:%.+]] = qref.get {{%.+}} + # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q5:%.+]] = qref.get {{%.+}} + + # PCPhase gets automatically canonicalized, so it will never have the `adj` attribute + # CHECK: [[theta:%.+]] = arith.negf {{%.+}} : f64 + # CHECK: qref.pcphase([[theta]], {{%.+}}) [[q3]], [[q4]] ctrls([[q5]]) ctrlvals([[false]]) : + # CHECK-SAME: !qref.bit, !qref.bit ctrls !qref.bit + + qp.ctrl(qp.adjoint(PCPhase(x, dim, (0, 1))), [2], [False]) return qp.state() From e3b96bb6131c02395297115185de432c8fb77917 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:25:36 -0400 Subject: [PATCH 39/48] Fix merge conflict --- frontend/catalyst/from_plxpr/qref_operator2_primitives.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 1b3fb8f352..8c13a2e22e 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -110,11 +110,9 @@ def _qref_operator_p_lowering( hybrid_lens = kwargs.pop("hybrid_lens") # pylint: disable=unused-variable hybrid_trees = kwargs.pop("hybrid_trees") # pylint: disable=unused-variable - adjoint = kwargs.pop("adjoint") # pylint: disable=unused-variable - n_ctrls = kwargs.pop("n_ctrls") # pylint: disable=unused-variable - wire_lens = kwargs.pop("wire_lens") adjoint = kwargs.pop("adjoint") n_ctrls = kwargs.pop("n_ctrls") + wire_lens = kwargs.pop("wire_lens") if n_ctrls: ctrl_qubits = args[-2 * n_ctrls : -n_ctrls] From 20998c6b7e20d483a1ea8560183c7f37a532f65a Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:25:38 -0400 Subject: [PATCH 40/48] [skip ci] Skip CI From aab165aa0d76674f1c4ac0a3f547a04ff46c6558 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:27:39 -0400 Subject: [PATCH 41/48] [skip ci] From 16f9aeeec338a9d079929a82a7d63f5e9a0f84eb Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:28:00 -0400 Subject: [PATCH 42/48] [skip ci] From e273e736115a90e67a9fd1b4d9dccd5fe99d6b94 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:55:09 -0400 Subject: [PATCH 43/48] Remove wayward comment --- frontend/catalyst/from_plxpr/qref_operator2_primitives.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 8c13a2e22e..75c1275362 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -137,7 +137,6 @@ def _qref_operator_p_lowering( adjoint=adjoint, **kwargs, ) - # will be used in future improvements name_attr = get_mlir_attribute_from_pyval(op_cls.__name__) From 429fda0c4d3d64f77dfa657ec5f74b7f2719782b Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:55:11 -0400 Subject: [PATCH 44/48] [skip ci] Skip CI From fd14cb9c6ebede97485b0e5551301939518c3597 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:57:25 -0400 Subject: [PATCH 45/48] Replace qp.ops.ControlledOp2 with qp.ctrl --- frontend/test/lit/test_operator.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 45b20310f6..89293f2c93 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -87,7 +87,7 @@ def c_controlled(): # CHECK-NEXT: static_data = {} # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} op0 = NoParams(reg=0) - qp.ops.ControlledOp2(op0, [1], control_values=[False]) + qp.ctrl(op0, [1], control_values=[False]) # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] # CHECK: [[q1_1:%.+]] = qref.get {{%.+}}[ 1] @@ -98,9 +98,7 @@ def c_controlled(): # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[false]], [[true]], [[true]]) # CHECK-NEXT: static_data = {} # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - qp.ops.ControlledOp2( - qp.ops.ControlledOp2(op0, [3], control_values=[True]), [1, 2], control_values=[False, True] - ) + qp.ctrl(qp.ctrl(op0, [3], control_values=[True]), [1, 2], control_values=[False, True]) return qp.state() @@ -121,7 +119,7 @@ def c_adjoint_and_controlled(): # CHECK-NEXT: static_data = {} # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} op0 = NoParams(reg=0) - qp.ops.ControlledOp2(qp.adjoint(op0), [1], [False]) + qp.ctrl(qp.adjoint(op0), [1], [False]) # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] @@ -129,7 +127,7 @@ def c_adjoint_and_controlled(): # CHECK-NEXT: ctrls([[q1]]) ctrl_vals([[false]]) # CHECK-NEXT: static_data = {} # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} - qp.adjoint(qp.ops.ControlledOp2(op0, [1], [False])) + qp.adjoint(qp.ctrl(op0, [1], [False])) return qp.state() From 6be5efbf422a8874e43f46a410d63fb41333b466 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Fri, 3 Jul 2026 17:57:28 -0400 Subject: [PATCH 46/48] [skip ci] Skip CI From 20a2867cf18f55f2fc20ecc3ba871c9a19364c78 Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Mon, 6 Jul 2026 11:05:34 -0400 Subject: [PATCH 47/48] Update PL dep --- .dep-versions | 2 +- doc/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.dep-versions b/.dep-versions index 6ae97bbca3..0c6ca0ce6c 100644 --- a/.dep-versions +++ b/.dep-versions @@ -8,7 +8,7 @@ enzyme=v0.0.238 # For a custom PL version, update the package version here and at # 'doc/requirements.txt' -pennylane=0.46.0.dev46 +pennylane=0.46.0.dev48 # For a custom LQ/LK version, update the package version here and at # 'doc/requirements.txt' diff --git a/doc/requirements.txt b/doc/requirements.txt index 7b92a16646..46d22b41dd 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -34,4 +34,4 @@ lxml_html_clean --extra-index-url https://test.pypi.org/simple/ pennylane-lightning-kokkos==0.46.0-dev15 pennylane-lightning==0.46.0-dev15 -pennylane==0.46.0.dev46 +pennylane==0.46.0.dev48 From 4b4d65f344787562d4c13112187281733816caec Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Tue, 7 Jul 2026 14:34:44 -0400 Subject: [PATCH 48/48] Address code review --- frontend/test/lit/test_operator.py | 111 ++++++++++++++++------------- mlir/lib/QRef/IR/QRefOps.cpp | 3 +- mlir/lib/Quantum/IR/QuantumOps.cpp | 3 +- 3 files changed, 62 insertions(+), 55 deletions(-) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index 89293f2c93..5afd98ba8e 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -35,14 +35,14 @@ def __init__(self, reg): def c_no_params(): # CHECK-LABEL: func.func public @c_no_params - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.operator "NoParams"() qubits([[q0]]) # CHECK: static_data = {} # CHECK: param_map = {} qubit_map = {reg = [0]} NoParams(reg=0) - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.operator "NoParams"() qubits([[q1]], [[q2]]) # CHECK: static_data = {} # CHECK: param_map = {} qubit_map = {reg = [0, 1]} @@ -91,11 +91,20 @@ def c_controlled(): # CHECK: [[q0_1:%.+]] = qref.get {{%.+}}[ 0] # CHECK: [[q1_1:%.+]] = qref.get {{%.+}}[ 1] - # CHECK: [[q2_1:%.+]] = qref.get {{%.+}}[ 2] - # CHECK: [[q3_1:%.+]] = qref.get {{%.+}}[ 3] - # CHECK: qref.operator "NoParams"() qubits([[q0_1]]) - # CHECK-NEXT: ctrls([[q1_1]], [[q2_1]], [[q3_1]]) ctrl_vals([[false]], [[true]], [[true]]) + # CHECK-NEXT: ctrls([[q1_1]]) ctrl_vals([[false]]) + # CHECK-NEXT: static_data = {} + # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} + op0 = NoParams(reg=0) + qp.ctrl(op0, [1], control_values=[0]) + + # CHECK: [[q0_2:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1_2:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2_2:%.+]] = qref.get {{%.+}}[ 2] + # CHECK: [[q3_2:%.+]] = qref.get {{%.+}}[ 3] + + # CHECK: qref.operator "NoParams"() qubits([[q0_2]]) + # CHECK-NEXT: ctrls([[q1_2]], [[q2_2]], [[q3_2]]) ctrl_vals([[false]], [[true]], [[true]]) # CHECK-NEXT: static_data = {} # CHECK-NEXT: param_map = {} qubit_map = {reg = [0]} qp.ctrl(qp.ctrl(op0, [3], control_values=[True]), [1, 2], control_values=[False, True]) @@ -146,17 +155,17 @@ def c_no_params_custom(): # CHECK-LABEL: func.func public @c_no_params_custom # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.custom "NoParamsCustomOp"() [[q0]] : !qref.bit NoParamsCustomOp(wires=0) - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.custom "NoParamsCustomOp"() [[q1]], [[q2]] : !qref.bit, !qref.bit NoParamsCustomOp(wires=(0, 1)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.custom "NoParamsCustomOp"() [[q3]] adj ctrls([[q4]]) # CHECK-SAME: ctrlvals([[false]]) : !qref.bit ctrls !qref.bit qp.ctrl(qp.adjoint(NoParamsCustomOp(wires=(0,))), [2], [False]) @@ -180,15 +189,15 @@ def __init__(self, x, reg): def c_single_param(x: float): # CHECK-LABEL: func.func public @c_single_param - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.operator "SingleParam"({{%.+}}: tensor) qubits([[q0]]) # CHECK: static_data = {} # CHECK: param_map = {x = [0]} qubit_map = {reg = [0]} SingleParam(x, 0) - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.operator "SingleParam"({{%.+}}: tensor<4x4xf64>) qubits([[q1]], [[q2]]) # CHECK: static_data = {} # CHECK: param_map = {x = [0]} qubit_map = {reg = [0, 1]} @@ -213,12 +222,12 @@ def __init__(self, x, wires): def c_single_param_custom(x: float): # CHECK-LABEL: func.func public @c_single_param_custom - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}) [[q0]] : !qref.bit SingleParamCustomOp(x, 0) - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.custom "SingleParamCustomOp"({{%.+}}) [[q1]], [[q2]] : !qref.bit, !qref.bit SingleParamCustomOp(0.5, (1, 2)) @@ -241,8 +250,8 @@ def __init__(self, a, b, thing, wires): def c_compilable(): # CHECK-LABEL: func.func public @c_compilable - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.operator "CompilableData"() qubits([[q1]], [[q2]]) # CHECK: static_data = {a = true, b = "some string", thing = [1, true, "string"]} # CHECK: param_map = {} qubit_map = {wires = [0, 1]} @@ -268,10 +277,10 @@ def __init__(self, reg1, reg2): def c_multiple_registers(): # CHECK-LABEL: func.func public @c_multiple_registers - # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 2] + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 3] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 4] # CHECK: qref.operator "MultipleRegisters"() qubits([[q0]], [[q2]], [[q3]], [[q4]]) # CHECK: static_data = {} @@ -298,7 +307,7 @@ def __init__(self, reg, a, b, c): def c_multi_params(): # CHECK-LABEL: func.func public @c_multi_params - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # pylint: disable=line-too-long # CHECK: qref.operator "MultiParams"({{%.+}}: tensor, {{%.+}}: tensor<4x2x1xf64>, {{%.+}}: tensor<3xi64>) qubits([[q0]]) @@ -324,7 +333,7 @@ def __init__(self, wires, a, b, c): @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_param_custom(): # CHECK-LABEL: func.func public @c_multi_param_custom - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # pylint: disable=line-too-long # CHECK: qref.custom "MultiParamsCustom"({{%.+}}, {{%.+}}, {{%.+}}) [[q0]] : !qref.bit @@ -349,16 +358,16 @@ def circuit_multirz(x: float): # CHECK-LABEL: func.func public @circuit_multirz # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.multirz({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit MultiRZ(x, (0, 1, 2)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} - # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q5:%.+]] = qref.get {{%.+}}[ 3] # MultiRZ gets automatically canonicalized, so it will never have the `adj` attribute # CHECK: [[theta:%.+]] = arith.negf {{%.+}} : f64 @@ -386,23 +395,23 @@ def circuit_paulirot(x: float): # CHECK-LABEL: func.func public @circuit_paulirot # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q0:%.+]] = qref.get {{%.+}} - # CHECK: [[q1:%.+]] = qref.get {{%.+}} - # CHECK: [[q2:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q1:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q2:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.paulirot ["X", "Y", "Z"]({{%.+}}) [[q0]], [[q1]], [[q2]] : !qref.bit, !qref.bit, !qref.bit PauliRot(x, "XYZ", (0, 1, 2)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} - # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q5:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.paulirot ["Y", "Z", "X"]({{%.+}}) [[q3]], [[q4]], [[q5]] : !qref.bit, !qref.bit, !qref.bit PauliRot(x, "YZX", (0, 1, 2)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} - # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q5:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.paulirot ["Y", "Z"]({{%.+}}) [[q3]], [[q4]] adj ctrls([[q5]]) # CHECK-SAME: ctrlvals([[false]]) : !qref.bit, !qref.bit ctrls !qref.bit @@ -432,7 +441,7 @@ def circuit_gphase(x: float): # CHECK: qref.gphase({{%.+}}) GlobalPhase(x) - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.gphase({{%.+}}) adj ctrls([[q0]]) ctrlvals([[false]]) : ctrls !qref.bit qp.ctrl(qp.adjoint(GlobalPhase(x)), [0], [False]) return qp.state() @@ -455,7 +464,7 @@ def circuit_qubitunitary(): # CHECK-LABEL: func.func public @circuit_qubitunitary # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q0:%.+]] = qref.get {{%.+}} + # CHECK: [[q0:%.+]] = qref.get {{%.+}}[ 0] # CHECK: qref.unitary({{%.+}} : tensor<2x2xcomplex>) [[q0]] : !qref.bit QubitUnitary(np.array([[0, 1], [1, 0]]), 0) @@ -466,9 +475,9 @@ def circuit_qubitunitary(): QubitUnitary(qp.CNOT.compute_matrix(), (0, 1)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} - # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q5:%.+]] = qref.get {{%.+}}[ 2] # CHECK: qref.unitary({{%.+}} : tensor<4x4xcomplex>) [[q3]], [[q4]] adj ctrls([[q5]]) # CHECK-SAME: ctrlvals([[false]]) : !qref.bit, !qref.bit ctrls !qref.bit @@ -493,15 +502,15 @@ def c_pcphase(x: float, dim: int): # CHECK-LABEL: func.func public @c_pcphase # CHECK: [[false:%.+]] = arith.constant false - # CHECK: [[q11:%.+]] = qref.get {{%.+}} - # CHECK: [[q12:%.+]] = qref.get {{%.+}} + # CHECK: [[q11:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q12:%.+]] = qref.get {{%.+}}[ 1] # CHECK: qref.pcphase({{%.+}}, {{%.+}}) [[q11]], [[q12]] : !qref.bit, !qref.bit PCPhase(x, dim, (0, 1)) - # CHECK: [[q3:%.+]] = qref.get {{%.+}} - # CHECK: [[q4:%.+]] = qref.get {{%.+}} - # CHECK: [[q5:%.+]] = qref.get {{%.+}} + # CHECK: [[q3:%.+]] = qref.get {{%.+}}[ 0] + # CHECK: [[q4:%.+]] = qref.get {{%.+}}[ 1] + # CHECK: [[q5:%.+]] = qref.get {{%.+}}[ 2] # PCPhase gets automatically canonicalized, so it will never have the `adj` attribute # CHECK: [[theta:%.+]] = arith.negf {{%.+}} : f64 diff --git a/mlir/lib/QRef/IR/QRefOps.cpp b/mlir/lib/QRef/IR/QRefOps.cpp index fa53be7400..06c4e99048 100644 --- a/mlir/lib/QRef/IR/QRefOps.cpp +++ b/mlir/lib/QRef/IR/QRefOps.cpp @@ -18,13 +18,12 @@ #include "llvm/Support/LogicalResult.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/OpImplementation.h" #include "QRef/IR/QRefDialect.h" #include "Quantum/IR/QuantumInterfaces.h" -#include - using namespace mlir; using namespace catalyst::qref; diff --git a/mlir/lib/Quantum/IR/QuantumOps.cpp b/mlir/lib/Quantum/IR/QuantumOps.cpp index 0f7e36e978..bc80ea377c 100644 --- a/mlir/lib/Quantum/IR/QuantumOps.cpp +++ b/mlir/lib/Quantum/IR/QuantumOps.cpp @@ -24,14 +24,13 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/OpImplementation.h" #include "QRef/IR/QRefOps.h" #include "Quantum/IR/QuantumAttrDefs.h" #include "Quantum/IR/QuantumDialect.h" -#include - using namespace mlir; using namespace catalyst::quantum;