wip

RUSTFLAGS="-C target-cpu=native" cargo run --profile release \
--features=integer --example manual_fft -p tfhe

tfhe/examples/manual_fft.rs

@@ -0,0 +1,37 @@
+use tfhe::core_crypto::fft_impl::fft64::math::fft::{setup_custom_fft_plan, FftAlgo, Method, Plan};
+use tfhe::prelude::*;
+use tfhe::{set_server_key, ClientKey, ConfigBuilder, FheUint64, ServerKey};
+
+pub fn main() {
+    let n = 2048;
+    let my_plan = Plan::new(
+        // n / 2 is due to how TFHE-rs handles ffts
+        n / 2,
+        Method::UserProvided {
+            // User responsibility to choose an algorithm compatible with their n
+            // Both for the algorithm and the base_n
+            base_algo: FftAlgo::Dif4,
+            base_n: n / 2,
+        },
+    );
+
+    setup_custom_fft_plan(my_plan);
+
+    let config = ConfigBuilder::default().build();
+    let cks = ClientKey::generate(config);
+    let sks = ServerKey::new(&cks);
+
+    let msg_a: u64 = 42;
+    let msg_b: u64 = 69;
+
+    let a = FheUint64::encrypt(msg_a, &cks);
+    let b = FheUint64::encrypt(msg_b, &cks);
+
+    set_server_key(sks);
+
+    let c = &a * &b;
+
+    let res: u64 = c.decrypt(&cks);
+
+    assert_eq!(res, msg_a.wrapping_mul(msg_b));
+}

tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs

@@ -19,7 +19,8 @@ use std::sync::{Arc, OnceLock, RwLock};
 #[cfg(not(feature = "experimental-force_fft_algo_dif4"))]
 use std::time::Duration;
 use tfhe_fft::c64;
-use tfhe_fft::unordered::{Method, Plan};
+pub use tfhe_fft::ordered::FftAlgo;
+pub use tfhe_fft::unordered::{Method, Plan};
 use tfhe_versionable::{Unversionize, UnversionizeError, Versionize, VersionizeOwned};
 
 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
@@ -105,6 +106,26 @@ fn plans() -> &'static PlanMap {
     PLANS.get_or_init(|| RwLock::new(HashMap::new()))
 }
 
+pub fn setup_custom_fft_plan(plan: Plan) {
+    let base_n = plan.fft_size();
+    let n = base_n * 2;
+
+    let plan = Arc::new((Twisties::new(base_n), plan));
+
+    let global_plans = plans();
+
+    let mut write = global_plans.write().unwrap();
+
+    match write.entry(n) {
+        Entry::Occupied(mut occupied_entry) => occupied_entry.get_mut().set(plan).unwrap(),
+        Entry::Vacant(vacant_entry) => {
+            let lock = OnceLock::new();
+            let _ = lock.get_or_init(|| plan);
+            vacant_entry.insert(Arc::new(lock));
+        }
+    }
+}
+
 /// Return the input slice, cast to the same type.
 ///
 /// This is useful when the fact that `From` and `To` are the same type cannot be proven in the