fix(gpu): fix default pbs with many luts

zama-ai · Oct 14, 2024 · 748ec04 · 748ec04
1 parent ff0609f
commit 748ec04
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 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
@@ -222,9 +222,9 @@ __global__ void __launch_bounds__(params::degree / params::opt)
         for (int i = 1; i < lut_count; i++) {
           auto next_lwe_array_out =
               lwe_array_out +
-              (i * gridDim.z * (glwe_dimension * polynomial_size + 1));
+              (i * gridDim.x * (glwe_dimension * polynomial_size + 1));
           auto next_block_lwe_array_out =
-              &next_lwe_array_out[lwe_output_indexes[blockIdx.z] *
+              &next_lwe_array_out[lwe_output_indexes[blockIdx.x] *
                                       (glwe_dimension * polynomial_size + 1) +
                                   blockIdx.y * polynomial_size];
 
@@ -239,9 +239,9 @@ __global__ void __launch_bounds__(params::degree / params::opt)
 
           auto next_lwe_array_out =
               lwe_array_out +
-              (i * gridDim.z * (glwe_dimension * polynomial_size + 1));
+              (i * gridDim.x * (glwe_dimension * polynomial_size + 1));
           auto next_block_lwe_array_out =
-              &next_lwe_array_out[lwe_output_indexes[blockIdx.z] *
+              &next_lwe_array_out[lwe_output_indexes[blockIdx.x] *
                                       (glwe_dimension * polynomial_size + 1) +
                                   blockIdx.y * polynomial_size];
 

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
@@ -331,9 +331,9 @@ __global__ void __launch_bounds__(params::degree / params::opt)
         for (int i = 1; i < lut_count; i++) {
           auto next_lwe_array_out =
               lwe_array_out +
-              (i * gridDim.z * (glwe_dimension * polynomial_size + 1));
+              (i * gridDim.x * (glwe_dimension * polynomial_size + 1));
           auto next_block_lwe_array_out =
-              &next_lwe_array_out[lwe_output_indexes[blockIdx.z] *
+              &next_lwe_array_out[lwe_output_indexes[blockIdx.x] *
                                       (glwe_dimension * polynomial_size + 1) +
                                   blockIdx.y * polynomial_size];
 
@@ -348,9 +348,9 @@ __global__ void __launch_bounds__(params::degree / params::opt)
 
           auto next_lwe_array_out =
               lwe_array_out +
-              (i * gridDim.z * (glwe_dimension * polynomial_size + 1));
+              (i * gridDim.x * (glwe_dimension * polynomial_size + 1));
           auto next_block_lwe_array_out =
-              &next_lwe_array_out[lwe_output_indexes[blockIdx.z] *
+              &next_lwe_array_out[lwe_output_indexes[blockIdx.x] *
                                       (glwe_dimension * polynomial_size + 1) +
                                   blockIdx.y * polynomial_size];
 

diff --git a/tfhe/src/integer/gpu/server_key/radix/mod.rs b/tfhe/src/integer/gpu/server_key/radix/mod.rs
@@ -976,8 +976,8 @@ impl CudaServerKey {
     ///     // Generate the lookup table for the functions
     ///     // f1: x -> x*x mod 4
     ///     // f2: x -> count_ones(x as binary) mod 4
-    ///     let f1 = |x: u64| x.pow(2) % 4;
-    ///     let f2 = |x: u64| x.count_ones() as u64 % 4;
+    ///     let f1 = |x: u64| x.pow(2) % 8;
+    ///     let f2 = |x: u64| x.count_ones() as u64 % 8;
     ///     // Easy to use for generation
     ///     let luts = sks.generate_many_lookup_table(&[&f1, &f2]);
     ///     let vec_res = unsafe { sks.apply_many_lookup_table_async(&d_ct.as_ref(), &luts, &stream) };