chore(gpu): replace cudaStreamCaptureModeGlobal by cudaStreamCaptureM…

…odeThreadLocal to avoid CPU multi-thread issues
zama-ai · Jan 6, 2025 · f9e94d9 · f9e94d9
1 parent d66e36b
commit f9e94d9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_classic.cuh b/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_classic.cuh
@@ -485,7 +485,7 @@ __host__ void host_programmable_bootstrap(
   cudaGraphExec_t instance;
   for (int i = 0; i < lwe_dimension; i++) {
     if (!graphCreated) {
-      cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
+      cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal);
       execute_step_one<Torus, params>(
           stream, gpu_index, lut_vector, lut_vector_indexes, lwe_array_in,
           lwe_input_indexes, bootstrapping_key, global_accumulator,

diff --git a/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh b/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh
@@ -657,7 +657,7 @@ __host__ void host_multi_bit_programmable_bootstrap(
        lwe_offset += lwe_chunk_size) {
 
     if (!graphCreated) {
-      cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
+      cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal);
       // Compute a keybundle
       execute_compute_keybundle<Torus, params>(
           stream, gpu_index, lwe_array_in, lwe_input_indexes, bootstrapping_key,