diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 49c84247be..e952657db7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -2,13 +2,13 @@
New features since last release
-* You can now dynamically prepare magic T states inside captured Catalyst workflows using
- ``qp.allocate(state="magic-T")`` and ``qp.allocate(state="magic-T-adj")``, which makes it
- easier to compile FTQC-style routines that need T-state ancillas on the fly (for example
- TemporaryAND) with ``qjit(capture=True)``.
+* You can now dynamically allocate magic T states (and their conjugates) with Catalyst on
+ both the capture and legacy tracing paths. This lets you build workflows that combine
+ dynamic wire allocation, mid-circuit measurements, and control flow — for example preparing
+ magic-T ancillas, applying a Pauli-product measurement, and branching with ``qp.cond``.
```python
- @qjit(capture=True)
+ @qjit # works with capture=True or capture=False
@qnode(dev)
def circuit():
qb = qp.allocate(state="magic-T")
@@ -19,6 +19,7 @@
Requires a PennyLane release that supports the new ``AllocateState`` values
(see [PennyLane #9846](https://github.com/PennyLaneAI/pennylane/pull/9846)).
[(#3029)](https://github.com/PennyLaneAI/catalyst/pull/3029)
+ [(#3027)](https://github.com/PennyLaneAI/catalyst/pull/3027)
* The `local-random` unitary folding option for :func:`~.mitigate_with_zne` is now implemented,
reproducing Mitiq's ``fold_gates_at_random``: every gate is folded ``floor((scale_factor-1)/2)``
diff --git a/frontend/catalyst/api_extensions/control_flow.py b/frontend/catalyst/api_extensions/control_flow.py
index 415c9bcde4..0beb1f1e99 100644
--- a/frontend/catalyst/api_extensions/control_flow.py
+++ b/frontend/catalyst/api_extensions/control_flow.py
@@ -1631,7 +1631,7 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
partial(inner_trace.new_arg, source_info=current_source_info()), [new_qreg]
)[0]
qrp_out = trace_quantum_operations(
- inner_tape, device, qreg_in, ctx, inner_trace, **kwargs
+ inner_tape, device, qreg_in, ctx, inner_trace, parent_qrp=qrp, **kwargs
)
qreg_out = qrp_out.actualize()
@@ -1683,8 +1683,10 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
apply_reverse_transform=self.apply_reverse_transform,
num_implicit_inputs=num_implicit_inputs,
preserve_dimensions=not expansion_strategy.input_unshare_variables,
- )
+ ),
+ num_device_wires=qrp.num_device_wires,
)
+ qrp2.inherit_dynamic_wire_state(qrp)
return qrp2
@@ -1754,7 +1756,7 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
partial(body_trace.new_arg, source_info=current_source_info()), [AbstractQreg()]
)[0]
qrp_out = trace_quantum_operations(
- body_tape, device, qreg_in, ctx, body_trace, **kwargs
+ body_tape, device, qreg_in, ctx, body_trace, parent_qrp=qrp, **kwargs
)
qreg_out = qrp_out.actualize()
arg_expanded_tracers = expand_args(
@@ -1804,8 +1806,10 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
body_nconsts=len(body_consts),
num_implicit_inputs=num_implicit_inputs,
preserve_dimensions=not expansion_strategy.input_unshare_variables,
- )
+ ),
+ num_device_wires=qrp.num_device_wires,
)
+ qrp2.inherit_dynamic_wire_state(qrp)
return qrp2
@@ -1888,7 +1892,13 @@ def trace_quantum_branches(op, ctx, device, trace, qrp, **kwargs) -> QRegPromise
partial(region.trace.new_arg, source_info=current_source_info()), [new_qreg]
)[0]
qreg_out = trace_quantum_operations(
- region.quantum_tape, device, qreg_in, ctx, region.trace, **kwargs
+ region.quantum_tape,
+ device,
+ qreg_in,
+ ctx,
+ region.trace,
+ parent_qrp=qrp,
+ **kwargs,
).actualize()
constants = []
@@ -1928,6 +1938,8 @@ def trace_quantum_branches(op, ctx, device, trace, qrp, **kwargs) -> QRegPromise
out_expanded_tracers=out_expanded_classical_tracers,
branch_jaxprs=branch_jaxprs,
num_implicit_outputs=num_implicit_outputs[0],
- )
+ ),
+ num_device_wires=qrp.num_device_wires,
)
+ qrp2.inherit_dynamic_wire_state(qrp)
return qrp2
diff --git a/frontend/catalyst/api_extensions/quantum_operators.py b/frontend/catalyst/api_extensions/quantum_operators.py
index 959f939bba..6b4f5ce5b8 100644
--- a/frontend/catalyst/api_extensions/quantum_operators.py
+++ b/frontend/catalyst/api_extensions/quantum_operators.py
@@ -418,6 +418,29 @@ def __init__(
# pylint: disable=too-many-arguments
def trace_quantum(self, ctx, device, trace, qrp, postselect_mode=None) -> QRegPromise:
qubit = qrp.extract(self.wires)[0]
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_measure_p,
+ )
+
+ kwargs = {}
+ if postselect_mode != "hw-like":
+ kwargs["postselect"] = self.postselect
+
+ original_binder = self.binder
+ self.binder = qref_measure_p.bind
+ try:
+ self.bind_overwrite_classical_tracers(
+ trace,
+ in_expanded_tracers=[qubit],
+ out_expanded_tracers=self.out_classical_tracers,
+ num_quantum_outs=0,
+ **kwargs,
+ )
+ finally:
+ self.binder = original_binder
+ return qrp
+
if postselect_mode == "hw-like":
qubit2 = self.bind_overwrite_classical_tracers(
trace,
@@ -460,6 +483,26 @@ def __init__(
def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
qubits = qrp.extract(self.wires)
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_pauli_measure_p,
+ )
+
+ original_binder = self.binder
+ self.binder = qref_pauli_measure_p.bind
+ try:
+ self.bind_overwrite_classical_tracers(
+ trace,
+ in_expanded_tracers=qubits,
+ out_expanded_tracers=self.out_classical_tracers,
+ num_quantum_outs=0,
+ pauli_word=self.pauli_word,
+ qubits_len=len(qubits),
+ )
+ finally:
+ self.binder = original_binder
+ return qrp
+
out_qubits = self.bind_overwrite_classical_tracers(
trace,
in_expanded_tracers=qubits,
@@ -594,7 +637,7 @@ def trace_quantum(self, ctx, device, _trace, qrp, **kwargs) -> QRegPromise:
partial(body_trace.new_arg, source_info=current_source_info()), [AbstractQreg()]
)[0]
qrp_out = trace_quantum_operations(
- body_tape, device, qreg_in, ctx, body_trace, **kwargs
+ body_tape, device, qreg_in, ctx, body_trace, parent_qrp=qrp, **kwargs
)
qreg_out = qrp_out.actualize()
body_jaxpr, _, body_consts = body_trace.frame.to_jaxpr2(
@@ -608,7 +651,8 @@ def trace_quantum(self, ctx, device, _trace, qrp, **kwargs) -> QRegPromise:
args_tree=args_tree,
jaxpr=ClosedJaxpr(convert_constvars_jaxpr(body_jaxpr), ()),
)
- qrp2 = QRegPromise(op_results[-1])
+ qrp2 = QRegPromise(op_results[-1], num_device_wires=qrp.num_device_wires)
+ qrp2.inherit_dynamic_wire_state(qrp)
return qrp2
@property
diff --git a/frontend/catalyst/jax_extras/tracing.py b/frontend/catalyst/jax_extras/tracing.py
index c34ca37a5e..728d882b81 100644
--- a/frontend/catalyst/jax_extras/tracing.py
+++ b/frontend/catalyst/jax_extras/tracing.py
@@ -208,9 +208,24 @@ def __lt__(self, other):
return self.id < other.id
+def _eqn_invar_counts(eqn: JaxprEqn) -> Set[int]:
+ """Return the set of JAX variable ids used as inputs to an equation."""
+ counts: Set[int] = set()
+ for v in eqn.invars:
+ if hasattr(v, "val") and isinstance(v.val, jax._src.core.Var):
+ actual_var = v.val
+ elif isinstance(v, jax._src.core.Var):
+ actual_var = v
+ else:
+ continue
+ counts.add(actual_var.count)
+ return counts
+
+
def sort_eqns(
eqns: List[JaxprEqn | Callable[[], JaxprEqn]],
forced_order_primitives: Set[JaxprPrimitive],
+ dealloc_after_use_primitives: Set[JaxprPrimitive] | None = None,
) -> List[JaxprEqn | Callable[[], JaxprEqn]]:
"""Topologically sort TracingEqns in a unsorted list of equations, based on their
input/output variables and additional criterias."""
@@ -253,6 +268,17 @@ def sort_eqns(
for b in boxes[i + 1 :]:
b.parents.append(q) # [3]
+ if dealloc_after_use_primitives:
+ for b in boxes:
+ if b.e.primitive not in dealloc_after_use_primitives:
+ continue
+ dealloc_invars = _eqn_invar_counts(b.e)
+ for other in boxes:
+ if other is b:
+ continue
+ if dealloc_invars & _eqn_invar_counts(other.e):
+ b.parents.append(other)
+
sorted_boxes = stable_toposort(boxes)
# Restore the original callables/wrappers for the sorted equations
diff --git a/frontend/catalyst/jax_tracer.py b/frontend/catalyst/jax_tracer.py
index a372684def..4ce619296f 100644
--- a/frontend/catalyst/jax_tracer.py
+++ b/frontend/catalyst/jax_tracer.py
@@ -118,6 +118,15 @@
from catalyst.utils.exceptions import CompileError
from catalyst.utils.patching import DictPatchWrapper, Patcher
+
+def _uses_compbasis_obs(obs_tracer: DynamicJaxprTracer) -> bool:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_compbasis_p,
+ )
+
+ return obs_tracer.primitive in (compbasis_p, qref_compbasis_p)
+
+
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())
@@ -403,9 +412,36 @@ class QRegPromise:
the insertions in order to re-use qubits later thus skipping the extractions."""
@debug_logger_init
- def __init__(self, qreg: DynamicJaxprTracer):
+ def __init__(self, qreg: DynamicJaxprTracer, num_device_wires: Optional[int] = None):
self.base: DynamicJaxprTracer = qreg
+ self.num_device_wires = num_device_wires
self.cache: Dict[Any, DynamicJaxprTracer] = {}
+ self.dynamic_wire_to_qubit: Dict[Any, DynamicJaxprTracer] = {}
+ self.qref_mode = False
+ self.qref_device_reg: Optional[DynamicJaxprTracer] = None
+
+ @debug_logger
+ def inherit_dynamic_wire_state(self, parent: "QRegPromise") -> None:
+ """Inherit dynamic wire mappings from an outer quantum tracing scope."""
+ if parent.qref_mode:
+ self.qref_mode = True
+ self.qref_device_reg = parent.qref_device_reg
+ self.dynamic_wire_to_qubit = parent.dynamic_wire_to_qubit
+ self.dynamic_wire_to_qreg = parent.dynamic_wire_to_qreg
+
+ @debug_logger
+ def enable_qref_mode(self) -> None:
+ """Switch to the qref primitive path for dynamic wire support."""
+ if self.qref_mode:
+ return
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_alloc_p,
+ )
+
+ if self.num_device_wires is None:
+ raise CompileError("Cannot use dynamic wire allocation without a static device size.")
+ self.qref_mode = True
+ self.qref_device_reg = qref_alloc_p.bind(static_num_qubits=self.num_device_wires)
@debug_logger
def extract(self, wires: List[Any], allow_reuse=False) -> List[DynamicJaxprTracer]:
@@ -419,7 +455,17 @@ def extract(self, wires: List[Any], allow_reuse=False) -> List[DynamicJaxprTrace
qrp.actualize()
qubits = []
for w in wires:
- if w in qrp.cache:
+ if isinstance(w, qp.wires.DynamicWire):
+ if w not in qrp.dynamic_wire_to_qubit:
+ raise CompileError(f"Dynamic wire {w} was used before being allocated.")
+ qubits.append(qrp.dynamic_wire_to_qubit[w])
+ elif qrp.qref_mode and w not in qrp.cache:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_get_p,
+ )
+
+ qubits.append(qref_get_p.bind(qrp.qref_device_reg, w))
+ elif w in qrp.cache:
qubit = qrp.cache[w]
assert (
qubit is not None
@@ -455,6 +501,45 @@ def actualize(self) -> DynamicJaxprTracer:
return qreg
+def _trace_allocate(qrp: QRegPromise, op: qp.allocation.Allocate) -> QRegPromise:
+ """Trace a dynamic wire allocation onto the tape."""
+ state = op.hyperparameters["state"]
+ restored = op.hyperparameters["restored"]
+ qrp.enable_qref_mode()
+ qubits = qp.allocation.allocate_prim.bind(
+ num_wires=len(op.wires),
+ state=state,
+ restored=restored,
+ )
+ for wire, qubit in zip(op.wires, qubits):
+ qrp.dynamic_wire_to_qubit[wire] = qubit
+ return qrp
+
+
+def _tape_uses_dynamic_wires(tape: QuantumTape) -> bool:
+ """Return whether a tape contains dynamic wire allocation or deallocation."""
+ for op in tape.operations:
+ if isinstance(op, (qp.allocation.Allocate, qp.allocation.Deallocate)):
+ return True
+ if isinstance(op, HybridOp):
+ for region in op.regions:
+ if region.quantum_tape and _tape_uses_dynamic_wires(region.quantum_tape):
+ return True
+ return False
+
+
+def _trace_deallocate(qrp: QRegPromise, op: qp.allocation.Deallocate) -> QRegPromise:
+ """Trace a dynamic wire deallocation from the tape."""
+ qubits = []
+ for wire in op.wires:
+ if wire not in qrp.dynamic_wire_to_qubit:
+ raise CompileError(f"Attempted to deallocate unknown dynamic wire {wire}.")
+ qubits.append(qrp.dynamic_wire_to_qubit.pop(wire))
+
+ qp.allocation.deallocate_prim.bind(*qubits)
+ return qrp
+
+
@dataclass
class HybridOpRegion:
"""A code region of a nested HybridOp operation containing a JAX trace manager, a quantum tape,
@@ -539,19 +624,22 @@ def bind_overwrite_classical_tracers(
as-is.
"""
assert self.binder is not None, "HybridOp should set a binder"
- assert num_quantum_outs >= 1, "num_quantum_outs must be at least 1"
+ assert num_quantum_outs >= 0, "num_quantum_outs must be non-negative"
# Here, we are binding any of the possible hybrid ops.
# which includes: for_loop, while_loop, cond, measure.
# This will place an equation at the very end of the current list of equations.
bind_outs = self.binder(*in_expanded_tracers, **kwargs)
- out_quantum_tracers = list(bind_outs[-num_quantum_outs:])
+ out_quantum_tracers = list(bind_outs[-num_quantum_outs:]) if num_quantum_outs else []
trace = EvaluationContext.get_current_trace()
eqn = _get_eqn_from_tracing_eqn(trace.frame.tracing_eqns[-1])
frame = trace.frame
- assert len(eqn.outvars[:-num_quantum_outs]) == len(
+ classical_outvars = (
+ eqn.outvars if num_quantum_outs == 0 else eqn.outvars[:-num_quantum_outs]
+ )
+ assert len(classical_outvars) == len(
out_expanded_tracers
), f"{eqn.outvars=}\n{out_expanded_tracers=}"
@@ -622,6 +710,8 @@ def bind_overwrite_classical_tracers(
# Now, the output variables can be considered as part of the current frame.
# This allows us to avoid importing all equations again next time.
jaxpr_variables.update(eqn.outvars)
+ if num_quantum_outs == 0:
+ return None
return out_quantum_tracers[0] if num_quantum_outs == 1 else out_quantum_tracers
@debug_logger
@@ -820,6 +910,7 @@ def trace_quantum_operations(
trace: DynamicJaxprTrace,
mcm_config: qp.devices.MCMConfig = qp.devices.MCMConfig(),
out_snapshot_tracer=None,
+ parent_qrp: Optional[QRegPromise] = None,
) -> QRegPromise:
"""Recursively trace ``quantum_tape``'s operations containing both PennyLane original and
Catalyst extension operations. Produce ``QRegPromise`` object holding the resulting quantum
@@ -913,20 +1004,42 @@ def bind_native_operation(qrp, op, controlled_wires, controlled_values, adjoint=
params = op.parameters
pcphase_dim = op.hyperparameters["dimension"][0] if isinstance(op, qp.PCPhase) else None
- qubits2 = qinst_p.bind(
- *[*qubits, *params, *controlled_qubits, *controlled_values],
- op=op.name,
- qubits_len=len(qubits),
- params_len=len(params),
- ctrl_len=len(controlled_qubits),
- adjoint=adjoint,
- pcphase_dim=pcphase_dim,
- )
- qrp.insert(op.wires, qubits2[: len(qubits)])
- qrp.insert(controlled_wires, qubits2[len(qubits) :])
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_qinst_p,
+ )
+
+ qref_qinst_p.bind(
+ *[*qubits, *params, *controlled_qubits, *controlled_values],
+ op=op.name,
+ qubits_len=len(qubits),
+ params_len=len(params),
+ ctrl_len=len(controlled_qubits),
+ adjoint=adjoint,
+ pcphase_dim=pcphase_dim,
+ )
+ else:
+ qubits2 = qinst_p.bind(
+ *[*qubits, *params, *controlled_qubits, *controlled_values],
+ op=op.name,
+ qubits_len=len(qubits),
+ params_len=len(params),
+ ctrl_len=len(controlled_qubits),
+ adjoint=adjoint,
+ pcphase_dim=pcphase_dim,
+ )
+ qrp.insert(op.wires, qubits2[: len(qubits)])
+ qrp.insert(controlled_wires, qubits2[len(qubits) :])
return qrp
- qrp = QRegPromise(qreg)
+ num_device_wires = (
+ len(device.wires)
+ if device.wires and not catalyst.device.qjit_device.is_dynamic_wires(device.wires)
+ else None
+ )
+ qrp = QRegPromise(qreg, num_device_wires=num_device_wires)
+ if parent_qrp is not None:
+ qrp.inherit_dynamic_wire_state(parent_qrp)
if isinstance(device, qp.devices.LegacyDevice):
# Old device API expands tapes here. Note: this way some ops might bypass the verification.
@@ -935,6 +1048,9 @@ def bind_native_operation(qrp, op, controlled_wires, controlled_values, adjoint=
else:
ops = quantum_tape.operations
+ if _tape_uses_dynamic_wires(quantum_tape):
+ qrp.enable_qref_mode()
+
for op in ops:
qrp2 = None
if isinstance(op, HybridOp):
@@ -945,6 +1061,10 @@ def bind_native_operation(qrp, op, controlled_wires, controlled_values, adjoint=
else:
kwargs = {"mcm_config": mcm_config} # propagate
qrp2 = op.trace_quantum(ctx, device, trace, qrp, **kwargs)
+ elif isinstance(op, qp.allocation.Allocate):
+ qrp2 = _trace_allocate(qrp, op)
+ elif isinstance(op, qp.allocation.Deallocate):
+ qrp2 = _trace_deallocate(qrp, op)
elif isinstance(op, MeasurementProcess):
qrp2 = qrp
else:
@@ -953,7 +1073,25 @@ def bind_native_operation(qrp, op, controlled_wires, controlled_values, adjoint=
assert qrp2 is not None
qrp = qrp2
trace = EvaluationContext.get_current_trace()
- trace.frame.tracing_eqns = sort_eqns(trace.frame.tracing_eqns, FORCED_ORDER_PRIMITIVES) # [1]
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_dealloc_p,
+ )
+ from catalyst.jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qdealloc_p,
+ qdealloc_qb_p,
+ )
+
+ defer_dealloc_primitives = {
+ qref_dealloc_p,
+ qp.allocation.deallocate_prim,
+ qdealloc_p,
+ qdealloc_qb_p,
+ }
+ trace.frame.tracing_eqns = sort_eqns(
+ trace.frame.tracing_eqns,
+ FORCED_ORDER_PRIMITIVES,
+ defer_dealloc_primitives,
+ ) # [1]
return qrp
@@ -994,12 +1132,27 @@ def trace_observables(
# If measuring all wires on the device, pass in the qreg to compbasis op
# TODO: "all wires on the device" is None when number of wires is static,
# but a tracer when dynamic. Update to handle dynamic case.
- qreg_out = qrp.actualize()
- obs_tracers = compbasis_p.bind(qreg_out, qreg_available=True)
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_compbasis_p,
+ )
+
+ obs_tracers = qref_compbasis_p.bind(qrp.qref_device_reg, qreg_available=True)
+ else:
+ qreg_out = qrp.actualize()
+ obs_tracers = compbasis_p.bind(qreg_out, qreg_available=True)
else:
qubits = qrp.extract(wires, allow_reuse=True)
- obs_tracers = compbasis_p.bind(*qubits)
- insert_observable_wires_into_cache(qrp, wires, qubits)
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_compbasis_p,
+ )
+
+ obs_tracers = qref_compbasis_p.bind(*qubits)
+ else:
+ obs_tracers = compbasis_p.bind(*qubits)
+ if not qrp.qref_mode:
+ insert_observable_wires_into_cache(qrp, wires, qubits)
elif isinstance(obs, KNOWN_NAMED_OBS):
# The MLIR NamedObs operation only takes in a single qubit value, instead of a variadic
# range.
@@ -1009,8 +1162,15 @@ def trace_observables(
# But of course we should fix this at some time.
# TODO: make the NamedObs op take in multiple qubit values
qubits = qrp.extract([wires[0]], allow_reuse=True)
- obs_tracers = namedobs_p.bind(qubits[0], kind=type(obs).__name__)
- insert_observable_wires_into_cache(qrp, [wires[0]], [qubits[0]])
+ if qrp.qref_mode:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_namedobs_p,
+ )
+
+ obs_tracers = qref_namedobs_p.bind(qubits[0], kind=type(obs).__name__)
+ else:
+ obs_tracers = namedobs_p.bind(qubits[0], kind=type(obs).__name__)
+ insert_observable_wires_into_cache(qrp, [wires[0]], [qubits[0]])
elif isinstance(obs, qp.Hermitian):
# TODO: remove once fixed upstream: https://github.com/PennyLaneAI/pennylane/issues/4263
qubits = qrp.extract(wires, allow_reuse=True)
@@ -1137,7 +1297,7 @@ def trace_quantum_measurements(
obs_tracers, nqubits = trace_observables(output.obs, qrp, m_wires)
nqubits = d_wires if nqubits is None else nqubits
- using_compbasis = obs_tracers.primitive == compbasis_p
+ using_compbasis = _uses_compbasis_obs(obs_tracers)
if (
mcm_config.mcm_method == "single-branch-statistics"
@@ -1547,11 +1707,7 @@ def _trace_classical_phase(
# Therefore we need to compute the tree with measurements as leaves and it comes
# with an extra computational cost
- if any(isinstance(wire, qp.wires.DynamicWire) for wire in quantum_tape.wires):
- msg = "qp.allocate() with qjit is only supported with program capture enabled."
- raise CompileError(msg)
-
- # 1. Recompute the original return
+ # Dynamic wire allocation is supported via qref allocation primitives.
with QueuingManager.stop_recording():
return_values = tree_unflatten(out_tree_promise(), return_values_flat)
@@ -1738,6 +1894,12 @@ def process_tape_output(tape, return_values_flat, return_values_tree):
transformed_results.append(meas_results)
# Deallocate the register and release the device after the current tape is finished.
+ if qrp_out.qref_mode and qrp_out.qref_device_reg is not None:
+ from catalyst.from_plxpr.qref_jax_primitives import ( # pylint: disable=import-outside-toplevel
+ qref_dealloc_p,
+ )
+
+ qref_dealloc_p.bind(qrp_out.qref_device_reg)
qdealloc_p.bind(qreg_out)
device_release_p.bind()
diff --git a/frontend/test/pytest/test_dynamic_qubit_allocation.py b/frontend/test/pytest/test_dynamic_qubit_allocation.py
index 51f0083220..ced799c444 100644
--- a/frontend/test/pytest/test_dynamic_qubit_allocation.py
+++ b/frontend/test/pytest/test_dynamic_qubit_allocation.py
@@ -250,6 +250,33 @@ def circuit(c):
assert np.allclose(expected, observed)
+@pytest.mark.parametrize("cond, expected", [(True, [0, 1, 0, 0]), (False, [1, 0, 0, 0])])
+def test_dynamic_wire_alloc_cond_outside_non_capture(cond, expected, backend):
+ """
+ Test passing dynamically allocated wires into qp.cond on the legacy pathway.
+ """
+
+ @qjit(capture=False)
+ @qp.qnode(qp.device(backend, wires=2))
+ def circuit(c):
+ q = qp.allocate(1)[0]
+ qp.X(q)
+
+ def true_fn():
+ qp.CNOT(wires=[q, 1])
+
+ def false_fn():
+ qp.Identity(0)
+
+ qp.cond(c, true_fn, false_fn)()
+ qp.deallocate(q)
+ return qp.probs(wires=[0, 1])
+
+ observed = circuit(cond)
+
+ assert np.allclose(expected, observed)
+
+
@pytest.mark.parametrize(
"num_iter, expected", [(3, [0, 0, 1, 0, 0, 0, 0, 0]), (4, [1, 0, 0, 0, 0, 0, 0, 0])]
)
@@ -532,6 +559,55 @@ def circuit():
assert np.allclose(observed, expected)
+def test_no_capture_zero_state(backend):
+ """Test that zero-state dynamic allocation works without capture enabled."""
+
+ @qjit
+ @qp.qnode(qp.device(backend, wires=2))
+ def circuit():
+ with qp.allocate(1) as q:
+ qp.X(q[0])
+ qp.CNOT(wires=[q[0], 0])
+ return qp.probs(wires=[0, 1])
+
+ assert np.allclose([0, 0, 1, 0], circuit())
+
+ @qjit(target="mlir")
+ @qp.qnode(qp.device(backend, wires=2))
+ def circuit_mlir():
+ with qp.allocate(1) as q:
+ qp.X(q[0])
+ qp.CNOT(wires=[q[0], 0])
+ return qp.probs(wires=[0, 1])
+
+ mlir_str = circuit_mlir.mlir
+ assert mlir_str.count("qref.dealloc") >= 2
+
+
+def test_no_capture_magic_state(backend):
+ """
+ Test that magic state allocation works without capture enabled.
+ """
+
+ @qjit
+ @qp.qnode(qp.device(backend, wires=2))
+ def baseline():
+ qp.H(1)
+ qp.T(1)
+ qp.CNOT(wires=[1, 0])
+ return qp.probs(wires=[0])
+
+ @qjit
+ @qp.qnode(qp.device(backend, wires=2))
+ def circuit():
+ q = qp.allocate(1, state="magic-T")
+ qp.CNOT(wires=[q[0], 0])
+ return qp.probs(wires=[0])
+
+ assert np.allclose(baseline(), circuit())
+
+
+@pytest.mark.parametrize("capture", (True, False))
@pytest.mark.parametrize(
("state", "prep"),
(
@@ -539,17 +615,17 @@ def circuit():
("magic-T-adj", lambda w: (qp.H(w), qp.adjoint(qp.T(w), lazy=False))),
),
)
-def test_magic_state_allocation(backend, state, prep):
+def test_magic_state_allocation(backend, capture, state, prep):
"""Test magic state allocation lowers correctly and produces expected states."""
- @qjit(capture=True)
+ @qjit(capture=capture)
@qp.qnode(qp.device(backend, wires=2))
def baseline():
prep(1)
qp.CNOT(wires=[1, 0])
return qp.probs(wires=[0])
- @qjit(capture=True)
+ @qjit(capture=capture)
@qp.qnode(qp.device(backend, wires=2))
def circuit():
q = qp.allocate(1, state=state)
@@ -624,6 +700,38 @@ def circuit():
assert np.allclose([1, 0], circuit())
+@pytest.mark.parametrize("capture", (True, False))
+def test_magic_state_pauli_measure(backend, capture):
+ """Test Pauli product measurement on a fabricated magic wire."""
+
+ @qjit(capture=capture, target="mlir")
+ @qp.qnode(qp.device(backend, wires=2))
+ def circuit():
+ qp.CNOT(wires=[0, 1])
+ magic = qp.allocate(1, state="magic-T")
+ qp.pauli_measure("ZZ", wires=[0, magic[0]])
+ return qp.expval(qp.Z(0))
+
+ mlir_str = circuit.mlir
+ assert "pbc.ref.fabricate" in mlir_str
+ assert "pauli" in mlir_str.lower()
+
+
+@pytest.mark.parametrize("capture", (True, False))
+def test_magic_state_mid_measure(backend, capture):
+ """Test mid-circuit measure on device wires after magic state allocation."""
+
+ @qjit(capture=capture, target="mlir")
+ @qp.qnode(qp.device(backend, wires=3))
+ def circuit():
+ magic = qp.allocate(1, state="magic-T")
+ qp.pauli_measure("Z", wires=[magic[0]])
+ qp.measure(2, reset=True)
+ return qp.expval(qp.Z(0))
+
+ circuit()
+
+
def test_magic_state_mlir_lowering(backend):
"""Test that magic state allocation lowers to pbc.ref.fabricate."""
diff --git a/mlir/lib/PBC/Transforms/ppr_to_ppm.cpp b/mlir/lib/PBC/Transforms/ppr_to_ppm.cpp
index d7d75e4ca6..5971166351 100644
--- a/mlir/lib/PBC/Transforms/ppr_to_ppm.cpp
+++ b/mlir/lib/PBC/Transforms/ppr_to_ppm.cpp
@@ -14,12 +14,14 @@
#define DEBUG_TYPE "ppr-to-ppm"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "PBC/Transforms/Patterns.h"
+#include "Quantum/IR/QuantumOps.h"
using namespace llvm;
using namespace mlir;
@@ -33,13 +35,48 @@ namespace pbc {
#define GEN_PASS_DEF_PPRTOPPMPASS
#include "PBC/Transforms/Passes.h.inc"
+using namespace catalyst::quantum;
+
+namespace {
+
+/// Move ``quantum.dealloc_qb`` operations to after their last use in each function.
+/// This ensures magic-state qubits remain mapped when ``ppr-to-ppm`` inserts late PPM users.
+void sinkQuantumDeallocs(ModuleOp module)
+{
+ module.walk([&](func::FuncOp func) {
+ SmallVector deallocOps;
+ func.walk([&](DeallocQubitOp deallocOp) { deallocOps.push_back(deallocOp); });
+
+ for (DeallocQubitOp deallocOp : deallocOps) {
+ mlir::Value qubit = deallocOp.getQubit();
+ Operation *lastUser = nullptr;
+ func.walk([&](Operation *op) {
+ if (op == deallocOp) {
+ return;
+ }
+ for (mlir::Value operand : op->getOperands()) {
+ if (operand == qubit) {
+ lastUser = op;
+ }
+ }
+ });
+
+ if (lastUser && lastUser != deallocOp->getPrevNode()) {
+ deallocOp->moveAfter(lastUser);
+ }
+ }
+ });
+}
+
+} // namespace
+
struct PPRToPPMPass : public impl::PPRToPPMPassBase {
using PPRToPPMPassBase::PPRToPPMPassBase;
void runOnOperation() final
{
auto ctx = &getContext();
- auto module = getOperation();
+ auto module = cast(getOperation());
RewritePatternSet non_clifford_patterns(ctx);
populateDecomposeNonCliffordPPRPatterns(non_clifford_patterns, decomposeMethod,
@@ -56,6 +93,8 @@ struct PPRToPPMPass : public impl::PPRToPPMPassBase {
if (failed(applyPatternsGreedily(module, std::move(clifford_patterns)))) {
return signalPassFailure();
}
+
+ sinkQuantumDeallocs(module);
}
};