From 205cc2fe4d768f45c757d2b259160a2057e2fb91 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 10:48:42 -0400 Subject: [PATCH 01/19] init; changelog --- doc/releases/changelog-dev.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 0e104520d..e590cd0c4 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -74,6 +74,10 @@ Array([0.25, 0.25, 0.25, 0.25], dtype=float64)) ``` +* Catalyst now supports automatic qubit management. + (Add explanation and examples here.) + [(#???)](https://github.com/PennyLaneAI/catalyst/pull/???) +

Improvements 🛠

* The behaviour of measurement processes executed on `null.qubit` with QJIT is now more in line with From cecffa70cdaff1ee919f379a8f62f2b95fdbfbb1 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 15:23:17 -0400 Subject: [PATCH 02/19] add core skeleton; add first simple test --- frontend/catalyst/device/qjit_device.py | 5 ++- frontend/catalyst/from_plxpr.py | 6 ++- frontend/catalyst/jax_primitives.py | 10 ++++- frontend/catalyst/jax_tracer.py | 25 ++++++++--- .../pytest/test_automatic_qubit_management.py | 36 ++++++++++++++++ mlir/include/Quantum/IR/QuantumOps.td | 13 ++++++ .../Quantum/Transforms/ConversionPatterns.cpp | 8 +++- runtime/include/RuntimeCAPI.h | 4 +- runtime/lib/capi/ExecutionContext.hpp | 4 ++ runtime/lib/capi/RuntimeCAPI.cpp | 42 ++++++++++++++----- 10 files changed, 130 insertions(+), 23 deletions(-) create mode 100644 frontend/test/pytest/test_automatic_qubit_management.py diff --git a/frontend/catalyst/device/qjit_device.py b/frontend/catalyst/device/qjit_device.py index 1f4b10066..85b7d3a69 100644 --- a/frontend/catalyst/device/qjit_device.py +++ b/frontend/catalyst/device/qjit_device.py @@ -548,6 +548,8 @@ def is_dynamic_wires(wires: qml.wires.Wires): If the number of wires is dynamic, the Wires object contains a single tracer that represents the number of wires. """ + # Automatic qubit management mode should not encounter this query + assert wires is not None return (len(wires) == 1) and (isinstance(wires[0], DynamicJaxprTracer)) @@ -555,7 +557,8 @@ def check_device_wires(wires): """Validate requirements Catalyst imposes on device wires.""" if wires is None: - raise AttributeError("Catalyst does not support device instances without set wires.") + # Automatic qubit management mode, nothing to check + return if len(wires) >= 2 or (not is_dynamic_wires(wires)): # A dynamic number of wires correspond to a single tracer for the number diff --git a/frontend/catalyst/from_plxpr.py b/frontend/catalyst/from_plxpr.py index ae3c67f73..01be1ad22 100644 --- a/frontend/catalyst/from_plxpr.py +++ b/frontend/catalyst/from_plxpr.py @@ -301,7 +301,11 @@ def __setattr__(self, __name: str, __value) -> None: def setup(self): """Initialize the stateref and bind the device.""" if self.stateref is None: - device_init_p.bind(self._shots, **_get_device_kwargs(self._device)) + device_init_p.bind( + self._shots, + auto_qubit_management=(self._device.wires is None), + **_get_device_kwargs(self._device), + ) self.stateref = {"qreg": qalloc_p.bind(len(self._device.wires)), "wire_map": {}} # pylint: disable=attribute-defined-outside-init diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 31bc82220..4238b28a3 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -811,12 +811,17 @@ def _zne_lowering(ctx, *args, folding, jaxpr, fn): # device_init # @device_init_p.def_abstract_eval -def _device_init_abstract_eval(shots, rtd_lib, rtd_name, rtd_kwargs): +def _device_init_abstract_eval(shots, auto_qubit_management, rtd_lib, rtd_name, rtd_kwargs): return () def _device_init_lowering( - jax_ctx: mlir.LoweringRuleContext, shots: ir.Value, rtd_lib, rtd_name, rtd_kwargs + jax_ctx: mlir.LoweringRuleContext, + shots: ir.Value, + auto_qubit_management, + rtd_lib, + rtd_name, + rtd_kwargs, ): ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True @@ -827,6 +832,7 @@ def _device_init_lowering( ir.StringAttr.get(rtd_name), ir.StringAttr.get(rtd_kwargs), shots=shots_value, + auto_qubit_management=auto_qubit_management, ) return () diff --git a/frontend/catalyst/jax_tracer.py b/frontend/catalyst/jax_tracer.py index 71da0334e..d3fcb0830 100644 --- a/frontend/catalyst/jax_tracer.py +++ b/frontend/catalyst/jax_tracer.py @@ -984,11 +984,19 @@ def trace_quantum_measurements( "Use qml.sample() instead." ) - d_wires = ( - device.wires[0] - if catalyst.device.qjit_device.is_dynamic_wires(device.wires) - else len(device.wires) - ) + # d_wires = ( + # device.wires[0] + # if catalyst.device.qjit_device.is_dynamic_wires(device.wires) + # else len(device.wires) + # ) + if device.wires is None: + # Automatic qubit management mode, TODO: what here??? + d_wires = 0 + pass + elif catalyst.device.qjit_device.is_dynamic_wires(device.wires): + d_wires = device.wires[0] + else: + d_wires = len(device.wires) m_wires = output.wires if output.wires else None obs_tracers, nqubits = trace_observables(output.obs, qrp, m_wires) nqubits = d_wires if nqubits is None else nqubits @@ -1385,11 +1393,16 @@ def is_leaf(obj): device_shots = get_device_shots(device) or 0 device_init_p.bind( device_shots, + auto_qubit_management=(device.wires is None), rtd_lib=device.backend_lib, rtd_name=device.backend_name, rtd_kwargs=str(device.backend_kwargs), ) - if catalyst.device.qjit_device.is_dynamic_wires(device.wires): + if device.wires is None: + # Automatic qubit management mode + # We start with 0 wires and allocate new wires in runtime as we encounter them. + qreg_in = qalloc_p.bind(0) + elif catalyst.device.qjit_device.is_dynamic_wires(device.wires): # When device has dynamic wires, the device.wires iterable object # has a single value, which is the tracer for the number of wires qreg_in = qalloc_p.bind(device.wires[0]) diff --git a/frontend/test/pytest/test_automatic_qubit_management.py b/frontend/test/pytest/test_automatic_qubit_management.py new file mode 100644 index 000000000..74eb87abf --- /dev/null +++ b/frontend/test/pytest/test_automatic_qubit_management.py @@ -0,0 +1,36 @@ +# Copyright 2025 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 file contains tests for automatic qubit management. +Automatic qubit management refers to when the user does not specify the total number of wires +during device initialization. +""" + +import numpy as np +import pennylane as qml +import pytest + +import catalyst +from catalyst import qjit + + +def test_probs(backend): + + def circuit(): + qml.PauliX(wires=0) + return qml.probs(wires=[2, 3]) + + ref = qjit(qml.qnode(qml.device(backend, wires=4))(circuit))() + observed = qjit(qml.qnode(qml.device(backend))(circuit))() + assert np.allclose(ref, observed) diff --git a/mlir/include/Quantum/IR/QuantumOps.td b/mlir/include/Quantum/IR/QuantumOps.td index 9efdd6628..780d63ded 100644 --- a/mlir/include/Quantum/IR/QuantumOps.td +++ b/mlir/include/Quantum/IR/QuantumOps.td @@ -81,6 +81,7 @@ def DeviceInitOp : Quantum_Op<"device"> { let arguments = (ins Optional:$shots, + UnitAttr:$auto_qubit_management, StrAttr:$lib, StrAttr:$device_name, StrAttr:$kwargs @@ -90,6 +91,18 @@ def DeviceInitOp : Quantum_Op<"device"> { (`shots` `(` $shots^ `)`)? `[` $lib `,` $device_name `,` $kwargs `]` attr-dict }]; + let builders = [ + OpBuilder< + // Convenience builder for a device not in automatic qubit management mode + (ins + "mlir::Value":$shots, + "mlir::StringAttr":$lib, + "mlir::StringAttr":$device_name, + "mlir::StringAttr":$kwargs + ),[{ + DeviceInitOp::build($_builder, $_state, shots, false, lib, device_name, kwargs); + }]>, + ]; } def DeviceReleaseOp : Quantum_Op<"device_release"> { diff --git a/mlir/lib/Quantum/Transforms/ConversionPatterns.cpp b/mlir/lib/Quantum/Transforms/ConversionPatterns.cpp index ccf03b741..6d037cff9 100644 --- a/mlir/lib/Quantum/Transforms/ConversionPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/ConversionPatterns.cpp @@ -187,11 +187,13 @@ struct DeviceInitOpPattern : public OpConversionPattern { Type charPtrType = LLVM::LLVMPointerType::get(rewriter.getContext()); Type int64Type = IntegerType::get(rewriter.getContext(), 64); + Type int1Type = IntegerType::get(rewriter.getContext(), 1); Type qirSignature = LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(ctx), {/* rtd_lib = */ charPtrType, /* rtd_name = */ charPtrType, /* rtd_kwargs = */ charPtrType, - /* shots = */ int64Type}); + /* shots = */ int64Type, + /*auto_qubit_management = */ int1Type}); LLVM::LLVMFuncOp fnDecl = catalyst::ensureFunctionDeclaration(rewriter, op, qirName, qirSignature); @@ -217,6 +219,10 @@ struct DeviceInitOpPattern : public OpConversionPattern { operands.push_back(shots); } + Value autoQubitManagement = rewriter.create(loc, rewriter.getI1Type(), + op.getAutoQubitManagement()); + operands.push_back(autoQubitManagement); + rewriter.create(loc, fnDecl, operands); rewriter.eraseOp(op); diff --git a/runtime/include/RuntimeCAPI.h b/runtime/include/RuntimeCAPI.h index 710cf3f47..33f26f03a 100644 --- a/runtime/include/RuntimeCAPI.h +++ b/runtime/include/RuntimeCAPI.h @@ -24,8 +24,8 @@ extern "C" { // Quantum Runtime Instructions void __catalyst__rt__fail_cstr(const char *); -void __catalyst__rt__initialize(uint32_t *seed); -void __catalyst__rt__device_init(int8_t *, int8_t *, int8_t *, int64_t shots); +void __catalyst__rt__initialize(uint32_t *); +void __catalyst__rt__device_init(int8_t *, int8_t *, int8_t *, int64_t, bool); void __catalyst__rt__device_release(); void __catalyst__rt__finalize(); void __catalyst__rt__toggle_recorder(bool); diff --git a/runtime/lib/capi/ExecutionContext.hpp b/runtime/lib/capi/ExecutionContext.hpp index dbe27706b..33aeaeb8b 100644 --- a/runtime/lib/capi/ExecutionContext.hpp +++ b/runtime/lib/capi/ExecutionContext.hpp @@ -166,6 +166,7 @@ class RTDevice { std::string rtd_lib; std::string rtd_name; std::string rtd_kwargs; + bool auto_qubit_management; std::unique_ptr rtd_dylib{nullptr}; std::unique_ptr rtd_qdevice{nullptr}; @@ -259,6 +260,9 @@ class RTDevice { void setDeviceStatus(RTDeviceStatus new_status) noexcept { status = new_status; } + void setDeviceAutoQubitManagementMode(bool mode) { auto_qubit_management = mode; } + bool getDeviceAutoQubitManagementMode() { return auto_qubit_management; } + [[nodiscard]] auto getDeviceStatus() const -> RTDeviceStatus { return status; } friend std::ostream &operator<<(std::ostream &os, const RTDevice &device) diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index 22f5f28c4..ce5693844 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -12,24 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include - -#include -#include - #include #include +#include #include +#include #include "mlir/ExecutionEngine/CRunnerUtils.h" #include "Exception.hpp" -#include "QuantumDevice.hpp" - #include "ExecutionContext.hpp" #include "MemRefUtils.hpp" +#include "QuantumDevice.hpp" #include "Timer.hpp" #include "RuntimeCAPI.h" @@ -243,7 +241,7 @@ void __catalyst__rt__finalize() } static int __catalyst__rt__device_init__impl(int8_t *rtd_lib, int8_t *rtd_name, int8_t *rtd_kwargs, - int64_t shots) + int64_t shots, bool auto_qubit_management) { // Device library cannot be a nullptr RT_FAIL_IF(!rtd_lib, "Invalid device library"); @@ -256,6 +254,7 @@ static int __catalyst__rt__device_init__impl(int8_t *rtd_lib, int8_t *rtd_name, (rtd_kwargs ? reinterpret_cast(rtd_kwargs) : "")}; RT_FAIL_IF(!initRTDevicePtr(args[0], args[1], args[2]), "Failed initialization of the backend device"); + RTD_PTR->setDeviceAutoQubitManagementMode(auto_qubit_management); getQuantumDevicePtr()->SetDeviceShots(shots); if (CTX->getDeviceRecorderStatus()) { getQuantumDevicePtr()->StartTapeRecording(); @@ -264,10 +263,10 @@ static int __catalyst__rt__device_init__impl(int8_t *rtd_lib, int8_t *rtd_name, } void __catalyst__rt__device_init(int8_t *rtd_lib, int8_t *rtd_name, int8_t *rtd_kwargs, - int64_t shots) + int64_t shots, bool auto_qubit_management) { timer::timer(__catalyst__rt__device_init__impl, "device_init", /* add_endl */ true, rtd_lib, - rtd_name, rtd_kwargs, shots); + rtd_name, rtd_kwargs, shots, auto_qubit_management); } static int __catalyst__rt__device_release__impl() @@ -1036,6 +1035,11 @@ int64_t __catalyst__rt__array_get_size_1d(QirArray *ptr) return qubit_vector_ptr->size(); } +static bool isDeviceAutomaticQubitManageMode() +{ + return RTD_PTR->getDeviceAutoQubitManagementMode(); +} + int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) { std::vector *qubit_vector_ptr = reinterpret_cast *>(ptr); @@ -1043,7 +1047,25 @@ int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) RT_ASSERT(idx >= 0); std::string error_msg = "The qubit register does not contain the requested wire: "; error_msg += std::to_string(idx); - RT_FAIL_IF(static_cast(idx) >= qubit_vector_ptr->size(), error_msg.c_str()); + + bool _ = isDeviceAutomaticQubitManageMode(); + std::cout << "auto management: " << _ << "\n"; + + if (static_cast(idx) >= qubit_vector_ptr->size()) { + if (!isDeviceAutomaticQubitManageMode()) { + RT_FAIL(error_msg.c_str()); + } + else { + // allocate a new qubit if we are in automatic qubit allocation mode + // `idx` is the new user wire index from frontend pennylane + // number of currently allocated qubits is `qubit_vector_ptr->size()` + while (qubit_vector_ptr->size() <= idx) { + auto new_alloced_id = + reinterpret_cast(__catalyst__rt__qubit_allocate()); + qubit_vector_ptr->push_back(new_alloced_id); + } + } + } QubitIdType *data = qubit_vector_ptr->data(); return (int8_t *)&data[idx]; From 8a12270b060041c624605124255e3b0b9af2d23d Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 16:25:46 -0400 Subject: [PATCH 03/19] add sample and counts test --- .../pytest/test_automatic_qubit_management.py | 44 +++++++++++++++++-- 1 file changed, 40 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/test_automatic_qubit_management.py b/frontend/test/pytest/test_automatic_qubit_management.py index 74eb87abf..c7ab59923 100644 --- a/frontend/test/pytest/test_automatic_qubit_management.py +++ b/frontend/test/pytest/test_automatic_qubit_management.py @@ -11,10 +11,21 @@ # 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 file contains tests for automatic qubit management. Automatic qubit management refers to when the user does not specify the total number of wires during device initialization. + +Note that, catalyst and pennylane handles device labels differently. For example, when a new label + qml.gate_or_measurement(wires=1000) +is encountered, core pennylane considers "1000" as a pure wire label, and interprets that as +*one* new wire, with the label "1000". However in catalyst we do not associate wires with +arbitrary labels and require wires to be continuous integers from zero, and we would interpret this +as "allocate new wires until we have 1001 wires, and act on wire[1000]". + +In other words, the reference runs of automatic qubit management should be qjit runs with wires +specified during device initialization, instead of non qjit runs. """ import numpy as np @@ -25,12 +36,37 @@ from catalyst import qjit -def test_probs(backend): +def test_partial_sample(backend): + + def circuit(): + qml.RX(0.0, wires=0) + return qml.sample(wires=[0, 2]) + + wires = [4, None] + devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert np.allclose(ref, observed) + + +def test_partial_counts(backend): + + def circuit(): + qml.RX(0.0, wires=0) + return qml.counts(wires=[0, 2]) + + wires = [4, None] + devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert np.allclose(ref, observed) + + +def test_partial_probs(backend): def circuit(): qml.PauliX(wires=0) - return qml.probs(wires=[2, 3]) + return qml.probs(wires=[0, 2]) - ref = qjit(qml.qnode(qml.device(backend, wires=4))(circuit))() - observed = qjit(qml.qnode(qml.device(backend))(circuit))() + wires = [4, None] + devices = [qml.device(backend, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) assert np.allclose(ref, observed) From a1e65be4ba46b02e9880d3debb5ac53babae6bf1 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 16:38:02 -0400 Subject: [PATCH 04/19] update lit test for new auto qubit manage mode unit attr --- mlir/test/Quantum/ConversionTest.mlir | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/mlir/test/Quantum/ConversionTest.mlir b/mlir/test/Quantum/ConversionTest.mlir index 035f0147d..0006c132b 100644 --- a/mlir/test/Quantum/ConversionTest.mlir +++ b/mlir/test/Quantum/ConversionTest.mlir @@ -44,7 +44,7 @@ func.func @finalize() { // ----- -// CHECK: llvm.func @__catalyst__rt__device_init(!llvm.ptr, !llvm.ptr, !llvm.ptr, i64) +// CHECK: llvm.func @__catalyst__rt__device_init(!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i1) // CHECK-LABEL: @device func.func @device() { @@ -61,7 +61,8 @@ func.func @device() { // CHECK: [[b1:%.+]] = llvm.getelementptr inbounds [[bo]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<16 x i8> // CHECK: [[d3:%.+]] = llvm.mlir.addressof @"{my_attr: my_attr_value}" : !llvm.ptr // CHECK: [[d4:%.+]] = llvm.getelementptr inbounds [[d3]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<25 x i8> - // CHECK: llvm.call @__catalyst__rt__device_init([[d1]], [[b1]], [[d4]], [[shots]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64) -> () + // CHECK: [[false:%.+]] = llvm.mlir.constant(false) : i1 + // CHECK: llvm.call @__catalyst__rt__device_init([[d1]], [[b1]], [[d4]], [[shots]], [[false]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i1) -> () %shots = llvm.mlir.constant(1000 : i64) : i64 quantum.device shots(%shots) ["rtd_lightning.so", "lightning.qubit", "{my_attr: my_attr_value}"] @@ -71,7 +72,8 @@ func.func @device() { // CHECK: [[e3:%.+]] = llvm.getelementptr inbounds [[e2]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<17 x i8> // CHECK: [[e4:%.+]] = llvm.mlir.addressof @"{my_other_attr: my_other_attr_value}" : !llvm.ptr // CHECK: [[e5:%.+]] = llvm.getelementptr inbounds [[e4]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<37 x i8> - // CHECK: llvm.call @__catalyst__rt__device_init([[e1]], [[e3]], [[e5]], [[shots]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64) -> () + // CHECK: [[false:%.+]] = llvm.mlir.constant(false) : i1 + // CHECK: llvm.call @__catalyst__rt__device_init([[e1]], [[e3]], [[e5]], [[shots]], [[false]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i1) -> () quantum.device shots(%shots) ["rtd_lightning.so", "lightning.kokkos", "{my_other_attr: my_other_attr_value}"] @@ -82,9 +84,14 @@ func.func @device() { // CHECK: [[d3:%.+]] = llvm.mlir.addressof @"{my_noshots_attr: my_noshots_attr_value}" : !llvm.ptr // CHECK: [[d4:%.+]] = llvm.getelementptr inbounds [[d3]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<41 x i8> // CHECK: [[zero_shots:%.+]] = llvm.mlir.constant(0 : i64) : i64 - // CHECK: llvm.call @__catalyst__rt__device_init([[d1]], [[b1]], [[d4]], [[zero_shots]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64) -> () + // CHECK: [[false:%.+]] = llvm.mlir.constant(false) : i1 + // CHECK: llvm.call @__catalyst__rt__device_init([[d1]], [[b1]], [[d4]], [[zero_shots]], [[false]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i1) -> () quantum.device ["rtd_lightning.so", "lightning.qubit", "{my_noshots_attr: my_noshots_attr_value}"] + // CHECK: [[true:%.+]] = llvm.mlir.constant(true) : i1 + // CHECK: llvm.call @__catalyst__rt__device_init({{%.+}}, {{%.+}}, {{%.+}}, {{%.+}}, [[true]]) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i1) -> () + quantum.device ["blah.so", "blah.qubit", ""] {auto_qubit_management} + return } From 8ede25f741e1dc00c6af8eb84bc6faf171eefc35 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 16:41:37 -0400 Subject: [PATCH 05/19] update runtime tests for new device init argument --- runtime/tests/Test_MBQC.cpp | 2 +- runtime/tests/Test_NullQubit.cpp | 8 ++++---- runtime/tests/Test_OpenQasmDevice.cpp | 14 ++++++++------ 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/runtime/tests/Test_MBQC.cpp b/runtime/tests/Test_MBQC.cpp index 8300794ec..a6bab019a 100644 --- a/runtime/tests/Test_MBQC.cpp +++ b/runtime/tests/Test_MBQC.cpp @@ -26,7 +26,7 @@ TEST_CASE("Test __catalyst__mbqc__measure_in_basis, device=null.qubit", "[MBQC]" __catalyst__rt__initialize(nullptr); const std::string rtd_name{"null.qubit"}; - __catalyst__rt__device_init((int8_t *)rtd_name.c_str(), nullptr, nullptr, 0); + __catalyst__rt__device_init((int8_t *)rtd_name.c_str(), nullptr, nullptr, 0, false); const size_t num_qubits = 1; QirArray *qs = __catalyst__rt__qubit_allocate_array(num_qubits); diff --git a/runtime/tests/Test_NullQubit.cpp b/runtime/tests/Test_NullQubit.cpp index 790751fe6..e1b3838af 100644 --- a/runtime/tests/Test_NullQubit.cpp +++ b/runtime/tests/Test_NullQubit.cpp @@ -43,7 +43,7 @@ TEST_CASE("Test __catalyst__rt__device_init registering device=null.qubit", "[Nu __catalyst__rt__initialize(nullptr); char rtd_name[11] = "null.qubit"; - __catalyst__rt__device_init((int8_t *)rtd_name, nullptr, nullptr, 0); + __catalyst__rt__device_init((int8_t *)rtd_name, nullptr, nullptr, 0, false); __catalyst__rt__device_release(); @@ -177,7 +177,7 @@ TEST_CASE("Test __catalyst__qis__Sample with num_qubits=2 and PartialSample call std::array{"null.qubit", "null_qubit", ""}; __catalyst__rt__initialize(nullptr); __catalyst__rt__device_init((int8_t *)rtd_lib.c_str(), (int8_t *)rtd_name.c_str(), - (int8_t *)rtd_kwargs.c_str(), 1000); + (int8_t *)rtd_kwargs.c_str(), 1000, false); QirArray *qs = __catalyst__rt__qubit_allocate_array(2); @@ -309,7 +309,7 @@ TEST_CASE("Test __catalyst__qis__Gradient_params Op=[Hadamard,RZ,RY,RZ,S,T,Param std::array{"null.qubit", "null_qubit", ""}; __catalyst__rt__device_init((int8_t *)rtd_lib.c_str(), (int8_t *)rtd_name.c_str(), - (int8_t *)rtd_kwargs.c_str(), 0); + (int8_t *)rtd_kwargs.c_str(), 0, false); QUBIT *q0 = __catalyst__rt__qubit_allocate(); QUBIT *q1 = __catalyst__rt__qubit_allocate(); @@ -372,7 +372,7 @@ TEST_CASE("Test __catalyst__rt__print_state", "[NullQubit]") auto [rtd_lib, rtd_name, rtd_kwargs] = std::array{"null.qubit", "null_qubit", ""}; __catalyst__rt__device_init((int8_t *)rtd_lib.c_str(), (int8_t *)rtd_name.c_str(), - (int8_t *)rtd_kwargs.c_str(), 0); + (int8_t *)rtd_kwargs.c_str(), 0, false); QirArray *qs = __catalyst__rt__qubit_allocate_array(2); diff --git a/runtime/tests/Test_OpenQasmDevice.cpp b/runtime/tests/Test_OpenQasmDevice.cpp index 6195a9e96..1009af47c 100644 --- a/runtime/tests/Test_OpenQasmDevice.cpp +++ b/runtime/tests/Test_OpenQasmDevice.cpp @@ -660,10 +660,11 @@ TEST_CASE("Test __catalyst__rt__device_init registering the OpenQasm device", "[ char device_aws[30] = "braket.aws.qubit"; #if __has_include("OpenQasmDevice.hpp") - __catalyst__rt__device_init((int8_t *)device_aws, nullptr, nullptr, 0); + __catalyst__rt__device_init((int8_t *)device_aws, nullptr, nullptr, 0, false); #else - REQUIRE_THROWS_WITH(__catalyst__rt__device_init((int8_t *)device_aws, nullptr, nullptr, 0), - ContainsSubstring("cannot open shared object file")); + REQUIRE_THROWS_WITH( + __catalyst__rt__device_init((int8_t *)device_aws, nullptr, nullptr, 0, false), + ContainsSubstring("cannot open shared object file")); #endif __catalyst__rt__finalize(); @@ -673,10 +674,11 @@ TEST_CASE("Test __catalyst__rt__device_init registering the OpenQasm device", "[ char device_local[30] = "braket.local.qubit"; #if __has_include("OpenQasmDevice.hpp") - __catalyst__rt__device_init((int8_t *)device_local, nullptr, nullptr, 0); + __catalyst__rt__device_init((int8_t *)device_local, nullptr, nullptr, 0, false); #else - REQUIRE_THROWS_WITH(__catalyst__rt__device_init((int8_t *)(int8_t *), nullptr, nullptr, 0), - ContainsSubstring("cannot open shared object file")); + REQUIRE_THROWS_WITH( + __catalyst__rt__device_init((int8_t *)(int8_t *), nullptr, nullptr, 0, false), + ContainsSubstring("cannot open shared object file")); #endif __catalyst__rt__finalize(); From 655f5d1f0f7f0bae8c73d4a74411bc2a7f3839de Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 16:46:00 -0400 Subject: [PATCH 06/19] remove some prints --- runtime/lib/capi/RuntimeCAPI.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index ce5693844..e052c4588 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -1048,9 +1048,6 @@ int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) std::string error_msg = "The qubit register does not contain the requested wire: "; error_msg += std::to_string(idx); - bool _ = isDeviceAutomaticQubitManageMode(); - std::cout << "auto management: " << _ << "\n"; - if (static_cast(idx) >= qubit_vector_ptr->size()) { if (!isDeviceAutomaticQubitManageMode()) { RT_FAIL(error_msg.c_str()); From a135e45457466a3f2585c0270352118406db58de Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 16:48:58 -0400 Subject: [PATCH 07/19] remove "catalyst does not support device with no wires" pytest --- frontend/test/pytest/test_device_api.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/frontend/test/pytest/test_device_api.py b/frontend/test/pytest/test_device_api.py index b8b760080..56103aa3b 100644 --- a/frontend/test/pytest/test_device_api.py +++ b/frontend/test/pytest/test_device_api.py @@ -56,17 +56,6 @@ def test_qjit_device(): device_qjit.execute(10, 2) -def test_qjit_device_no_wires(): - """Test the qjit device from a device using the new api without wires set.""" - device = NullQubit(shots=2032) - - with pytest.raises( - AttributeError, match="Catalyst does not support device instances without set wires." - ): - # Create qjit device - QJITDevice(device) - - @pytest.mark.parametrize( "wires", ( From 213326387879171a872124aaac7fd7e265fa1bd2 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 5 Jun 2025 17:45:15 -0400 Subject: [PATCH 08/19] clang warning about comparing between signed and unsigned --- runtime/lib/capi/RuntimeCAPI.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index e052c4588..6fec55258 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -1056,7 +1056,7 @@ int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) // allocate a new qubit if we are in automatic qubit allocation mode // `idx` is the new user wire index from frontend pennylane // number of currently allocated qubits is `qubit_vector_ptr->size()` - while (qubit_vector_ptr->size() <= idx) { + while (qubit_vector_ptr->size() <= static_cast(idx)) { auto new_alloced_id = reinterpret_cast(__catalyst__rt__qubit_allocate()); qubit_vector_ptr->push_back(new_alloced_id); From a7963bffa6688c6e5e4c5dd9dca015ee798efe4c Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 6 Jun 2025 15:54:21 -0400 Subject: [PATCH 09/19] add tests for non partial measurements --- .../pytest/test_automatic_qubit_management.py | 57 ++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/test_automatic_qubit_management.py b/frontend/test/pytest/test_automatic_qubit_management.py index c7ab59923..4324f982f 100644 --- a/frontend/test/pytest/test_automatic_qubit_management.py +++ b/frontend/test/pytest/test_automatic_qubit_management.py @@ -30,9 +30,7 @@ import numpy as np import pennylane as qml -import pytest -import catalyst from catalyst import qjit @@ -45,6 +43,7 @@ def circuit(): wires = [4, None] devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert ref.shape == observed.shape assert np.allclose(ref, observed) @@ -57,6 +56,7 @@ def circuit(): wires = [4, None] devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert (ref[i].shape == observed[i].shape for i in (0, 1)) assert np.allclose(ref, observed) @@ -69,4 +69,57 @@ def circuit(): wires = [4, None] devices = [qml.device(backend, wires=wire) for wire in wires] ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert ref.shape == observed.shape + assert np.allclose(ref, observed) + + +def test_sample(backend): + + def circuit(): + qml.RX(0.0, wires=3) + return qml.sample() + + wires = [4, None] + devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert ref.shape == observed.shape + assert np.allclose(ref, observed) + + +def test_counts(backend): + + def circuit(): + qml.RX(0.0, wires=3) + return qml.counts() + + wires = [4, None] + devices = [qml.device(backend, shots=10, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert (ref[i].shape == observed[i].shape for i in (0, 1)) + assert np.allclose(ref, observed) + + +def test_probs(backend): + + def circuit(): + qml.PauliX(wires=3) + return qml.probs() + + wires = [4, None] + devices = [qml.device(backend, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert ref.shape == observed.shape + assert np.allclose(ref, observed) + + +def test_state(backend): + + def circuit(): + qml.PauliX(wires=3) + return qml.state() + + wires = [4, None] + devices = [qml.device(backend, wires=wire) for wire in wires] + ref, observed = (qjit(qml.qnode(dev)(circuit))() for dev in devices) + assert ref.shape == observed.shape assert np.allclose(ref, observed) From 328ddcc04b97ce9679118f3df9679ac496b284dc Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 6 Jun 2025 15:56:40 -0400 Subject: [PATCH 10/19] update oqd runtime test's device init call --- runtime/tests/Test_OQDDevice.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtime/tests/Test_OQDDevice.cpp b/runtime/tests/Test_OQDDevice.cpp index f1e6fa308..22903d280 100644 --- a/runtime/tests/Test_OQDDevice.cpp +++ b/runtime/tests/Test_OQDDevice.cpp @@ -928,7 +928,7 @@ TEST_CASE("Test OpenAPL Program generation", "[oqd]") __catalyst__rt__initialize(nullptr); __catalyst__rt__device_init((int8_t *)rtd_lib.c_str(), (int8_t *)rtd_name.c_str(), - (int8_t *)rtd_kwargs.c_str(), 1000); + (int8_t *)rtd_kwargs.c_str(), 1000, false); QirArray *qs = __catalyst__rt__qubit_allocate_array(num_qubits); From 5881c987f8d597ba557377a141aa1f82f32d1734 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 6 Jun 2025 16:06:08 -0400 Subject: [PATCH 11/19] changelog --- doc/releases/changelog-dev.md | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index bf865823d..37018e36a 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -75,8 +75,34 @@ ``` * Catalyst now supports automatic qubit management. - (Add explanation and examples here.) - [(#???)](https://github.com/PennyLaneAI/catalyst/pull/???) + [(#1788)](https://github.com/PennyLaneAI/catalyst/pull/1788) + + The number of wires does not need to be speficied during device initialization, + and can be automatically allocated whenever a new wire is encountered. + + ```python + @qjit + def workflow(): + dev = qml.device("lightning.qubit") # no wires here! + @qml.qnode(dev) + def circuit(): + qml.PauliX(wires=2) + return qml.probs() + return circuit() + + print(workflow()) + ``` + + ```pycon + [0. 1. 0. 0. 0. 0. 0. 0.] + ``` + + In this example, the number of wires is not specified at device initialization. + When we encounter an X gate on `wires=2`, catalyst automatically allocates 3 wires, and applies + an X gate to the wire with index `2`. The result is the state |001>, giving the probabilities + above. + + This feature can be turned on by simply not supplying a `wires` argument to the device.

Improvements 🛠

From 09db6d5348e66151788d553923e66c8c79c86fd6 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 6 Jun 2025 16:30:33 -0400 Subject: [PATCH 12/19] typo --- frontend/catalyst/jax_tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/jax_tracer.py b/frontend/catalyst/jax_tracer.py index 7f6530885..01a0e39aa 100644 --- a/frontend/catalyst/jax_tracer.py +++ b/frontend/catalyst/jax_tracer.py @@ -988,7 +988,7 @@ def trace_quantum_measurements( if device.wires is None: d_wires = num_qubits_p.bind() elif catalyst.device.qjit_device.is_dynamic_wires(device.wires): - d_wires = device.wires[0] + d_wires = num_qubits_p.bind() else: d_wires = len(device.wires) From d61f6aa3e036afe0b83bccf98052336bff5b2925 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 09:46:30 -0400 Subject: [PATCH 13/19] codefactor --- frontend/catalyst/jax_primitives.py | 1 + .../pytest/test_automatic_qubit_management.py | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 5c6933f86..a09dfbfff 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -817,6 +817,7 @@ def _device_init_abstract_eval(shots, auto_qubit_management, rtd_lib, rtd_name, return () +# pylint: disable=too-many-arguments, too-many-positional-arguments def _device_init_lowering( jax_ctx: mlir.LoweringRuleContext, shots: ir.Value, diff --git a/frontend/test/pytest/test_automatic_qubit_management.py b/frontend/test/pytest/test_automatic_qubit_management.py index 4324f982f..eade14ccc 100644 --- a/frontend/test/pytest/test_automatic_qubit_management.py +++ b/frontend/test/pytest/test_automatic_qubit_management.py @@ -35,6 +35,10 @@ def test_partial_sample(backend): + """ + Test that a `sample` terminal measurement with wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.RX(0.0, wires=0) @@ -48,6 +52,10 @@ def circuit(): def test_partial_counts(backend): + """ + Test that a `counts` terminal measurement with wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.RX(0.0, wires=0) @@ -61,6 +69,10 @@ def circuit(): def test_partial_probs(backend): + """ + Test that a `probs` terminal measurement with wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.PauliX(wires=0) @@ -74,6 +86,10 @@ def circuit(): def test_sample(backend): + """ + Test that a `sample` terminal measurement without wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.RX(0.0, wires=3) @@ -87,6 +103,10 @@ def circuit(): def test_counts(backend): + """ + Test that a `counts` terminal measurement without wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.RX(0.0, wires=3) @@ -100,6 +120,10 @@ def circuit(): def test_probs(backend): + """ + Test that a `probs` terminal measurement without wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.PauliX(wires=3) @@ -113,6 +137,10 @@ def circuit(): def test_state(backend): + """ + Test that a `state` terminal measurement without wires specified can be executed + correctly with automatic qubit management. + """ def circuit(): qml.PauliX(wires=3) From 6e14e85081d74c9cc87ad952f23e4666138805e0 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 09:56:32 -0400 Subject: [PATCH 14/19] add py lit etst --- frontend/test/lit/test_measurements.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/frontend/test/lit/test_measurements.py b/frontend/test/lit/test_measurements.py index 502f660f4..827f7cab8 100644 --- a/frontend/test/lit/test_measurements.py +++ b/frontend/test/lit/test_measurements.py @@ -678,3 +678,26 @@ def circ(): state_dynamic(10) print(state_dynamic.mlir) + + +# CHECK-LABEL: @automatic_qubit_management +@qjit(target="mlir") +def automatic_qubit_management(): + @qml.qnode(qml.device("lightning.qubit")) + def circ(): + # CHECK: [[qreg:%.+]] = quantum.alloc( 0) : !quantum.reg + # CHECK: [[in_qubit:%.+]] = quantum.extract [[qreg]][ 2] : !quantum.reg -> !quantum.bit + # CHECK: [[out_qubit:%.+]] = quantum.custom "Hadamard"() [[in_qubit]] : !quantum.bit + qml.Hadamard(wires=2) + + # CHECK: [[nqubits:%.+]] = quantum.num_qubits : i64 + # CHECK: [[toTensor:%.+]] = tensor.from_elements [[nqubits]] : tensor + # CHECK: [[probs_shape:%.+]] = stablehlo.shift_left {{%.+}}, [[toTensor]] : tensor + # CHECK: [[deTensor:%.+]] = tensor.extract [[probs_shape]][] : tensor + # CHECK: {{%.+}} = quantum.probs {{%.+}} shape [[deTensor]] : tensor + return qml.probs() + + return circ() + + +print(automatic_qubit_management.mlir) From c159c440672b2bcfb5930c10d1b9c47872d17f1f Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 10:57:55 -0400 Subject: [PATCH 15/19] move mode to RTDevice constructor --- runtime/lib/capi/ExecutionContext.hpp | 25 +++++++++++++------------ runtime/lib/capi/RuntimeCAPI.cpp | 14 ++++---------- 2 files changed, 17 insertions(+), 22 deletions(-) diff --git a/runtime/lib/capi/ExecutionContext.hpp b/runtime/lib/capi/ExecutionContext.hpp index 33aeaeb8b..5b18c4db8 100644 --- a/runtime/lib/capi/ExecutionContext.hpp +++ b/runtime/lib/capi/ExecutionContext.hpp @@ -209,16 +209,17 @@ class RTDevice { public: explicit RTDevice(std::string _rtd_lib, std::string _rtd_name = {}, - std::string _rtd_kwargs = {}) + std::string _rtd_kwargs = {}, bool _auto_qubit_management = false) : rtd_lib(std::move(_rtd_lib)), rtd_name(std::move(_rtd_name)), - rtd_kwargs(std::move(_rtd_kwargs)) + rtd_kwargs(std::move(_rtd_kwargs)), auto_qubit_management(_auto_qubit_management) { _pl2runtime_device_info(rtd_lib, rtd_name); } explicit RTDevice(std::string_view _rtd_lib, std::string_view _rtd_name, - std::string_view _rtd_kwargs) - : rtd_lib(_rtd_lib), rtd_name(_rtd_name), rtd_kwargs(_rtd_kwargs) + std::string_view _rtd_kwargs, bool _auto_qubit_management) + : rtd_lib(_rtd_lib), rtd_name(_rtd_name), rtd_kwargs(_rtd_kwargs), + auto_qubit_management(_auto_qubit_management) { _pl2runtime_device_info(rtd_lib, rtd_name); } @@ -260,8 +261,7 @@ class RTDevice { void setDeviceStatus(RTDeviceStatus new_status) noexcept { status = new_status; } - void setDeviceAutoQubitManagementMode(bool mode) { auto_qubit_management = mode; } - bool getDeviceAutoQubitManagementMode() { return auto_qubit_management; } + bool getQubitManagementMode() { return auto_qubit_management; } [[nodiscard]] auto getDeviceStatus() const -> RTDeviceStatus { return status; } @@ -317,12 +317,13 @@ class ExecutionContext final { } [[nodiscard]] auto getOrCreateDevice(std::string_view rtd_lib, std::string_view rtd_name, - std::string_view rtd_kwargs) + std::string_view rtd_kwargs, bool auto_qubit_management) -> const std::shared_ptr & { std::lock_guard lock(pool_mu); - auto device = std::make_shared(rtd_lib, rtd_name, rtd_kwargs); + auto device = + std::make_shared(rtd_lib, rtd_name, rtd_kwargs, auto_qubit_management); const size_t key = device_pool.size(); for (size_t i = 0; i < key; i++) { @@ -348,13 +349,13 @@ class ExecutionContext final { return device_pool[key]; } - [[nodiscard]] auto getOrCreateDevice(const std::string &rtd_lib, - const std::string &rtd_name = {}, - const std::string &rtd_kwargs = {}) + [[nodiscard]] auto + getOrCreateDevice(const std::string &rtd_lib, const std::string &rtd_name = {}, + const std::string &rtd_kwargs = {}, bool auto_qubit_management = false) -> const std::shared_ptr & { return getOrCreateDevice(std::string_view{rtd_lib}, std::string_view{rtd_name}, - std::string_view{rtd_kwargs}); + std::string_view{rtd_kwargs}, auto_qubit_management); } [[nodiscard]] auto getDevice(size_t device_key) -> const std::shared_ptr & diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index 6fec55258..0f5933b0e 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -76,9 +76,9 @@ std::vector getModifiersControlledValues(const Modifiers *modifiers) * to the new initialized device pointer. */ [[nodiscard]] bool initRTDevicePtr(std::string_view rtd_lib, std::string_view rtd_name, - std::string_view rtd_kwargs) + std::string_view rtd_kwargs, bool auto_qubit_management) { - auto &&device = CTX->getOrCreateDevice(rtd_lib, rtd_name, rtd_kwargs); + auto &&device = CTX->getOrCreateDevice(rtd_lib, rtd_name, rtd_kwargs, auto_qubit_management); if (device) { RTD_PTR = device.get(); return RTD_PTR ? true : false; @@ -252,9 +252,8 @@ static int __catalyst__rt__device_init__impl(int8_t *rtd_lib, int8_t *rtd_name, const std::vector args{ reinterpret_cast(rtd_lib), (rtd_name ? reinterpret_cast(rtd_name) : ""), (rtd_kwargs ? reinterpret_cast(rtd_kwargs) : "")}; - RT_FAIL_IF(!initRTDevicePtr(args[0], args[1], args[2]), + RT_FAIL_IF(!initRTDevicePtr(args[0], args[1], args[2], auto_qubit_management), "Failed initialization of the backend device"); - RTD_PTR->setDeviceAutoQubitManagementMode(auto_qubit_management); getQuantumDevicePtr()->SetDeviceShots(shots); if (CTX->getDeviceRecorderStatus()) { getQuantumDevicePtr()->StartTapeRecording(); @@ -1035,11 +1034,6 @@ int64_t __catalyst__rt__array_get_size_1d(QirArray *ptr) return qubit_vector_ptr->size(); } -static bool isDeviceAutomaticQubitManageMode() -{ - return RTD_PTR->getDeviceAutoQubitManagementMode(); -} - int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) { std::vector *qubit_vector_ptr = reinterpret_cast *>(ptr); @@ -1049,7 +1043,7 @@ int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) error_msg += std::to_string(idx); if (static_cast(idx) >= qubit_vector_ptr->size()) { - if (!isDeviceAutomaticQubitManageMode()) { + if (!RTD_PTR->getQubitManagementMode()) { RT_FAIL(error_msg.c_str()); } else { From ce8e78e609d5bd3ec83b89597a2007132784a77a Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 11:04:23 -0400 Subject: [PATCH 16/19] apply changelog suggestions --- doc/releases/changelog-dev.md | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 37018e36a..b3bd0c295 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -78,7 +78,7 @@ [(#1788)](https://github.com/PennyLaneAI/catalyst/pull/1788) The number of wires does not need to be speficied during device initialization, - and can be automatically allocated whenever a new wire is encountered. + and instead will be automatically managed by the Catalyst Runtime. ```python @qjit @@ -98,11 +98,13 @@ ``` In this example, the number of wires is not specified at device initialization. - When we encounter an X gate on `wires=2`, catalyst automatically allocates 3 wires, and applies - an X gate to the wire with index `2`. The result is the state |001>, giving the probabilities - above. + When we encounter an X gate on `wires=2`, catalyst automatically expands the size + of the qubit register to include the requested wire index. + Here, the register will contain (at least) 3 qubits after the X operation. + As a result, we can see the QNode returning the probabilities for the state |001>, + meaning 3 wires were allocated in total. - This feature can be turned on by simply not supplying a `wires` argument to the device. + This feature can be turned on by omitting the `wires` argument to the device.

Improvements 🛠

From 68291ae9b4744240a11e8097396b12017a6c8515 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 11:46:56 -0400 Subject: [PATCH 17/19] use batch alloc and outline --- runtime/lib/capi/RuntimeCAPI.cpp | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index 0f5933b0e..f6659a203 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -102,6 +102,20 @@ void deactivateDevice() CTX->deactivateDevice(RTD_PTR); RTD_PTR = nullptr; } + +static void autoQubitManagementAllocate(std::vector *qubit_vector_ptr, int64_t idx) +{ + // allocate new qubits if we are in automatic qubit allocation mode + // and encountered a new user wire index + // `idx` is the new user wire index from frontend pennylane + // number of currently allocated qubits is `qubit_vector_ptr->size()` + QirArray *new_qubits = __catalyst__rt__qubit_allocate_array( + idx + 1 - static_cast(qubit_vector_ptr->size())); + std::vector *new_qubits_vector = + reinterpret_cast *>(new_qubits); + qubit_vector_ptr->insert(qubit_vector_ptr->end(), new_qubits_vector->begin(), + new_qubits_vector->end()); +} } // namespace Catalyst::Runtime extern "C" { @@ -1047,14 +1061,7 @@ int8_t *__catalyst__rt__array_get_element_ptr_1d(QirArray *ptr, int64_t idx) RT_FAIL(error_msg.c_str()); } else { - // allocate a new qubit if we are in automatic qubit allocation mode - // `idx` is the new user wire index from frontend pennylane - // number of currently allocated qubits is `qubit_vector_ptr->size()` - while (qubit_vector_ptr->size() <= static_cast(idx)) { - auto new_alloced_id = - reinterpret_cast(__catalyst__rt__qubit_allocate()); - qubit_vector_ptr->push_back(new_alloced_id); - } + autoQubitManagementAllocate(qubit_vector_ptr, idx); } } From 43a72a2c97a0757077054c0898ecea7da4a0fc58 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Mon, 9 Jun 2025 14:06:16 -0400 Subject: [PATCH 18/19] update other relavant methods in RTDevice to include the new auto_qubit_management field --- runtime/lib/capi/ExecutionContext.hpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/runtime/lib/capi/ExecutionContext.hpp b/runtime/lib/capi/ExecutionContext.hpp index 5b18c4db8..56b754398 100644 --- a/runtime/lib/capi/ExecutionContext.hpp +++ b/runtime/lib/capi/ExecutionContext.hpp @@ -233,7 +233,8 @@ class RTDevice { auto operator==(const RTDevice &other) const -> bool { return (this->rtd_lib == other.rtd_lib && this->rtd_name == other.rtd_name) && - this->rtd_kwargs == other.rtd_kwargs; + this->rtd_kwargs == other.rtd_kwargs && + this->auto_qubit_management == other.auto_qubit_management; } [[nodiscard]] auto getQuantumDevicePtr() -> const std::unique_ptr & @@ -252,9 +253,10 @@ class RTDevice { return rtd_qdevice; } - [[nodiscard]] auto getDeviceInfo() const -> std::tuple + [[nodiscard]] auto getDeviceInfo() const + -> std::tuple { - return {rtd_lib, rtd_name, rtd_kwargs}; + return {rtd_lib, rtd_name, rtd_kwargs, auto_qubit_management}; } [[nodiscard]] auto getDeviceName() const -> const std::string & { return rtd_name; } @@ -268,7 +270,8 @@ class RTDevice { friend std::ostream &operator<<(std::ostream &os, const RTDevice &device) { os << "RTD, name: " << device.rtd_name << " lib: " << device.rtd_lib - << " kwargs: " << device.rtd_kwargs; + << " kwargs: " << device.rtd_kwargs + << "auto_qubit_management: " << device.auto_qubit_management; return os; } }; From 1bb541f2a96190f7974c8f086955ada1a527abb5 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Tue, 10 Jun 2025 10:27:20 -0400 Subject: [PATCH 19/19] add runtime test --- runtime/tests/Test_NullQubit.cpp | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/runtime/tests/Test_NullQubit.cpp b/runtime/tests/Test_NullQubit.cpp index e1b3838af..5eb22b438 100644 --- a/runtime/tests/Test_NullQubit.cpp +++ b/runtime/tests/Test_NullQubit.cpp @@ -50,6 +50,37 @@ TEST_CASE("Test __catalyst__rt__device_init registering device=null.qubit", "[Nu __catalyst__rt__finalize(); } +TEST_CASE("Test automatic qubit management", "[NullQubit]") +{ + constexpr size_t shots = 10; + const auto [rtd_lib, rtd_name, rtd_kwargs] = + std::array{"null.qubit", "null_qubit", ""}; + __catalyst__rt__initialize(nullptr); + __catalyst__rt__device_init((int8_t *)rtd_lib.c_str(), (int8_t *)rtd_name.c_str(), + (int8_t *)rtd_kwargs.c_str(), shots, + /*auto_qubit_management=*/true); + + QirArray *qs = __catalyst__rt__qubit_allocate_array(0); + + // a new index `2` will mean 3 new allocations + QUBIT **target = (QUBIT **)__catalyst__rt__array_get_element_ptr_1d(qs, 2); + + __catalyst__qis__Hadamard(*target, NO_MODIFIERS); + + size_t n = __catalyst__rt__num_qubits(); + CHECK(n == 3); + + double buffer[shots * n]; + MemRefT_double_2d result = {buffer, buffer, 0, {shots, n}, {n, 1}}; + __catalyst__qis__Sample(&result, n); + + __catalyst__rt__qubit_release_array(qs); + CHECK(__catalyst__rt__num_qubits() == 0); + + __catalyst__rt__device_release(); + __catalyst__rt__finalize(); +} + TEST_CASE("Test NullQubit qubit allocation is successful.", "[NullQubit]") { std::unique_ptr sim = std::make_unique();