Add dem_sampling CPU/GPU across C++ and Python

.github/actions/build-lib/build_qec.sh

@@ -26,8 +26,10 @@ cmake -S libs/qec -B "$1" \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache \
   -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DCUDAQ_DIR=/cudaq-install/lib/cmake/cudaq/ \
+  -DCUQUANTUM_INSTALL_PREFIX=/cudaq-install \
   -DCUDAQX_INCLUDE_TESTS=ON \
   -DCUDAQX_BINDINGS_PYTHON=ON \
+  -DCUDAQ_QEC_REQUIRE_CUSTABILIZER=ON \
   -DCMAKE_INSTALL_PREFIX="$2" \
   -DCUDAQ_REALTIME_ROOT=$CUDAQ_REALTIME_ROOT
 

cmake/Modules/FindcuStabilizer.cmake

@@ -0,0 +1,96 @@
+# ============================================================================ #
+# Copyright (c) 2024 - 2026 NVIDIA Corporation & Affiliates.                   #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.    #
+# ============================================================================ #
+
+#[=======================================================================[.rst:
+FindcuStabilizer
+----------------
+
+Find the cuStabilizer library (shipped inside cuQuantum).
+
+Uses the same ``CUQUANTUM_INSTALL_PREFIX`` convention as CUDA-Q.
+
+Imported Targets
+^^^^^^^^^^^^^^^^
+
+``cuStabilizer::cuStabilizer``
+  The cuStabilizer library.
+
+Result Variables
+^^^^^^^^^^^^^^^^
+
+``cuStabilizer_FOUND``
+``cuStabilizer_INCLUDE_DIR``
+``cuStabilizer_LIBRARY``
+
+Hints
+^^^^^
+
+``CUSTABILIZER_ROOT``
+  Preferred search prefix.
+``CUQUANTUM_INSTALL_PREFIX``
+  cuQuantum installation prefix (same as CUDA-Q convention).
+
+#]=======================================================================]
+
+if(NOT CUSTABILIZER_ROOT AND NOT CUQUANTUM_INSTALL_PREFIX)
+  set(CUQUANTUM_INSTALL_PREFIX "$ENV{CUQUANTUM_INSTALL_PREFIX}" CACHE PATH
+    "Path to cuQuantum installation")
+endif()
+
+find_path(cuStabilizer_INCLUDE_DIR
+  NAMES custabilizer.h
+  HINTS
+    ${CUSTABILIZER_ROOT}/include
+    ${CUQUANTUM_INSTALL_PREFIX}/include
+)
+
+find_library(cuStabilizer_LIBRARY
+  NAMES custabilizer libcustabilizer.so.0
+  HINTS
+    ${CUSTABILIZER_ROOT}/lib64
+    ${CUSTABILIZER_ROOT}/lib
+    ${CUQUANTUM_INSTALL_PREFIX}/lib64
+    ${CUQUANTUM_INSTALL_PREFIX}/lib
+)
+
+set(cuStabilizer_VERSION "")
+if(cuStabilizer_INCLUDE_DIR AND EXISTS "${cuStabilizer_INCLUDE_DIR}/custabilizer.h")
+  file(STRINGS "${cuStabilizer_INCLUDE_DIR}/custabilizer.h"
+       _custab_major_line REGEX "^#define[ \t]+CUSTABILIZER_MAJOR[ \t]+[0-9]+")
+  file(STRINGS "${cuStabilizer_INCLUDE_DIR}/custabilizer.h"
+       _custab_minor_line REGEX "^#define[ \t]+CUSTABILIZER_MINOR[ \t]+[0-9]+")
+  file(STRINGS "${cuStabilizer_INCLUDE_DIR}/custabilizer.h"
+       _custab_patch_line REGEX "^#define[ \t]+CUSTABILIZER_PATCH[ \t]+[0-9]+")
+  if(_custab_major_line AND _custab_minor_line AND _custab_patch_line)
+    string(REGEX REPLACE "^#define[ \t]+CUSTABILIZER_MAJOR[ \t]+([0-9]+)$" "\\1"
+           _custab_major "${_custab_major_line}")
+    string(REGEX REPLACE "^#define[ \t]+CUSTABILIZER_MINOR[ \t]+([0-9]+)$" "\\1"
+           _custab_minor "${_custab_minor_line}")
+    string(REGEX REPLACE "^#define[ \t]+CUSTABILIZER_PATCH[ \t]+([0-9]+)$" "\\1"
+           _custab_patch "${_custab_patch_line}")
+    set(cuStabilizer_VERSION "${_custab_major}.${_custab_minor}.${_custab_patch}")
+  endif()
+endif()
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(cuStabilizer
+  REQUIRED_VARS cuStabilizer_INCLUDE_DIR cuStabilizer_LIBRARY
+  VERSION_VAR cuStabilizer_VERSION
+)
+
+if(cuStabilizer_FOUND AND NOT TARGET cuStabilizer::cuStabilizer)
+  add_library(cuStabilizer::cuStabilizer INTERFACE IMPORTED GLOBAL)
+  set_target_properties(cuStabilizer::cuStabilizer PROPERTIES
+    INTERFACE_INCLUDE_DIRECTORIES "${cuStabilizer_INCLUDE_DIR}"
+  )
+  if(cuStabilizer_LIBRARY)
+    set_target_properties(cuStabilizer::cuStabilizer PROPERTIES
+      INTERFACE_LINK_LIBRARIES "${cuStabilizer_LIBRARY}"
+    )
+  endif()
+endif()

libs/qec/CMakeLists.txt

@@ -61,6 +61,11 @@ option(CUDAQX_QEC_INSTALL_PYTHON
        "Install python files alongside the library."
        ${CUDAQX_INSTALL_PYTHON})
 
+# Require cuStabilizer + CUDA (with sparse DEM APIs) for QEC builds.
+option(CUDAQ_QEC_REQUIRE_CUSTABILIZER
+       "Fail configuration when cuStabilizer/CUDA or required DEM APIs are unavailable."
+       ON)
+
 # Option to control TRT decoder build (default: ON)
 option(CUDAQ_QEC_BUILD_TRT_DECODER "Build the TensorRT decoder plugin" ON)
 

libs/qec/include/cudaq/qec/dem_sampling.h

@@ -0,0 +1,71 @@
+/****************************************************************-*- C++ -*-****
+ * Copyright (c) 2024 - 2026 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+#pragma once
+
+#include "cuda-qx/core/tensor.h"
+#include <cstdint>
+#include <optional>
+#include <tuple>
+#include <vector>
+
+namespace cudaq::qec::dem_sampler {
+
+/// CPU implementation of DEM sampling. This is the existing implementation
+/// that uses std::bernoulli_distribution and tensor dot product.
+namespace cpu {
+
+/// @brief Sample measurements from a check matrix (CPU, per-mechanism probs)
+/// @param check_matrix Binary matrix [num_checks × num_error_mechanisms]
+/// @param numShots Number of measurement shots
+/// @param error_probabilities Per-error-mechanism probabilities
+/// @return (checks [numShots × num_checks], errors [numShots × num_mechanisms])
+std::tuple<cudaqx::tensor<uint8_t>, cudaqx::tensor<uint8_t>>
+sample_dem(const cudaqx::tensor<uint8_t> &check_matrix, std::size_t numShots,
+           const std::vector<double> &error_probabilities);
+
+/// @brief Sample measurements from a check matrix with seed (CPU)
+std::tuple<cudaqx::tensor<uint8_t>, cudaqx::tensor<uint8_t>>
+sample_dem(const cudaqx::tensor<uint8_t> &check_matrix, std::size_t numShots,
+           const std::vector<double> &error_probabilities, unsigned seed);
+
+} // namespace cpu
+
+/// GPU implementation of DEM sampling via cuStabilizer C API.
+///
+/// The caller provides device pointers. The function composes:
+///   1. pack_check_matrix_rowwise (dense uint8 → bitpacked uint32)
+///   2. custabilizerSampleProbArraySparseCompute (sparse Bernoulli sampling)
+///   3. custabilizerGF2SparseDenseMatrixMultiply (syndrome = errors * H^T)
+///   4. unpack_syndromes_gpu (bitpacked uint32 → uint8)
+///   5. csr_to_dense_fused (CSR errors → dense uint8)
+///
+/// All GPU memory is allocated and freed internally per call.
+namespace gpu {
+
+/// @brief GPU DEM sampling with device pointers
+///
+/// @param d_check_matrix Device pointer [num_checks × num_error_mechanisms]
+/// @param num_checks Number of checks (rows of H)
+/// @param num_error_mechanisms Number of error mechanisms (columns of H)
+/// @param d_error_probabilities Device pointer [num_error_mechanisms]
+/// @param num_shots Number of samples
+/// @param seed RNG seed
+/// @param d_checks_out Device pointer [num_shots × num_checks] (OUTPUT)
+/// @param d_errors_out Device pointer [num_shots × num_error_mechanisms] (OUT)
+/// @param stream_handle Optional CUDA stream handle (uintptr_t cast), 0 for
+///        default stream
+/// @return true on success, false if cuStabilizer is unavailable
+bool sample_dem(const uint8_t *d_check_matrix, size_t num_checks,
+                size_t num_error_mechanisms,
+                const double *d_error_probabilities, size_t num_shots,
+                unsigned seed, uint8_t *d_checks_out, uint8_t *d_errors_out,
+                std::uintptr_t stream_handle = 0);
+
+} // namespace gpu
+
+} // namespace cudaq::qec::dem_sampler

libs/qec/include/cudaq/qec/experiments.h

@@ -9,6 +9,8 @@
 
 #include "cudaq/qec/code.h"
 #include "cudaq/qec/detector_error_model.h"
+#include <tuple>
+#include <vector>
 
 namespace cudaq::qec {
 
@@ -61,6 +63,25 @@ std::tuple<cudaqx::tensor<uint8_t>, cudaqx::tensor<uint8_t>>
 sample_code_capacity(const code &code, std::size_t numShots,
                      double error_probability);
 
+/// @brief Sample a detector error model on CPU (legacy API).
+/// @param check_matrix Binary matrix [num_checks x num_error_mechanisms]
+/// @param numShots Number of independent Monte-Carlo shots
+/// @param error_probabilities Per-error-mechanism Bernoulli probabilities
+/// @return Tuple of (checks, errors)
+std::tuple<cudaqx::tensor<uint8_t>, cudaqx::tensor<uint8_t>>
+dem_sampling(const cudaqx::tensor<uint8_t> &check_matrix, std::size_t numShots,
+             const std::vector<double> &error_probabilities);
+
+/// @brief Sample a detector error model on CPU (legacy API, seeded).
+/// @param check_matrix Binary matrix [num_checks x num_error_mechanisms]
+/// @param numShots Number of independent Monte-Carlo shots
+/// @param error_probabilities Per-error-mechanism Bernoulli probabilities
+/// @param seed RNG seed for reproducibility
+/// @return Tuple of (checks, errors)
+std::tuple<cudaqx::tensor<uint8_t>, cudaqx::tensor<uint8_t>>
+dem_sampling(const cudaqx::tensor<uint8_t> &check_matrix, std::size_t numShots,
+             const std::vector<double> &error_probabilities, unsigned seed);
+
 /// @brief Sample syndrome measurements with circuit-level noise
 /// @param statePrep Initial state preparation operation
 /// @param numShots Number of measurement shots

libs/qec/lib/CMakeLists.txt

@@ -10,8 +10,75 @@ set(LIBRARY_NAME cudaq-qec)
 
 add_compile_options(-Wno-attributes) 
 
-# FIXME?: This must be a shared library. Trying to build a static one will fail.
-add_library(${LIBRARY_NAME} SHARED
+# cuStabilizer GPU DEM sampling support (required by default).
+find_package(cuStabilizer QUIET)
+set(_cuStabilizer_location
+    "include='${cuStabilizer_INCLUDE_DIR}', library='${cuStabilizer_LIBRARY}'")
+if(cuStabilizer_VERSION)
+  string(APPEND _cuStabilizer_location ", version='${cuStabilizer_VERSION}'")
+endif()
+set(CUDAQ_QEC_CUSTABILIZER_HAS_SPARSE_DEM_APIS FALSE)
+if(cuStabilizer_FOUND AND CMAKE_CUDA_COMPILER)
+  include(CheckCXXSourceCompiles)
+  set(CMAKE_REQUIRED_INCLUDES "${cuStabilizer_INCLUDE_DIR}")
+  if(CUDAToolkit_INCLUDE_DIRS)
+    list(APPEND CMAKE_REQUIRED_INCLUDES ${CUDAToolkit_INCLUDE_DIRS})
+  endif()
+  check_cxx_source_compiles(
+    "
+    #include <custabilizer.h>
+    int main() {
+      (void)CUSTABILIZER_STATUS_INSUFFICIENT_SPARSE_STORAGE;
+      auto prepare = &custabilizerSampleProbArraySparsePrepare;
+      auto compute = &custabilizerSampleProbArraySparseCompute;
+      auto spmm = &custabilizerGF2SparseDenseMatrixMultiply;
+      (void)prepare;
+      (void)compute;
+      (void)spmm;
+      return 0;
+    }
+    "
+    CUDAQ_QEC_CUSTABILIZER_HAS_SPARSE_DEM_APIS
+  )
+  unset(CMAKE_REQUIRED_INCLUDES)
+endif()
+
+if(cuStabilizer_FOUND AND CMAKE_CUDA_COMPILER AND
+   CUDAQ_QEC_CUSTABILIZER_HAS_SPARSE_DEM_APIS)
+  message(STATUS
+          "cuStabilizer found (${_cuStabilizer_location}) - GPU DEM sampling enabled")
+  set(CUDAQ_QEC_HAS_CUSTABILIZER TRUE)
+  set(CUDAQ_QEC_HAS_CUSTABILIZER TRUE PARENT_SCOPE)
+elseif(cuStabilizer_FOUND AND CMAKE_CUDA_COMPILER)
+  message(WARNING
+    "cuStabilizer found (${_cuStabilizer_location}), but sparse DEM APIs are "
+    "missing. GPU DEM sampling disabled. Required APIs: "
+    "custabilizerSampleProbArraySparsePrepare, "
+    "custabilizerSampleProbArraySparseCompute, "
+    "custabilizerGF2SparseDenseMatrixMultiply."
+  )
+  set(CUDAQ_QEC_HAS_CUSTABILIZER FALSE)
+  set(CUDAQ_QEC_HAS_CUSTABILIZER FALSE PARENT_SCOPE)
+else()
+  message(STATUS "cuStabilizer not found or CUDA unavailable – GPU DEM sampling disabled")
+  set(CUDAQ_QEC_HAS_CUSTABILIZER FALSE)
+  set(CUDAQ_QEC_HAS_CUSTABILIZER FALSE PARENT_SCOPE)
+endif()
+
+if(CUDAQ_QEC_REQUIRE_CUSTABILIZER AND NOT CUDAQ_QEC_HAS_CUSTABILIZER)
+  message(FATAL_ERROR
+    "CUDAQ_QEC_REQUIRE_CUSTABILIZER is ON, but required cuStabilizer/CUDA "
+    "support was not detected.\n"
+    "Detected: cuStabilizer_FOUND='${cuStabilizer_FOUND}', "
+    "${_cuStabilizer_location}, "
+    "cuda_compiler='${CMAKE_CUDA_COMPILER}', "
+    "sparse_dem_apis='${CUDAQ_QEC_CUSTABILIZER_HAS_SPARSE_DEM_APIS}'.\n"
+    "Required APIs: custabilizerSampleProbArraySparsePrepare, "
+    "custabilizerSampleProbArraySparseCompute, "
+    "custabilizerGF2SparseDenseMatrixMultiply.")
+endif()
+
+set(QEC_SOURCES
   code.cpp
   decoder.cpp
   detector_error_model.cpp
@@ -24,6 +91,16 @@ add_library(${LIBRARY_NAME} SHARED
   version.cpp
 )
 
+if(CUDAQ_QEC_HAS_CUSTABILIZER)
+  list(APPEND QEC_SOURCES
+    dem_sampling/dem_sampling_utils.cu
+    dem_sampling/dem_sampling_gpu.cpp
+  )
+endif()
+
+# FIXME?: This must be a shared library. Trying to build a static one will fail.
+add_library(${LIBRARY_NAME} SHARED ${QEC_SOURCES})
+
 add_subdirectory(decoders/plugins/example)
 add_subdirectory(decoders/plugins/pymatching)
 
@@ -38,7 +115,10 @@ if (CUDAQX_QEC_USE_FLOAT)
    target_compile_definitions(${LIBRARY_NAME} PUBLIC -DCUDAQX_QEC_FLOAT_TYPE=float)
 endif()
 
+
+
 target_include_directories(${LIBRARY_NAME}
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}
   PUBLIC 
     $<BUILD_INTERFACE:${CUDAQX_QEC_INCLUDE_DIR}>
     $<INSTALL_INTERFACE:${CUDAQ_INCLUDE_DIR}>
@@ -57,6 +137,18 @@ target_link_libraries(${LIBRARY_NAME}
     cudaq::cudaq-common
 )
 
+if(CUDAQ_QEC_HAS_CUSTABILIZER)
+  target_compile_definitions(${LIBRARY_NAME} PUBLIC CUDAQ_QEC_HAS_CUSTABILIZER)
+  target_link_libraries(${LIBRARY_NAME} PRIVATE
+    cuStabilizer::cuStabilizer
+    CUDA::cudart
+  )
+  set_source_files_properties(
+    dem_sampling/dem_sampling_utils.cu
+    PROPERTIES LANGUAGE CUDA
+  )
+endif()
+
 set_target_properties(${LIBRARY_NAME} PROPERTIES
   LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)