Skip to content

Commit 25876de

Browse files
committed
Add support for cuStateVecEx.
1 parent 9f00b45 commit 25876de

37 files changed

Lines changed: 3141 additions & 39 deletions

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: 15 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
@@ -140,6 +143,10 @@ cuda-tests:
140143
custatevec-tests: | check-cuquantum-root-set
141144
$(MAKE) -C tests/ custatevec-tests
142145

146+
.PHONY: custatevecex-tests
147+
custatevecex-tests: | check-cuquantum-root-set
148+
$(MAKE) -C tests/ custatevecex-tests
149+
143150
.PHONY: hip-tests
144151
hip-tests:
145152
$(MAKE) -C tests/ hip-tests
@@ -156,6 +163,14 @@ run-cuda-tests: cuda-tests
156163
run-custatevec-tests: custatevec-tests
157164
$(MAKE) -C tests/ run-custatevec-tests
158165

166+
.PHONY: run-custatevecex-tests
167+
run-custatevecex-tests: custatevecex-tests
168+
$(MAKE) -C tests/ run-custatevecex-tests
169+
170+
.PHONY: run-custatevecex-mpi-tests
171+
run-custatevecex-mpi-tests: custatevecex-tests
172+
$(MAKE) -C tests/ run-custatevecex-mpi-tests
173+
159174
.PHONY: run-hip-tests
160175
run-hip-tests: hip-tests
161176
$(MAKE) -C tests/ run-hip-tests

apps/Makefile

Lines changed: 6 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+
CUSTATEVEC_TARGETS = $(shell find . -maxdepth 1 -name "*custatevecex.cu")
11+
CUSTATEVEC_TARGETS := $(CUSTATEVEC_TARGETS:%custatevec.cu=%custatevecex.x)
12+
1013
HIP_TARGETS = $(shell find . -maxdepth 1 -name '*cuda.cu')
1114
HIP_TARGETS := $(HIP_TARGETS:%cuda.cu=%hip.x)
1215

@@ -31,6 +34,9 @@ qsim-hip: $(HIP_TARGETS)
3134
%custatevec.x: %custatevec.cu
3235
$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
3336

37+
%custatevecex.x: %custatevecex.cu
38+
$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
39+
3440
%hip.x: %cuda.cu
3541
$(HIPCC) -o ./$@ $< $(HIPCCFLAGS)
3642

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: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
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+
break;
58+
case 'v':
59+
opt.verbosity = std::atoi(optarg);
60+
break;
61+
default:
62+
qsim::IO::errorf(usage);
63+
exit(1);
64+
}
65+
}
66+
67+
return opt;
68+
}
69+
70+
bool ValidateOptions(const Options& opt) {
71+
if (opt.circuit_file.empty()) {
72+
qsim::IO::errorf("circuit file is not provided.\n");
73+
return false;
74+
}
75+
76+
return true;
77+
}
78+
79+
template <typename StateSpace, typename State>
80+
void PrintAmplitudes(
81+
unsigned num_qubits, const StateSpace& state_space, const State& state) {
82+
static constexpr char const* bits[8] = {
83+
"000", "001", "010", "011", "100", "101", "110", "111",
84+
};
85+
86+
uint64_t size = std::min(uint64_t{8}, uint64_t{1} << num_qubits);
87+
unsigned s = 3 - std::min(unsigned{3}, num_qubits);
88+
89+
for (uint64_t i = 0; i < size; ++i) {
90+
auto a = state_space.GetAmpl(state, i);
91+
qsim::IO::messagef("%s:%16.8g%16.8g%16.8g\n",
92+
bits[i] + s, std::real(a), std::imag(a), std::norm(a));
93+
}
94+
}
95+
96+
int main(int argc, char* argv[]) {
97+
using namespace qsim;
98+
99+
auto opt = GetOptions(argc, argv);
100+
if (!ValidateOptions(opt)) {
101+
return 1;
102+
}
103+
104+
using fp_type = float;
105+
106+
Circuit<GateQSim<fp_type>> circuit;
107+
if (!CircuitQsimParser<IOFile>::FromFile(opt.maxtime, opt.circuit_file,
108+
circuit)) {
109+
return 1;
110+
}
111+
112+
struct Factory {
113+
using Simulator = qsim::SimulatorCuStateVecEx<fp_type>;
114+
using StateSpace = Simulator::StateSpace;
115+
116+
explicit Factory(unsigned verbosity = 0) : verbosity(verbosity) {
117+
mp.initialize();
118+
}
119+
120+
StateSpace CreateStateSpace() const {
121+
StateSpace::Parameter param;
122+
param.verbosity = verbosity;
123+
124+
return StateSpace{mp, param};
125+
}
126+
127+
Simulator CreateSimulator() const {
128+
return Simulator{};
129+
}
130+
131+
MultiProcessCuStateVecEx mp;
132+
unsigned verbosity;
133+
};
134+
135+
using Simulator = Factory::Simulator;
136+
using StateSpace = Simulator::StateSpace;
137+
using State = StateSpace::State;
138+
using Runner = CuStateVecExRunner<IO, Factory>;
139+
140+
Factory factory(opt.verbosity);
141+
142+
StateSpace state_space = factory.CreateStateSpace();
143+
State state = state_space.Create(circuit.num_qubits);
144+
145+
if (state_space.IsNull(state)) {
146+
IO::errorf("not enough memory: is the number of qubits too large?\n");
147+
return 1;
148+
}
149+
150+
state_space.SetStateZero(state);
151+
152+
Runner::Parameter param;
153+
param.seed = opt.seed;
154+
param.verbosity = opt.verbosity;
155+
156+
if (Runner::Run(param, factory, circuit, state)) {
157+
PrintAmplitudes(circuit.num_qubits, state_space, state);
158+
}
159+
160+
return 0;
161+
}

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+

lib/BUILD

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,8 +186,10 @@ cuda_library(
186186
"matrix.h",
187187
"mps_simulator.h",
188188
"mps_statespace.h",
189+
"multiprocess_custatevecex.h",
189190
"parfor.h",
190191
"qtrajectory.h",
192+
"run_custatevecex.h",
191193
"run_qsim.h",
192194
"run_qsimh.h",
193195
"seqfor.h",
@@ -198,12 +200,14 @@ cuda_library(
198200
"simulator_avx512.h",
199201
"simulator_basic.h",
200202
"simulator_custatevec.h",
203+
"simulator_custatevecex.h",
201204
"simulator_sse.h",
202205
"statespace.h",
203206
"statespace_avx.h",
204207
"statespace_avx512.h",
205208
"statespace_basic.h",
206209
"statespace_custatevec.h",
210+
"statespace_custatevecex.h",
207211
"statespace_sse.h",
208212
"umux.h",
209213
"unitary_calculator_avx.h",
@@ -219,8 +223,10 @@ cuda_library(
219223
"util_cpu.h",
220224
"util_cuda.h",
221225
"util_custatevec.h",
226+
"util_custatevecex.h",
222227
"vectorspace.h",
223228
"vectorspace_cuda.h",
229+
"vectorspace_custatevecex.h",
224230
],
225231
copts = ["-D__CUSTATEVEC__"],
226232
deps = [
@@ -357,6 +363,11 @@ cuda_library(
357363
hdrs = ["util_custatevec.h"],
358364
)
359365

366+
cuda_library(
367+
name = "util_custatevecex",
368+
hdrs = ["util_custatevecex.h"],
369+
)
370+
360371
### Input/output libraries ###
361372

362373
cc_library(
@@ -506,6 +517,29 @@ cc_library(
506517
],
507518
)
508519

520+
cuda_library(
521+
name = "run_custatevecex",
522+
hdrs = ["run_custatevecex.h"],
523+
deps = [
524+
":circuit",
525+
":util",
526+
":util_custatevec",
527+
":util_custatevecex",
528+
],
529+
)
530+
531+
### Multi-process library ###
532+
533+
cuda_library(
534+
name = "multiprocess_custatevecex",
535+
hdrs = ["multiprocess_custatevecex.h"],
536+
deps = [
537+
":io",
538+
":util_custatevec",
539+
":util_custatevecex",
540+
],
541+
)
542+
509543
### Vectorspace libraries ###
510544

511545
cc_library(
@@ -518,6 +552,19 @@ cuda_library(
518552
hdrs = ["vectorspace_cuda.h"],
519553
)
520554

555+
cuda_library(
556+
name = "vectorspace_custatevecex",
557+
hdrs = ["vectorspace_custatevecex.h"],
558+
deps = [
559+
"io",
560+
":multiprocess_custatevecex",
561+
":util_cuda",
562+
":util_custatevec",
563+
":util_custatevecex",
564+
":vectorspace_custatevecex",
565+
],
566+
)
567+
521568
### Statespace libraries ###
522569

523570
cc_library(
@@ -591,6 +638,20 @@ cuda_library(
591638
],
592639
)
593640

641+
cuda_library(
642+
name = "statespace_custatevecex",
643+
hdrs = [
644+
"statespace_custatevecex.h",
645+
],
646+
deps = [
647+
":multiprocess_custatevecex",
648+
":statespace",
649+
":util_custatevec",
650+
":util_custatevecex",
651+
":vectorspace_custatevecex",
652+
],
653+
)
654+
594655
### Simulator libraries ###
595656

596657
cc_library(
@@ -660,6 +721,19 @@ cuda_library(
660721
],
661722
)
662723

724+
cuda_library(
725+
name = "simulator_custatevecex",
726+
hdrs = [
727+
"simulator_custatevecex.h",
728+
],
729+
deps = [
730+
":io",
731+
":statespace_custatevecex",
732+
":util_custatevec",
733+
":util_custatevecex",
734+
],
735+
)
736+
663737
# All three state-vector simulators with multiplexer
664738
cc_library(
665739
name = "simulator",

0 commit comments

Comments
 (0)