Add experimental support for cuStateVecEx.

CMakeLists.txt

@@ -64,6 +64,7 @@ if(NOT CMAKE_APPLE_SILICON_PROCESSOR)
         add_subdirectory(pybind_interface/cuda)
         if(DEFINED ENV{CUQUANTUM_ROOT})
             add_subdirectory(pybind_interface/custatevec)
+            add_subdirectory(pybind_interface/custatevecex)
         endif()
     elseif(has_hipcc)
         add_subdirectory(pybind_interface/hip)

Makefile

@@ -94,7 +94,10 @@ ifneq (,$(strip $(CUQUANTUM_ROOT)))
         CUSVFLAGS += -lcustatevec -lcublas
         CUSTATEVECFLAGS ?= $(CUSVFLAGS)
         TARGETS += qsim-custatevec
+        TARGETS += qsim-custatevecex
         TESTS += run-custatevec-tests
+        TESTS += run-custatevecex-tests
+	TESTS += run-custatevecex-mpi-tests
     else
         $(warning $$CUQUANTUM_ROOT is set, but the path does not seem to exist)
     endif
@@ -120,6 +123,10 @@ qsim-cuda:
 qsim-custatevec: | check-cuquantum-root-set
 	$(MAKE) -C apps/ qsim-custatevec
 
+.PHONY: qsim-custatevecex
+qsim-custatevecex: | check-cuquantum-root-set
+	$(MAKE) -C apps/ qsim-custatevecex
+
 .PHONY: qsim-hip
 qsim-hip:
 	$(MAKE) -C apps/ qsim-hip
@@ -140,6 +147,10 @@ cuda-tests:
 custatevec-tests: | check-cuquantum-root-set
 	$(MAKE) -C tests/ custatevec-tests
 
+.PHONY: custatevecex-tests
+custatevecex-tests: | check-cuquantum-root-set
+	$(MAKE) -C tests/ custatevecex-tests
+
 .PHONY: hip-tests
 hip-tests:
 	$(MAKE) -C tests/ hip-tests
@@ -156,6 +167,14 @@ run-cuda-tests: cuda-tests
 run-custatevec-tests: custatevec-tests
 	$(MAKE) -C tests/ run-custatevec-tests
 
+.PHONY: run-custatevecex-tests
+run-custatevecex-tests: custatevecex-tests
+	$(MAKE) -C tests/ run-custatevecex-tests
+
+.PHONY: run-custatevecex-mpi-tests
+run-custatevecex-mpi-tests: custatevecex-tests
+	$(MAKE) -C tests/ run-custatevecex-mpi-tests
+
 .PHONY: run-hip-tests
 run-hip-tests: hip-tests
 	$(MAKE) -C tests/ run-hip-tests

apps/Makefile

@@ -7,6 +7,9 @@ CUDA_TARGETS := $(CUDA_TARGETS:%cuda.cu=%cuda.x)
 CUSTATEVEC_TARGETS = $(shell find . -maxdepth 1 -name "*custatevec.cu")
 CUSTATEVEC_TARGETS := $(CUSTATEVEC_TARGETS:%custatevec.cu=%custatevec.x)
 
+CUSTATEVECEX_TARGETS = $(shell find . -maxdepth 1 -name "*custatevecex.cu")
+CUSTATEVECEX_TARGETS := $(CUSTATEVECEX_TARGETS:%custatevecex.cu=%custatevecex.x)
+
 HIP_TARGETS = $(shell find . -maxdepth 1 -name '*cuda.cu')
 HIP_TARGETS := $(HIP_TARGETS:%cuda.cu=%hip.x)
 
@@ -19,6 +22,9 @@ qsim-cuda: $(CUDA_TARGETS)
 .PHONY: qsim-custatevec
 qsim-custatevec: $(CUSTATEVEC_TARGETS)
 
+.PHONY: qsim-custatevecex
+qsim-custatevecex: $(CUSTATEVECEX_TARGETS)
+
 .PHONY: qsim-hip
 qsim-hip: $(HIP_TARGETS)
 
@@ -31,6 +37,9 @@ qsim-hip: $(HIP_TARGETS)
 %custatevec.x: %custatevec.cu
 	$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
 
+%custatevecex.x: %custatevecex.cu
+	$(NVCC) -o ./$@ $< $(NVCCFLAGS) $(CUSTATEVECFLAGS)
+
 %hip.x: %cuda.cu
 	$(HIPCC) -o ./$@ $< $(HIPCCFLAGS)
 

apps/make.sh

@@ -37,7 +37,8 @@ if command -v nvcc &>/dev/null; then
         )
         nvcc -O3 "${CUSTATEVECFLAGS[@]}" \
              -o qsim_base_custatevec.x qsim_base_custatevec.cu
-
+        nvcc -O3 "${CUSTATEVECFLAGS[@]}" \
+             -o qsim_base_custatevecex.x qsim_base_custatevecex.cu
     fi
 elif command -v hipcc &>/dev/null; then
     hipcc -O3 -o qsim_base_hip.x qsim_base_cuda.cu

apps/qsim_base_custatevecex.cu

@@ -0,0 +1,160 @@
+// Copyright 2025 Google LLC. All Rights Reserved.
+//
+// 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
+//
+//     https://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.
+
+#include <unistd.h>
+
+#include <algorithm>
+#include <complex>
+#include <limits>
+#include <string>
+
+#include "../lib/circuit_qsim_parser.h"
+#include "../lib/formux.h"
+#include "../lib/gates_qsim.h"
+#include "../lib/io_file.h"
+#include "../lib/multiprocess_custatevecex.h"
+#include "../lib/run_custatevecex.h"
+#include "../lib/simulator_custatevecex.h"
+#include "../lib/util_custatevec.h"
+
+struct Options {
+  std::string circuit_file;
+  unsigned maxtime = std::numeric_limits<unsigned>::max();
+  unsigned seed = 1;
+  unsigned verbosity = 0;
+};
+
+Options GetOptions(int argc, char* argv[]) {
+  constexpr char usage[] = "usage:\n  ./qsim_base -c circuit -d maxtime "
+                           "-s seed -v verbosity\n";
+
+  Options opt;
+
+  int k;
+
+  while ((k = getopt(argc, argv, "c:d:s:v:")) != -1) {
+    switch (k) {
+      case 'c':
+        opt.circuit_file = optarg;
+        break;
+      case 'd':
+        opt.maxtime = std::atoi(optarg);
+        break;
+      case 's':
+        opt.seed = std::atoi(optarg);
+        break;
+      case 'v':
+        opt.verbosity = std::atoi(optarg);
+        break;
+      default:
+        qsim::IO::errorf(usage);
+        exit(1);
+    }
+  }
+
+  return opt;
+}
+
+bool ValidateOptions(const Options& opt) {
+  if (opt.circuit_file.empty()) {
+    qsim::IO::errorf("circuit file is not provided.\n");
+    return false;
+  }
+
+  return true;
+}
+
+template <typename StateSpace, typename State>
+void PrintAmplitudes(
+    unsigned num_qubits, const StateSpace& state_space, const State& state) {
+  static constexpr char const* bits[8] = {
+    "000", "001", "010", "011", "100", "101", "110", "111",
+  };
+
+  uint64_t size = std::min(uint64_t{8}, uint64_t{1} << num_qubits);
+  unsigned s = 3 - std::min(unsigned{3}, num_qubits);
+
+  for (uint64_t i = 0; i < size; ++i) {
+    auto a = state_space.GetAmpl(state, i);
+    qsim::IO::messagef("%s:%16.8g%16.8g%16.8g\n",
+                       bits[i] + s, std::real(a), std::imag(a), std::norm(a));
+  }
+}
+
+int main(int argc, char* argv[]) {
+  using namespace qsim;
+
+  auto opt = GetOptions(argc, argv);
+  if (!ValidateOptions(opt)) {
+    return 1;
+  }
+
+  using fp_type = float;
+
+  Circuit<GateQSim<fp_type>> circuit;
+  if (!CircuitQsimParser<IOFile>::FromFile(opt.maxtime, opt.circuit_file,
+                                           circuit)) {
+    return 1;
+  }
+
+  struct Factory {
+    using Simulator = qsim::SimulatorCuStateVecEx<fp_type>;
+    using StateSpace = Simulator::StateSpace;
+
+    explicit Factory(unsigned verbosity = 0) : verbosity(verbosity) {
+      mp.initialize();
+    }
+
+    StateSpace CreateStateSpace() const {
+      StateSpace::Parameter param;
+      param.verbosity = verbosity;
+
+      return StateSpace{mp, param};
+    }
+
+    Simulator CreateSimulator() const {
+      return Simulator{};
+    }
+
+    MultiProcessCuStateVecEx mp;
+    unsigned verbosity;
+  };
+
+  using Simulator = Factory::Simulator;
+  using StateSpace = Simulator::StateSpace;
+  using State = StateSpace::State;
+  using Runner = CuStateVecExRunner<IO, Factory>;
+
+  Factory factory(opt.verbosity);
+
+  StateSpace state_space = factory.CreateStateSpace();
+  State state = state_space.Create(circuit.num_qubits);
+
+  if (state_space.IsNull(state)) {
+    IO::errorf("not enough memory: is the number of qubits too large?\n");
+    return 1;
+  }
+
+  state_space.SetStateZero(state);
+
+  Runner::Parameter param;
+  param.seed = opt.seed;
+  param.verbosity = opt.verbosity;
+
+  if (Runner::Run(param, factory, circuit, state)) {
+    PrintAmplitudes(circuit.num_qubits, state_space, state);
+  }
+
+  return 0;
+}

docs/cirq_interface.md

@@ -186,8 +186,11 @@ library.
 `QSimOptions` provides five parameters to configure GPU execution. `use_gpu`
 is required to enable GPU execution:
 * `use_gpu`: if True, use GPU instead of CPU for simulation.
-* `gpu_mode`: use CUDA if set to 0 (default value) or use the NVIDIA cuStateVec
-library if set to any other value.
+* `gpu_mode`: use CUDA if set to 0 (default value), use the NVIDIA cuStateVec
+if set to 1 or use the NVIDIA cuStateVecEx library if set to any other value.
+
+In the case of the NVIDIA cuStateVecEx library, simulations can be performed
+in multi-device / multi-node environments.
 
 If `use_gpu` is set and `gpu_mode` is set to 0, the remaining parameters can
 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.
 StateSpace. This must be a power of 2 in the range [32, 1024].
 * `gpu_data_blocks`: number of data blocks to use for the GPU StateSpace.
 Below 16 data blocks, performance is noticeably reduced.
+