Skip to content

Commit f79fc3f

Browse files
authored
Add experimental support for cuStateVecEx. (#965)
1 parent d0df9fc commit f79fc3f

39 files changed

+3167
-47
lines changed

CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ if(NOT CMAKE_APPLE_SILICON_PROCESSOR)
6464
add_subdirectory(pybind_interface/cuda)
6565
if(DEFINED ENV{CUQUANTUM_ROOT})
6666
add_subdirectory(pybind_interface/custatevec)
67+
add_subdirectory(pybind_interface/custatevecex)
6768
endif()
6869
elseif(has_hipcc)
6970
add_subdirectory(pybind_interface/hip)

Makefile

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,10 @@ ifneq (,$(strip $(CUQUANTUM_ROOT)))
9494
CUSVFLAGS += -lcustatevec -lcublas
9595
CUSTATEVECFLAGS ?= $(CUSVFLAGS)
9696
TARGETS += qsim-custatevec
97+
TARGETS += qsim-custatevecex
9798
TESTS += run-custatevec-tests
99+
TESTS += run-custatevecex-tests
100+
TESTS += run-custatevecex-mpi-tests
98101
else
99102
$(warning $$CUQUANTUM_ROOT is set, but the path does not seem to exist)
100103
endif
@@ -120,6 +123,10 @@ qsim-cuda:
120123
qsim-custatevec: | check-cuquantum-root-set
121124
$(MAKE) -C apps/ qsim-custatevec
122125

126+
.PHONY: qsim-custatevecex
127+
qsim-custatevecex: | check-cuquantum-root-set
128+
$(MAKE) -C apps/ qsim-custatevecex
129+
123130
.PHONY: qsim-hip
124131
qsim-hip:
125132
$(MAKE) -C apps/ qsim-hip
@@ -140,6 +147,10 @@ cuda-tests:
140147
custatevec-tests: | check-cuquantum-root-set
141148
$(MAKE) -C tests/ custatevec-tests
142149

150+
.PHONY: custatevecex-tests
151+
custatevecex-tests: | check-cuquantum-root-set
152+
$(MAKE) -C tests/ custatevecex-tests
153+
143154
.PHONY: hip-tests
144155
hip-tests:
145156
$(MAKE) -C tests/ hip-tests
@@ -156,6 +167,14 @@ run-cuda-tests: cuda-tests
156167
run-custatevec-tests: custatevec-tests
157168
$(MAKE) -C tests/ run-custatevec-tests
158169

170+
.PHONY: run-custatevecex-tests
171+
run-custatevecex-tests: custatevecex-tests
172+
$(MAKE) -C tests/ run-custatevecex-tests
173+
174+
.PHONY: run-custatevecex-mpi-tests
175+
run-custatevecex-mpi-tests: custatevecex-tests
176+
$(MAKE) -C tests/ run-custatevecex-mpi-tests
177+
159178
.PHONY: run-hip-tests
160179
run-hip-tests: hip-tests
161180
$(MAKE) -C tests/ run-hip-tests

apps/Makefile

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ CUDA_TARGETS := $(CUDA_TARGETS:%cuda.cu=%cuda.x)
77
CUSTATEVEC_TARGETS = $(shell find . -maxdepth 1 -name "*custatevec.cu")
88
CUSTATEVEC_TARGETS := $(CUSTATEVEC_TARGETS:%custatevec.cu=%custatevec.x)
99

10+
CUSTATEVECEX_TARGETS = $(shell find . -maxdepth 1 -name "*custatevecex.cu")
11+
CUSTATEVECEX_TARGETS := $(CUSTATEVECEX_TARGETS:%custatevecex.cu=%custatevecex.x)
12+
1013
HIP_TARGETS = $(shell find . -maxdepth 1 -name '*cuda.cu')
1114
HIP_TARGETS := $(HIP_TARGETS:%cuda.cu=%hip.x)
1215

@@ -19,6 +22,9 @@ qsim-cuda: $(CUDA_TARGETS)
1922
.PHONY: qsim-custatevec
2023
qsim-custatevec: $(CUSTATEVEC_TARGETS)
2124

25+
.PHONY: qsim-custatevecex
26+
qsim-custatevecex: $(CUSTATEVECEX_TARGETS)
27+
2228
.PHONY: qsim-hip
2329
qsim-hip: $(HIP_TARGETS)
2430

@@ -31,6 +37,9 @@ qsim-hip: $(HIP_TARGETS)
3137
%custatevec.x: %custatevec.cu
3238
$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
3339

40+
%custatevecex.x: %custatevecex.cu
41+
$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
42+
3443
%hip.x: %cuda.cu
3544
$(HIPCC) -o ./$@ $< $(HIPCCFLAGS)
3645

apps/make.sh

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@ if command -v nvcc &>/dev/null; then
3737
)
3838
nvcc -O3 "${CUSTATEVECFLAGS[@]}" \
3939
-o qsim_base_custatevec.x qsim_base_custatevec.cu
40-
40+
nvcc -O3 "${CUSTATEVECFLAGS[@]}" \
41+
-o qsim_base_custatevecex.x qsim_base_custatevecex.cu
4142
fi
4243
elif command -v hipcc &>/dev/null; then
4344
hipcc -O3 -o qsim_base_hip.x qsim_base_cuda.cu

apps/qsim_base_custatevecex.cu

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
// Copyright 2025 Google LLC. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include <unistd.h>
16+
17+
#include <algorithm>
18+
#include <complex>
19+
#include <limits>
20+
#include <string>
21+
22+
#include "../lib/circuit_qsim_parser.h"
23+
#include "../lib/formux.h"
24+
#include "../lib/gates_qsim.h"
25+
#include "../lib/io_file.h"
26+
#include "../lib/multiprocess_custatevecex.h"
27+
#include "../lib/run_custatevecex.h"
28+
#include "../lib/simulator_custatevecex.h"
29+
#include "../lib/util_custatevec.h"
30+
31+
struct Options {
32+
std::string circuit_file;
33+
unsigned maxtime = std::numeric_limits<unsigned>::max();
34+
unsigned seed = 1;
35+
unsigned verbosity = 0;
36+
};
37+
38+
Options GetOptions(int argc, char* argv[]) {
39+
constexpr char usage[] = "usage:\n ./qsim_base -c circuit -d maxtime "
40+
"-s seed -v verbosity\n";
41+
42+
Options opt;
43+
44+
int k;
45+
46+
while ((k = getopt(argc, argv, "c:d:s:v:")) != -1) {
47+
switch (k) {
48+
case 'c':
49+
opt.circuit_file = optarg;
50+
break;
51+
case 'd':
52+
opt.maxtime = std::atoi(optarg);
53+
break;
54+
case 's':
55+
opt.seed = std::atoi(optarg);
56+
break;
57+
case 'v':
58+
opt.verbosity = std::atoi(optarg);
59+
break;
60+
default:
61+
qsim::IO::errorf(usage);
62+
exit(1);
63+
}
64+
}
65+
66+
return opt;
67+
}
68+
69+
bool ValidateOptions(const Options& opt) {
70+
if (opt.circuit_file.empty()) {
71+
qsim::IO::errorf("circuit file is not provided.\n");
72+
return false;
73+
}
74+
75+
return true;
76+
}
77+
78+
template <typename StateSpace, typename State>
79+
void PrintAmplitudes(
80+
unsigned num_qubits, const StateSpace& state_space, const State& state) {
81+
static constexpr char const* bits[8] = {
82+
"000", "001", "010", "011", "100", "101", "110", "111",
83+
};
84+
85+
uint64_t size = std::min(uint64_t{8}, uint64_t{1} << num_qubits);
86+
unsigned s = 3 - std::min(unsigned{3}, num_qubits);
87+
88+
for (uint64_t i = 0; i < size; ++i) {
89+
auto a = state_space.GetAmpl(state, i);
90+
qsim::IO::messagef("%s:%16.8g%16.8g%16.8g\n",
91+
bits[i] + s, std::real(a), std::imag(a), std::norm(a));
92+
}
93+
}
94+
95+
int main(int argc, char* argv[]) {
96+
using namespace qsim;
97+
98+
auto opt = GetOptions(argc, argv);
99+
if (!ValidateOptions(opt)) {
100+
return 1;
101+
}
102+
103+
using fp_type = float;
104+
105+
Circuit<GateQSim<fp_type>> circuit;
106+
if (!CircuitQsimParser<IOFile>::FromFile(opt.maxtime, opt.circuit_file,
107+
circuit)) {
108+
return 1;
109+
}
110+
111+
struct Factory {
112+
using Simulator = qsim::SimulatorCuStateVecEx<fp_type>;
113+
using StateSpace = Simulator::StateSpace;
114+
115+
explicit Factory(unsigned verbosity = 0) : verbosity(verbosity) {
116+
mp.initialize();
117+
}
118+
119+
StateSpace CreateStateSpace() const {
120+
StateSpace::Parameter param;
121+
param.verbosity = verbosity;
122+
123+
return StateSpace{mp, param};
124+
}
125+
126+
Simulator CreateSimulator() const {
127+
return Simulator{};
128+
}
129+
130+
MultiProcessCuStateVecEx mp;
131+
unsigned verbosity;
132+
};
133+
134+
using Simulator = Factory::Simulator;
135+
using StateSpace = Simulator::StateSpace;
136+
using State = StateSpace::State;
137+
using Runner = CuStateVecExRunner<IO, Factory>;
138+
139+
Factory factory(opt.verbosity);
140+
141+
StateSpace state_space = factory.CreateStateSpace();
142+
State state = state_space.Create(circuit.num_qubits);
143+
144+
if (state_space.IsNull(state)) {
145+
IO::errorf("not enough memory: is the number of qubits too large?\n");
146+
return 1;
147+
}
148+
149+
state_space.SetStateZero(state);
150+
151+
Runner::Parameter param;
152+
param.seed = opt.seed;
153+
param.verbosity = opt.verbosity;
154+
155+
if (Runner::Run(param, factory, circuit, state)) {
156+
PrintAmplitudes(circuit.num_qubits, state_space, state);
157+
}
158+
159+
return 0;
160+
}

docs/cirq_interface.md

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,8 +186,11 @@ library.
186186
`QSimOptions` provides five parameters to configure GPU execution. `use_gpu`
187187
is required to enable GPU execution:
188188
* `use_gpu`: if True, use GPU instead of CPU for simulation.
189-
* `gpu_mode`: use CUDA if set to 0 (default value) or use the NVIDIA cuStateVec
190-
library if set to any other value.
189+
* `gpu_mode`: use CUDA if set to 0 (default value), use the NVIDIA cuStateVec
190+
if set to 1 or use the NVIDIA cuStateVecEx library if set to any other value.
191+
192+
In the case of the NVIDIA cuStateVecEx library, simulations can be performed
193+
in multi-device / multi-node environments.
191194

192195
If `use_gpu` is set and `gpu_mode` is set to 0, the remaining parameters can
193196
optionally be set to fine-tune StateSpace performance for a specific device.
@@ -196,3 +199,4 @@ In most cases, the default values provide good performance.
196199
StateSpace. This must be a power of 2 in the range [32, 1024].
197200
* `gpu_data_blocks`: number of data blocks to use for the GPU StateSpace.
198201
Below 16 data blocks, performance is noticeably reduced.
202+

0 commit comments

Comments
 (0)