From ae907ff05f82464e9218af4d1bed93788c0cef4c Mon Sep 17 00:00:00 2001
From: Agnes Leroy <agnes.leroy@zama.ai>
Date: Thu, 1 Aug 2024 10:15:48 +0200
Subject: [PATCH] chore(gpu): define higher values for the sm size based on
 compute capability

---
 backends/tfhe-cuda-backend/cuda/src/device.cu | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/backends/tfhe-cuda-backend/cuda/src/device.cu b/backends/tfhe-cuda-backend/cuda/src/device.cu
index 29bea24de6..14aad18a70 100644
--- 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;
 }