chore(gpu): define higher values for the sm size based on compute cap…

…ability
zama-ai · Aug 1, 2024 · ae907ff · ae907ff
1 parent 02521c0
commit ae907ff
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/backends/tfhe-cuda-backend/cuda/src/device.cu b/backends/tfhe-cuda-backend/cuda/src/device.cu
@@ -247,5 +247,12 @@ int cuda_get_max_shared_memory(uint32_t gpu_index) {
   cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock,
                          gpu_index);
   check_cuda_error(cudaGetLastError());
+#if CUDA_ARCH == 900
+  max_shared_memory = 227000;
+#elif CUDA_ARCH == 800
+  max_shared_memory = 163000;
+#elif CUDA_ARCH == 700
+  max_shared_memory = 96000;
+#endif
   return max_shared_memory;
 }