From 14d49f08919a2a90c8d23960a041fafaf84c82bc Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Thu, 9 Oct 2025 15:21:14 +0200
Subject: [PATCH] chore: add possibility to manually populate tfhe FFT plan
 cache

---
 tfhe/Cargo.toml                               |  5 +++
 tfhe/examples/manual_fft.rs                   | 39 +++++++++++++++++++
 .../fft_impl/fft64/math/fft/mod.rs            | 24 +++++++++++-
 3 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 tfhe/examples/manual_fft.rs

diff --git a/tfhe/Cargo.toml b/tfhe/Cargo.toml
index deaec3ace..23efa2771 100644
--- a/tfhe/Cargo.toml
+++ b/tfhe/Cargo.toml
@@ -199,6 +199,11 @@ required-features = ["shortint", "internal-keycache"]
 
 # Real use-case examples
 
+[[example]]
+name = "manual_fft"
+path = "examples/manual_fft.rs"
+required-features = ["integer"]
+
 [[example]]
 name = "dark_market"
 required-features = ["integer", "internal-keycache"]
diff --git a/tfhe/examples/manual_fft.rs b/tfhe/examples/manual_fft.rs
new file mode 100644
index 000000000..36bd1ce16
--- /dev/null
+++ b/tfhe/examples/manual_fft.rs
@@ -0,0 +1,39 @@
+use tfhe::core_crypto::fft_impl::fft64::math::fft::{
+    setup_custom_fft_plan, FftAlgo, Method, Plan, PolynomialSize,
+};
+use tfhe::prelude::*;
+use tfhe::{set_server_key, ClientKey, ConfigBuilder, FheUint64, ServerKey};
+
+pub fn main() {
+    let n = PolynomialSize(2048);
+    let fourier_polynomial_size = n.to_fourier_polynomial_size();
+    let my_plan = Plan::new(
+        fourier_polynomial_size.0,
+        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: fourier_polynomial_size.0,
+        },
+    );
+
+    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));
+}
diff --git a/tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs b/tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs
index 2cb864c98..50c8d4bdc 100644
--- a/tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs
+++ b/tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs
@@ -4,7 +4,8 @@ use crate::core_crypto::backward_compatibility::fft_impl::{
 };
 use crate::core_crypto::commons::math::torus::UnsignedTorus;
 use crate::core_crypto::commons::numeric::CastInto;
-use crate::core_crypto::commons::parameters::{PolynomialCount, PolynomialSize};
+pub use crate::core_crypto::commons::parameters::PolynomialSize;
+use crate::core_crypto::commons::parameters::{FourierPolynomialSize, PolynomialCount};
 use crate::core_crypto::commons::traits::{Container, ContainerMut, IntoContainerOwned};
 use crate::core_crypto::commons::utils::izip_eq;
 use crate::core_crypto::entities::*;
@@ -19,7 +20,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 +107,24 @@ fn plans() -> &'static PlanMap {
     PLANS.get_or_init(|| RwLock::new(HashMap::new()))
 }
 
+pub fn setup_custom_fft_plan(plan: Plan) {
+    let base_n = FourierPolynomialSize(plan.fft_size());
+    let n = base_n.to_standard_polynomial_size();
+
+    let plan = Arc::new((Twisties::new(base_n.0), plan));
+
+    let global_plans = plans();
+
+    let mut write = global_plans.write().unwrap();
+
+    match write.entry(n.0) {
+        Entry::Occupied(mut occupied_entry) => occupied_entry.get_mut().set(plan).unwrap(),
+        Entry::Vacant(vacant_entry) => {
+            vacant_entry.insert(Arc::new(OnceLock::from(plan)));
+        }
+    }
+}
+
 /// 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