Add zero-copy GPU state vector access (#836)

dev_tools/compiler_probe.bzl

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Repository rule to probe compiler support for SFrame flags."""
+"""Repository rule to probe the C++ compiler's SFrame support and -march=native."""
 
 def _compiler_probe_impl(repository_ctx):
     # Try to determine the compiler. Prefer CC from environment.
@@ -33,8 +33,35 @@ def _compiler_probe_impl(repository_ctx):
 
     supports_gsframe = (res.return_code == 0)
 
+    # Identify what -march=native selects on this machine. Code compiled with
+    # -march=native depends on the build machine's CPU, but the compile command
+    # line does not, so machines with the same toolchain compute the same action
+    # key for it. A shared disk or remote cache can then hand an object built on
+    # a CPU with, e.g., AVX-512 to a machine without it, where the test dies with
+    # an illegal instruction. Adding this ID to such command lines prevents that.
+    # The ID is a hash of the compiler's predefined macros, which name the
+    # instruction set extensions it may use (__AVX2__, __AVX512F__, ...), sorted
+    # so that their order doesn't matter. Use the same compiler as Bazel's
+    # auto-configured C++ toolchain: $CC if set, otherwise gcc.
+    res = repository_ctx.execute([
+        repository_ctx.os.environ.get("CC", "gcc"),
+        "-march=native",
+        "-dM",
+        "-E",
+        "-x",
+        "c++",
+        "/dev/null",
+    ])
+    if res.return_code == 0:
+        native_arch_id = str(hash("\n".join(sorted(res.stdout.splitlines()))))
+    else:
+        native_arch_id = "unknown"
+
     repository_ctx.file("BUILD.bazel", "package(default_visibility = ['//visibility:public'])\n")
-    repository_ctx.file("compiler_config.bzl", "SUPPORTS_GSFRAME = %s\n" % supports_gsframe)
+    repository_ctx.file(
+        "compiler_config.bzl",
+        "SUPPORTS_GSFRAME = %s\nNATIVE_ARCH_ID = \"%s\"\n" % (supports_gsframe, native_arch_id),
+    )
 
 compiler_probe = repository_rule(
     implementation = _compiler_probe_impl,

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 {

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<void*>(ptr_.get());
+    }
+
     unsigned num_qubits() const {
       return num_qubits_;
     }

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,

pybind_interface/cuda/CMakeLists.txt

@@ -52,8 +52,16 @@ else()  # means pybind11 has been fetched in GetPybind11.cmake
 endif()
 
 add_library(qsim_cuda MODULE pybind_main_cuda.cpp)
+# nvcc's cudafe++ rewrites C++ anonymous namespaces into named
+# `_GLOBAL__N__<hash>` namespaces with default ELF visibility unless
+# CUDA_VISIBILITY_PRESET hidden is set. Setting hidden visibility prevents
+# internal SimulatorHelper / DeviceStateVector symbols from leaking into
+# .dynsym and matches pybind11's standard module visibility configuration.
 set_target_properties(qsim_cuda PROPERTIES
     CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}"
+    CXX_VISIBILITY_PRESET hidden
+    CUDA_VISIBILITY_PRESET hidden
+    VISIBILITY_INLINES_HIDDEN ON
     PREFIX "${PYTHON_MODULE_PREFIX}"
     SUFFIX "${PYTHON_MODULE_EXTENSION}"
 )

pybind_interface/cuda/pybind_main_cuda.cpp

@@ -12,6 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Enable zero-copy device state-vector bindings (issue #836).
+// Must precede the header include below: pybind_main.h branches on
+// this macro and would otherwise emit a conflicting inline stub.
+#define QSIM_DEVICE_STATE_BINDINGS
+
 #include "pybind_main_cuda.h"
 
 #include "../../lib/fuser_mqubit.h"

pybind_interface/custatevec/CMakeLists.txt

@@ -47,7 +47,12 @@ include_directories($ENV{CUQUANTUM_ROOT}/include)
 link_directories($ENV{CUQUANTUM_ROOT}/lib $ENV{CUQUANTUM_ROOT}/lib64)
 
 add_library(qsim_custatevec MODULE pybind_main_custatevec.cpp)
+# See pybind_interface/cuda/CMakeLists.txt for why hidden visibility is set
+# on nvcc-compiled pybind11 extension modules.
 set_target_properties(qsim_custatevec PROPERTIES
+    CXX_VISIBILITY_PRESET hidden
+    CUDA_VISIBILITY_PRESET hidden
+    VISIBILITY_INLINES_HIDDEN ON
     PREFIX "${PYTHON_MODULE_PREFIX}"
     SUFFIX "${PYTHON_MODULE_EXTENSION}"
 )

pybind_interface/custatevec/pybind_main_custatevec.cpp

@@ -15,6 +15,11 @@
 #include <cublas_v2.h>
 #include <custatevec.h>
 
+// Enable zero-copy device state-vector bindings (issue #836).
+// Must precede the header include below: pybind_main.h branches on
+// this macro and would otherwise emit a conflicting inline stub.
+#define QSIM_DEVICE_STATE_BINDINGS
+
 #include "pybind_main_custatevec.h"
 
 #include "../../lib/fuser_mqubit.h"

pybind_interface/custatevecex/CMakeLists.txt

@@ -47,7 +47,12 @@ include_directories($ENV{CUQUANTUM_ROOT}/include)
 link_directories($ENV{CUQUANTUM_ROOT}/lib $ENV{CUQUANTUM_ROOT}/lib64)
 
 add_library(qsim_custatevecex MODULE pybind_main_custatevecex.cpp)
+# See pybind_interface/cuda/CMakeLists.txt for why hidden visibility is set
+# on nvcc-compiled pybind11 extension modules.
 set_target_properties(qsim_custatevecex PROPERTIES
+    CXX_VISIBILITY_PRESET hidden
+    CUDA_VISIBILITY_PRESET hidden
+    VISIBILITY_INLINES_HIDDEN ON
     PREFIX "${PYTHON_MODULE_PREFIX}"
     SUFFIX "${PYTHON_MODULE_EXTENSION}"
 )

pybind_interface/custatevecex/pybind_main_custatevecex.cpp

@@ -16,6 +16,11 @@
 
 #include <custatevecEx.h>
 
+// Enable zero-copy device state-vector bindings (issue #836).
+// Must precede the header include below: pybind_main.h branches on
+// this macro and would otherwise emit a conflicting inline stub.
+#define QSIM_DEVICE_STATE_BINDINGS
+
 #include "pybind_main_custatevecex.h"
 
 #include "../../lib/fuser_mqubit.h"

pybind_interface/hip/pybind_main_hip.cpp

@@ -12,6 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Enable zero-copy device state-vector bindings (issue #836).
+// Must precede the header include below: pybind_main.h branches on
+// this macro and would otherwise emit a conflicting inline stub.
+#define QSIM_DEVICE_STATE_BINDINGS
+
 #include "pybind_main_hip.h"
 
 #include "../../lib/fuser_mqubit.h"