generate many lut gpu

zama-ai · Sep 11, 2024 · 7f75c3f · 7f75c3f
1 parent 488c942
commit 7f75c3f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 8 deletions.
diff --git a/tfhe/src/integer/gpu/server_key/radix/mod.rs b/tfhe/src/integer/gpu/server_key/radix/mod.rs
@@ -18,8 +18,8 @@ use crate::integer::gpu::{
     CudaServerKey, PBSType,
 };
 use crate::shortint::ciphertext::{Degree, NoiseLevel};
-use crate::shortint::engine::fill_accumulator;
-use crate::shortint::server_key::{BivariateLookupTableOwned, LookupTableOwned};
+use crate::shortint::engine::{fill_accumulator, fill_many_lut_accumulator};
+use crate::shortint::server_key::{BivariateLookupTableOwned, LookupTableOwned, ManyLookupTableOwned};
 use crate::shortint::PBSOrder;
 
 mod add;
@@ -676,6 +676,33 @@ impl CudaServerKey {
         }
     }
 
+
+    pub(crate) fn generate_many_lookup_table<F>(&self,
+        functions: &[&dyn Fn(u64) -> u64]) -> ManyLookupTableOwned
+    where
+        F: Fn(u64) -> u64,
+    {
+        let (glwe_size, polynomial_size) = match &self.bootstrapping_key {
+            CudaBootstrappingKey::Classic(d_bsk) => {
+                (d_bsk.glwe_dimension.to_glwe_size(), d_bsk.polynomial_size)
+            }
+            CudaBootstrappingKey::MultiBit(d_bsk) => {
+                (d_bsk.glwe_dimension.to_glwe_size(), d_bsk.polynomial_size)
+            }
+        };
+        let mut acc = GlweCiphertext::new(0, glwe_size, polynomial_size, self.ciphertext_modulus);
+
+        let (input_max_degree, sample_extraction_stride, per_function_output_degree) =
+        fill_many_lut_accumulator(&mut acc, polynomial_size, glwe_size, self.message_modulus, self.carry_modulus, functions);
+
+        ManyLookupTableOwned {
+            acc,
+            input_max_degree,
+            sample_extraction_stride,
+            per_function_output_degree,
+        }
+    }
+
     /// Generates a bivariate accumulator
     pub(crate) fn generate_lookup_table_bivariate<F>(&self, f: F) -> BivariateLookupTableOwned
     where

diff --git a/tfhe/src/shortint/engine/mod.rs b/tfhe/src/shortint/engine/mod.rs
@@ -162,30 +162,33 @@ pub(crate) fn fill_accumulator_no_encoding<F, C>(
 /// Fills a GlweCiphertext for use in a ManyLookupTable setting
 pub(crate) fn fill_many_lut_accumulator<C>(
     accumulator: &mut GlweCiphertext<C>,
-    server_key: &ServerKey,
+    polynomial_size: PolynomialSize,
+    glwe_size: GlweSize,
+    message_modulus: MessageModulus,
+    carry_modulus: CarryModulus,
     functions: &[&dyn Fn(u64) -> u64],
 ) -> (MaxDegree, usize, Vec<Degree>)
 where
     C: ContainerMut<Element = u64>,
 {
     assert_eq!(
         accumulator.polynomial_size(),
-        server_key.bootstrapping_key.polynomial_size()
+        polynomial_size
     );
     assert_eq!(
         accumulator.glwe_size(),
-        server_key.bootstrapping_key.glwe_size()
+        glwe_size
     );
 
     let mut accumulator_view = accumulator.as_mut_view();
 
     accumulator_view.get_mut_mask().as_mut().fill(0);
 
     // Modulus of the msg contained in the msg bits and operations buffer
-    let modulus_sup = server_key.message_modulus.0 * server_key.carry_modulus.0;
+    let modulus_sup = message_modulus.0 * carry_modulus.0;
 
     // N/(p/2) = size of each block
-    let box_size = server_key.bootstrapping_key.polynomial_size().0 / modulus_sup;
+    let box_size = polynomial_size.0 / modulus_sup;
 
     // Value of the delta we multiply our messages by
     let delta = (1_u64 << 63) / (modulus_sup as u64);

diff --git a/tfhe/src/shortint/server_key/mod.rs b/tfhe/src/shortint/server_key/mod.rs
@@ -45,6 +45,7 @@ use crate::core_crypto::commons::traits::*;
 use crate::core_crypto::entities::*;
 use crate::core_crypto::fft_impl::fft64::math::fft::Fft;
 use crate::core_crypto::prelude::ComputationBuffers;
+use crate::integer::encryption::KnowsMessageModulus;
 use crate::shortint::ciphertext::{Ciphertext, Degree, MaxDegree, MaxNoiseLevel, NoiseLevel};
 use crate::shortint::client_key::ClientKey;
 use crate::shortint::engine::{
@@ -883,7 +884,8 @@ impl ServerKey {
             self.ciphertext_modulus,
         );
         let (input_max_degree, sample_extraction_stride, per_function_output_degree) =
-            fill_many_lut_accumulator(&mut acc, self, functions);
+            fill_many_lut_accumulator(&mut acc, self.bootstrapping_key.polynomial_size(),
+            self.bootstrapping_key.glwe_size(), self.message_modulus(), self.carry_modulus, functions);
 
         ManyLookupTableOwned {
             acc,