chore(gpu): try resetting device to gpu 0 after sync local streams

zama-ai · Nov 25, 2024 · de3e09d · de3e09d
1 parent 2a18d6f
commit de3e09d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/div_rem.cuh b/backends/tfhe-cuda-backend/cuda/src/integer/div_rem.cuh
@@ -592,6 +592,7 @@ __host__ void host_unsigned_integer_div_rem_kb(
     cuda_synchronize_stream(mem_ptr->sub_streams_1[j], gpu_indexes[j]);
     cuda_synchronize_stream(mem_ptr->sub_streams_2[j], gpu_indexes[j]);
   }
+  cudaSetDevice(gpu_indexes[0]);
 }
 
 template <typename Torus>
@@ -688,6 +689,7 @@ __host__ void host_integer_div_rem_kb(cudaStream_t const *streams,
       cuda_synchronize_stream(int_mem_ptr->sub_streams_1[j], gpu_indexes[j]);
       cuda_synchronize_stream(int_mem_ptr->sub_streams_2[j], gpu_indexes[j]);
     }
+    cudaSetDevice(gpu_indexes[0]);
   } else {
     host_unsigned_integer_div_rem_kb<Torus>(
         streams, gpu_indexes, gpu_count, quotient, remainder, numerator,

diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/scalar_shifts.cuh b/backends/tfhe-cuda-backend/cuda/src/integer/scalar_shifts.cuh
@@ -235,6 +235,7 @@ __host__ void host_integer_radix_arithmetic_scalar_shift_kb_inplace(
         cuda_synchronize_stream(mem->local_streams_1[j], gpu_indexes[j]);
         cuda_synchronize_stream(mem->local_streams_2[j], gpu_indexes[j]);
       }
+      cudaSetDevice(gpu_indexes[0]);
     }
   } else {
     PANIC("Cuda error (scalar shift): left scalar shift is never of the "