From 3092fdab78b7c085978ddfc1c0f65b6b4e81e397 Mon Sep 17 00:00:00 2001 From: Jonathan Wang Date: Mon, 3 Aug 2026 20:14:13 +0000 Subject: [PATCH 1/2] Add zero-copy GPU state vector access (#836) This PR adds an opt-in zero-copy API, `QSimSimulator.simulate_into_device_array(...)`, to extract the final state vector from GPU simulation without copying data from GPU to host memory (issue #836). The returned `DeviceStateVector` object owns the GPU allocation and exposes it through the CUDA Array Interface (`__cuda_array_interface__`, v3), allowing downstream GPU frameworks (CuPy, PyTorch, Numba) to consume the device buffer directly without a device -> host -> device round trip. --- lib/statespace_custatevecex.h | 4 + lib/vectorspace_cuda.h | 6 + lib/vectorspace_custatevecex.h | 11 ++ pybind_interface/cuda/pybind_main_cuda.cpp | 3 + .../custatevec/pybind_main_custatevec.cpp | 3 + .../custatevecex/pybind_main_custatevecex.cpp | 3 + pybind_interface/hip/pybind_main_hip.cpp | 3 + pybind_interface/pybind_main.cpp | 181 ++++++++++++++++++ pybind_interface/pybind_main.h | 9 + qsimcirq/qsim_simulator.py | 115 +++++++++++ qsimcirq_tests/qsimcirq_test.py | 157 +++++++++++++++ 11 files changed, 495 insertions(+) diff --git a/lib/statespace_custatevecex.h b/lib/statespace_custatevecex.h index 8aa6167f2..8590f16d8 100644 --- a/lib/statespace_custatevecex.h +++ b/lib/statespace_custatevecex.h @@ -78,6 +78,10 @@ class StateSpaceCuStateVecEx : void InternalToNormalOrder(State& state) const { state.to_normal_order(); + + // to_normal_order() is asynchronous; synchronize so that callers may + // observe the raw device buffer as soon as this method returns. + ErrorCheck(custatevecExStateVectorSynchronize(state.get())); } void NormalToInternalOrder(State& state) const { diff --git a/lib/vectorspace_cuda.h b/lib/vectorspace_cuda.h index f1efdd51d..16a012990 100644 --- a/lib/vectorspace_cuda.h +++ b/lib/vectorspace_cuda.h @@ -68,6 +68,12 @@ class VectorSpaceCUDA { return ptr_.release(); } + // Raw pointer to device (GPU) memory. Ownership is retained by this + // vector; the pointer is invalidated when the vector is destroyed. + void* device_ptr() const { + return static_cast(ptr_.get()); + } + unsigned num_qubits() const { return num_qubits_; } diff --git a/lib/vectorspace_custatevecex.h b/lib/vectorspace_custatevecex.h index 69f23affc..5ea05ba10 100644 --- a/lib/vectorspace_custatevecex.h +++ b/lib/vectorspace_custatevecex.h @@ -160,6 +160,17 @@ class VectorSpaceCuStateVecEx { return true; } + // Raw pointer to device (GPU) memory. Only meaningful in single-device + // mode; the state has no single contiguous device buffer when it is + // distributed across multiple devices or processes. Ownership is + // retained by this vector. + void* device_ptr() const { + if (distr_type_ != kSingleDevice) { + return nullptr; + } + return get_resources(0).device_ptr; + } + const auto& get_wire_ordering() const { ErrorCheck(custatevecExStateVectorGetProperty( ptr_, CUSTATEVEC_EX_SV_PROP_WIRE_ORDERING, diff --git a/pybind_interface/cuda/pybind_main_cuda.cpp b/pybind_interface/cuda/pybind_main_cuda.cpp index 8d0e82dda..9d8d061e2 100644 --- a/pybind_interface/cuda/pybind_main_cuda.cpp +++ b/pybind_interface/cuda/pybind_main_cuda.cpp @@ -54,4 +54,7 @@ namespace qsim { inline void ClearFlushToZeroAndDenormalsAreZeros() {} } +// Enable zero-copy device state-vector bindings (issue #836). +#define QSIM_DEVICE_STATE_BINDINGS + #include "../pybind_main.cpp" diff --git a/pybind_interface/custatevec/pybind_main_custatevec.cpp b/pybind_interface/custatevec/pybind_main_custatevec.cpp index 5e9a53e7b..8c0813943 100644 --- a/pybind_interface/custatevec/pybind_main_custatevec.cpp +++ b/pybind_interface/custatevec/pybind_main_custatevec.cpp @@ -63,4 +63,7 @@ namespace qsim { inline void ClearFlushToZeroAndDenormalsAreZeros() {} } +// Enable zero-copy device state-vector bindings (issue #836). +#define QSIM_DEVICE_STATE_BINDINGS + #include "../pybind_main.cpp" diff --git a/pybind_interface/custatevecex/pybind_main_custatevecex.cpp b/pybind_interface/custatevecex/pybind_main_custatevecex.cpp index 688b4391e..594246e11 100644 --- a/pybind_interface/custatevecex/pybind_main_custatevecex.cpp +++ b/pybind_interface/custatevecex/pybind_main_custatevecex.cpp @@ -107,4 +107,7 @@ namespace qsim { inline void ClearFlushToZeroAndDenormalsAreZeros() {} } +// Enable zero-copy device state-vector bindings (issue #836). +#define QSIM_DEVICE_STATE_BINDINGS + #include "../pybind_main.cpp" diff --git a/pybind_interface/hip/pybind_main_hip.cpp b/pybind_interface/hip/pybind_main_hip.cpp index fc2e6789b..80e0d8a8a 100644 --- a/pybind_interface/hip/pybind_main_hip.cpp +++ b/pybind_interface/hip/pybind_main_hip.cpp @@ -54,4 +54,7 @@ namespace qsim { inline void ClearFlushToZeroAndDenormalsAreZeros() {} } +// Enable zero-copy device state-vector bindings (issue #836). +#define QSIM_DEVICE_STATE_BINDINGS + #include "../pybind_main.cpp" diff --git a/pybind_interface/pybind_main.cpp b/pybind_interface/pybind_main.cpp index 3ca152a60..be7b50ce8 100644 --- a/pybind_interface/pybind_main.cpp +++ b/pybind_interface/pybind_main.cpp @@ -17,7 +17,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -590,6 +592,50 @@ class SimulatorHelper { return results; } +#ifdef QSIM_DEVICE_STATE_BINDINGS + // Runs a fullstate simulation and returns a heap-allocated helper whose + // final state stays alive in device (GPU) memory. Returns nullptr if the + // simulation fails. + template + static std::unique_ptr simulate_fullstate_device( + const py::dict &options, bool is_noisy, const StateType& input_state) { + std::unique_ptr helper( + new SimulatorHelper(options, is_noisy)); + // IsNull distinguishes allocation failure here from the multi-device + // case, which device_ptr_normal_order() reports separately. + if (!helper->is_valid || StateSpace::IsNull(helper->state) || + !helper->simulate(input_state)) { + return nullptr; + } + return helper; + } + + // Returns the device pointer to the final state vector, converting the + // buffer to normal (interleaved complex) order in place on first call. + // The state must not be used for further simulation afterwards. Returns + // nullptr, without reordering the state, if the state has no single + // contiguous device buffer (e.g. multi-device or multi-process + // cuStateVecEx states). + void* device_ptr_normal_order() { + void* ptr = state.device_ptr(); + if (ptr == nullptr) { + return nullptr; + } + if (!normal_order_done_) { + StateSpace state_space = factory.CreateStateSpace(); + state_space.InternalToNormalOrder(state); + StateSpace::DeviceSync(); + normal_order_done_ = true; + } + // Re-query in case the reorder relocated the buffer. + return state.device_ptr(); + } + + unsigned get_num_qubits() const { + return num_qubits; + } +#endif + private: SimulatorHelper(const py::dict &options, bool noisy) : factory(Factory(options)), @@ -773,8 +819,143 @@ class SimulatorHelper { // Only set to "true" once initialization is complete. bool is_valid; + +#ifdef QSIM_DEVICE_STATE_BINDINGS + // Once set, the internal state layout has been irreversibly mutated to + // normal order (for the native CUDA backend); the state must not be used + // for further simulation or sampling. + bool normal_order_done_ = false; +#endif +}; + +#ifdef QSIM_DEVICE_STATE_BINDINGS + +// Owns the final state of a fullstate simulation in device (GPU) memory and +// exposes it through the CUDA Array Interface (version 3), so that libraries +// such as CuPy, Numba and PyTorch can consume the state without copying it +// to the host. See https://github.com/quantumlib/qsim/issues/836. +class DeviceStateVector { + public: + explicit DeviceStateVector(std::unique_ptr helper) + : helper_(std::move(helper)) {} + + template + static std::unique_ptr simulate( + const py::dict &options, bool is_noisy, const StateType& input_state) { + auto helper = SimulatorHelper::simulate_fullstate_device( + options, is_noisy, input_state); + if (helper == nullptr) { + throw std::runtime_error("qsim simulation errored out."); + } + return std::unique_ptr( + new DeviceStateVector(std::move(helper))); + } + + unsigned num_qubits() const { + check_not_freed(); + return helper_->get_num_qubits(); + } + + bool is_freed() const { + return helper_ == nullptr; + } + + // Releases the device memory immediately instead of waiting for garbage + // collection. Idempotent. + void free() { + helper_.reset(); + } + + py::dict cuda_array_interface() { + check_not_freed(); + void* ptr = helper_->device_ptr_normal_order(); + if (ptr == nullptr) { + throw std::runtime_error( + "This state has no single contiguous device buffer (multi-device " + "or multi-process simulation is not supported)."); + } + + py::dict interface; + // The "stream" key is intentionally omitted per CAI v3: the data is + // fully synchronized (DeviceSync) before the pointer is exposed. + interface["shape"] = + py::make_tuple(uint64_t{1} << helper_->get_num_qubits()); + interface["typestr"] = "(ptr), false); + interface["version"] = 3; + return interface; + } + + private: + void check_not_freed() const { + if (helper_ == nullptr) { + throw std::runtime_error("This DeviceStateVector has been freed."); + } + } + + std::unique_ptr helper_; }; +void bind_device_state_vector(py::module_& m) { + // module_local: several GPU modules (e.g. qsim_cuda and qsim_custatevec) + // are loaded into the same process and each registers this class; without + // it, pybind11's shared type registry rejects the second registration and + // `import qsimcirq` fails. Instances never cross modules. + py::class_( + m, "DeviceStateVector", py::module_local(), + "Final state vector of a simulation, held in device (GPU) memory. " + "After free(), all properties and methods except is_freed and " + "free() itself raise RuntimeError.") + .def_property_readonly("__cuda_array_interface__", + &DeviceStateVector::cuda_array_interface) + .def_property_readonly("num_qubits", &DeviceStateVector::num_qubits, + "Number of qubits in the state vector. Raises " + "RuntimeError if the state has been freed.") + .def_property_readonly("is_freed", &DeviceStateVector::is_freed, + "Whether the device memory has been released. " + "Always safe to access, even after free().") + .def("free", &DeviceStateVector::free, + "Releases the device memory held by this object. Idempotent; " + "afterwards all other properties and methods raise RuntimeError. " + "WARNING: any view previously created from " + "__cuda_array_interface__ (e.g. via cupy.asarray) becomes a " + "dangling device pointer; drop all such views, and synchronize " + "any streams still reading the buffer, before calling free()."); + + m.def( + "qsim_simulate_fullstate_device", + [](const py::dict &options, uint64_t input_state) { + return DeviceStateVector::simulate(options, false, input_state); + }, + "Call the qsim simulator, keeping the final state in device memory"); + m.def( + "qsim_simulate_fullstate_device", + [](const py::dict &options, const py::array_t &input_vector) { + return DeviceStateVector::simulate(options, false, input_vector); + }, + "Call the qsim simulator, keeping the final state in device memory"); + m.def( + "qtrajectory_simulate_fullstate_device", + [](const py::dict &options, uint64_t input_state) { + return DeviceStateVector::simulate(options, true, input_state); + }, + "Call the qtrajectory simulator, keeping the final state in device " + "memory"); + m.def( + "qtrajectory_simulate_fullstate_device", + [](const py::dict &options, const py::array_t &input_vector) { + return DeviceStateVector::simulate(options, true, input_vector); + }, + "Call the qtrajectory simulator, keeping the final state in device " + "memory"); +} + +#endif // QSIM_DEVICE_STATE_BINDINGS + // Methods for simulating full state vectors. py::array_t qsim_simulate_fullstate( diff --git a/pybind_interface/pybind_main.h b/pybind_interface/pybind_main.h index 92a98c942..83699f35b 100644 --- a/pybind_interface/pybind_main.h +++ b/pybind_interface/pybind_main.h @@ -159,6 +159,12 @@ qtrajectory_simulate_moment_expectation_values( // Hybrid simulator. std::vector> qsimh_simulate(const py::dict &options); +#ifdef QSIM_DEVICE_STATE_BINDINGS +void bind_device_state_vector(py::module_& m); +#else +inline void bind_device_state_vector(py::module_& m) {} +#endif + template T ParseOptions(const py::dict& options, const char* key) { if (!options.contains(key)) { @@ -393,6 +399,9 @@ T ParseOptions(const py::dict& options, const char* key) { &qtrajectory_simulate_fullstate), \ "Call the qtrajectory simulator for full state vector simulation"); \ \ + /* Zero-copy access to the final state in device memory (issue #836) */ \ + bind_device_state_vector(m); \ + \ /* Methods for returning samples */ \ m.def("qsim_sample", &qsim_sample, "Call the qsim sampler"); \ m.def("qsim_sample_final", &qsim_sample_final, \ diff --git a/qsimcirq/qsim_simulator.py b/qsimcirq/qsim_simulator.py index 2b538167c..b2b20f34c 100644 --- a/qsimcirq/qsim_simulator.py +++ b/qsimcirq/qsim_simulator.py @@ -550,6 +550,121 @@ def simulate_into_1d_array( params = cirq.study.ParamResolver(param_resolver) return next(self._simulate_impl(program, params, qubit_order, initial_state)) + def simulate_into_device_array( + self, + program: cirq.AbstractCircuit, + param_resolver: cirq.ParamResolverOrSimilarType = None, + qubit_order: cirq.QubitOrderOrList = cirq.ops.QubitOrder.DEFAULT, + initial_state: Optional[Union[int, np.ndarray]] = None, + ) -> Tuple[cirq.ParamResolver, Any, Sequence[cirq.Qid]]: + """Simulates the circuit, leaving the final state in GPU memory. + + Requires a GPU backend (`use_gpu=True`). Unlike `simulate`, the final + state vector is not copied to host memory. Instead, this returns an + object exposing `__cuda_array_interface__`, which GPU libraries such + as CuPy, Numba or PyTorch can consume without a copy, e.g. + `cupy.asarray(device_state)`. + + The returned object owns the device buffer; the buffer is released + when the object is garbage-collected or when its `free()` method is + called. Consumers such as CuPy and Numba keep the returned object + alive via the `__cuda_array_interface__` owner reference. However, + `free()` releases the device buffer immediately and must not be + called while any consumer view of the buffer exists (doing so leads + to use-after-free on the device); synchronize any streams still + reading the buffer first. After `free()`, all properties and + methods of the returned object except `is_freed` and `free()` itself + raise RuntimeError. Release the state (via `free()` or by dropping + all references) before interpreter shutdown; freeing device memory + during interpreter teardown, after the CUDA context is destroyed, + is unsafe. + + The state has the same layout and qubit-ordering convention as the + numpy array returned by `simulate_into_1d_array`. + + Multi-device and multi-process simulations (`gpu_mode >= 2` running + on more than one GPU) have no single device buffer; accessing + `__cuda_array_interface__` on their results raises a RuntimeError. + + Returns: + Tuple of (param resolver, device state, qubit order). The device + state is a `DeviceStateVector` (a per-backend pybind11 class) + exposing `__cuda_array_interface__`, `num_qubits`, `is_freed` + and `free()`. The qubit order is the cirq-ordered qubits + (`Sequence[cirq.Qid]`) corresponding to the state's qubit + indexing. + + Raises: + ValueError: if this simulator was not configured with + `use_gpu=True`, or if `initial_state` is a vector whose size + does not match the number of qubits. + TypeError: if `initial_state` is neither an int nor a + `np.complex64` numpy array. + RuntimeError: if the simulation fails (e.g. the device state + could not be allocated). Additionally, accessing + `__cuda_array_interface__` on the returned object raises + RuntimeError if the state is spread across multiple devices + or has been freed. + """ + if not self.qsim_options["g"]: + raise ValueError( + "simulate_into_device_array requires GPU execution. " + "Set use_gpu=True in QSimOptions." + ) + + if initial_state is None: + initial_state = 0 + if not isinstance(initial_state, (int, np.ndarray)): + raise TypeError("initial_state must be an int or state vector.") + + # Add noise to the circuit if a noise model was provided. + all_qubits = program.all_qubits() + program = qsimc.QSimCircuit( + ( + self.noise.noisy_moments(program, sorted(all_qubits)) + if self.noise is not cirq.NO_NOISE + else program + ), + ) + + options = {} + options.update(self.qsim_options) + + prs = cirq.study.ParamResolver(param_resolver) + cirq_order = cirq.QubitOrder.as_qubit_order(qubit_order).order_for(all_qubits) + num_qubits = len(cirq_order) + if isinstance(initial_state, np.ndarray): + if initial_state.dtype != np.complex64: + raise TypeError("initial_state vector must have dtype np.complex64.") + input_vector = initial_state.view(np.float32) + if len(input_vector) != 2**num_qubits * 2: + raise ValueError( + "initial_state vector size must match number of qubits. " + f"Expected: {2**num_qubits * 2} Received: {len(input_vector)}" + ) + + if _needs_trajectories(program): + translator_fn_name = "translate_cirq_to_qtrajectory" + simulator_fn = self._sim_module.qtrajectory_simulate_fullstate_device + else: + translator_fn_name = "translate_cirq_to_qsim" + simulator_fn = self._sim_module.qsim_simulate_fullstate_device + + solved_circuit = cirq.resolve_parameters(program, prs) + options["c"], _ = self._translate_circuit( + solved_circuit, + translator_fn_name, + cirq_order, + ) + options["s"] = self.get_seed() + + if isinstance(initial_state, int): + device_state = simulator_fn(options, initial_state) + else: + device_state = simulator_fn(options, input_vector) + + return prs, device_state, cirq_order + def simulate_sweep_iter( self, program: cirq.Circuit, diff --git a/qsimcirq_tests/qsimcirq_test.py b/qsimcirq_tests/qsimcirq_test.py index df7c9d9dd..1ca017a78 100644 --- a/qsimcirq_tests/qsimcirq_test.py +++ b/qsimcirq_tests/qsimcirq_test.py @@ -1514,6 +1514,163 @@ def test_qsim_gpu_input_state(): assert cirq.approx_eq(state_vector[i], 0, atol=1e-6) +def _device_state_to_numpy(device_state): + """Copies a __cuda_array_interface__ buffer back to host via CuPy.""" + cupy = pytest.importorskip("cupy") + return cupy.asnumpy(cupy.asarray(device_state)) + + +def test_simulate_into_device_array_requires_gpu(): + cpu_sim = qsimcirq.QSimSimulator() + a, b = cirq.LineQubit.range(2) + circuit = cirq.Circuit(cirq.H(a), cirq.CNOT(a, b)) + with pytest.raises(ValueError, match="requires GPU execution"): + cpu_sim.simulate_into_device_array(circuit) + + +def test_cirq_qsim_gpu_simulate_into_device_array(): + if qsimcirq.qsim_gpu is None: + pytest.skip("GPU is not available for testing.") + pytest.importorskip("cupy") + + qubits = cirq.LineQubit.range(5) + circuit = cirq.testing.random_circuit( + qubits, n_moments=10, op_density=1.0, random_state=11 + ) + + gpu_options = qsimcirq.QSimOptions(use_gpu=True) + sim = qsimcirq.QSimSimulator(qsim_options=gpu_options) + + _, device_state, _ = sim.simulate_into_device_array(circuit) + + interface = device_state.__cuda_array_interface__ + assert interface["shape"] == (2 ** len(qubits),) + assert interface["typestr"] == " 1: + # With gpu_mode=2, cuStateVecEx spreads the state across all visible + # devices; there is no single contiguous device buffer to expose. + with pytest.raises(RuntimeError, match="no single contiguous device buffer"): + _ = device_state.__cuda_array_interface__ + else: + expected = sim.simulate(circuit).final_state_vector + actual = _device_state_to_numpy(device_state) + assert np.allclose(actual, expected, atol=1e-6) + + def test_cirq_qsim_custatevec_amplitudes(): if qsimcirq.qsim_custatevec is None: pytest.skip("cuStateVec library is not available for testing.") From 1e86a7d3bbc263902e13a8c57da77296d8b1fc4e Mon Sep 17 00:00:00 2001 From: Jonathan Wang Date: Thu, 6 Aug 2026 22:20:57 +0000 Subject: [PATCH 2/2] Address review comments for zero-copy device state (#836) - Ensure initial_state NumPy array is C-contiguous (via np.ascontiguousarray) before creating float32 view to prevent out-of-bounds reads on non-contiguous array slices. - Use std::make_unique instead of bare new expression per Google C++ Style Guide. - Add test_cirq_qsim_gpu_simulate_into_device_array_with_non_contiguous_input_state test case. --- pybind_interface/pybind_main.cpp | 3 +-- qsimcirq/qsim_simulator.py | 1 + qsimcirq_tests/qsimcirq_test.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/pybind_interface/pybind_main.cpp b/pybind_interface/pybind_main.cpp index be7b50ce8..8fb98e1af 100644 --- a/pybind_interface/pybind_main.cpp +++ b/pybind_interface/pybind_main.cpp @@ -847,8 +847,7 @@ class DeviceStateVector { if (helper == nullptr) { throw std::runtime_error("qsim simulation errored out."); } - return std::unique_ptr( - new DeviceStateVector(std::move(helper))); + return std::make_unique(std::move(helper)); } unsigned num_qubits() const { diff --git a/qsimcirq/qsim_simulator.py b/qsimcirq/qsim_simulator.py index b2b20f34c..be05686dd 100644 --- a/qsimcirq/qsim_simulator.py +++ b/qsimcirq/qsim_simulator.py @@ -636,6 +636,7 @@ def simulate_into_device_array( if isinstance(initial_state, np.ndarray): if initial_state.dtype != np.complex64: raise TypeError("initial_state vector must have dtype np.complex64.") + initial_state = np.ascontiguousarray(initial_state) input_vector = initial_state.view(np.float32) if len(input_vector) != 2**num_qubits * 2: raise ValueError( diff --git a/qsimcirq_tests/qsimcirq_test.py b/qsimcirq_tests/qsimcirq_test.py index 1ca017a78..a8c4a9823 100644 --- a/qsimcirq_tests/qsimcirq_test.py +++ b/qsimcirq_tests/qsimcirq_test.py @@ -1581,6 +1581,37 @@ def test_cirq_qsim_gpu_simulate_into_device_array_with_input_state(): assert cirq.approx_eq(state_vector[i], 0, atol=1e-6) +def test_cirq_qsim_gpu_simulate_into_device_array_with_non_contiguous_input_state(): + if qsimcirq.qsim_gpu is None: + pytest.skip("GPU is not available for testing.") + pytest.importorskip("cupy") + + num_qubits = 2 + qubits = cirq.LineQubit.range(num_qubits) + circuit = cirq.Circuit(cirq.H.on_each(*qubits)) + + gpu_options = qsimcirq.QSimOptions(use_gpu=True) + sim = qsimcirq.QSimSimulator(qsim_options=gpu_options) + + # Create non-contiguous sliced initial_state array + full_array = np.zeros(8, dtype=np.complex64) + full_array[0] = 0.5 + full_array[2] = 0.5 + full_array[4] = 0.5 + full_array[6] = 0.5 + initial_state_strided = full_array[::2] + assert not initial_state_strided.flags.c_contiguous + + _, device_state, _ = sim.simulate_into_device_array( + circuit, initial_state=initial_state_strided + ) + state_vector = _device_state_to_numpy(device_state) + + assert cirq.approx_eq(state_vector[0], 1, atol=1e-6) + for i in range(1, 4): + assert cirq.approx_eq(state_vector[i], 0, atol=1e-6) + + def test_device_state_vector_free(): if qsimcirq.qsim_gpu is None: pytest.skip("GPU is not available for testing.")