github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-09 14:47:56 -05:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): use typed PolynomialSize for Plan

Browse Source
This commit is contained in:
Mayeul@Zama
2025-12-11 16:54:50 +01:00
parent aa6284314b
commit 59fc25f39a
4 changed files with 32 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tfhe/src/core_crypto/commons/math/ntt/ntt64.rs
View File

@@ -24,7 +24,7 @@ impl Ntt64 {
}
// Key is (polynomial size, modulus).
type PlanMap = crate::core_crypto::commons::plan::PlanMap<(usize, u64), Plan>;
type PlanMap = crate::core_crypto::commons::plan::PlanMap<(PolynomialSize, u64), Plan>;
pub(crate) static PLANS: OnceLock<PlanMap> = OnceLock::new();
@@ -34,21 +34,25 @@ fn plans() -> &'static PlanMap {
impl Ntt64 {
/// Real polynomial of size `size`.
pub fn new(modulus: CiphertextModulus<u64>, size: PolynomialSize) -> Self {
pub fn new(modulus: CiphertextModulus<u64>, polynomial_size: PolynomialSize) -> Self {
let global_plans = plans();
assert_eq!(modulus.kind(), CiphertextModulusKind::Other);
let n = size.0;
let modulus = modulus.get_custom_modulus() as u64;
let plan = new_from_plan_map(global_plans, (n, modulus), |(n, modulus)| {
Plan::try_new(n, modulus).unwrap_or_else(|| {
panic!(
"could not generate an NTT plan for the given (size, modulus) ({n}, {modulus})"
let plan = new_from_plan_map(
global_plans,
(polynomial_size, modulus),
|(polynomial_size, modulus)| {
Plan::try_new(polynomial_size.0, modulus).unwrap_or_else(|| {
panic!(
"could not generate an NTT plan for the given (size, modulus) ({}, {modulus})",
polynomial_size.0
)
})
});
})
},
);
Self { plan }
}

2
tfhe/src/core_crypto/commons/parameters.rs
View File

@@ -135,7 +135,7 @@ impl GlweDimension {
///
/// Assuming a polynomial $a\_0 + a\_1X + /dots + a\_{N-1}X^{N-1}$, this new-type contains $N$.
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Versionize,
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Versionize, Hash,
)]
#[versionize(PolynomialSizeVersions)]
pub struct PolynomialSize(pub usize);

10
tfhe/src/core_crypto/fft_impl/fft128/math/fft/mod.rs
View File

@@ -46,7 +46,7 @@ impl Fft128 {
}
}
type PlanMap = crate::core_crypto::commons::plan::PlanMap<usize, PlanWrapper>;
type PlanMap = crate::core_crypto::commons::plan::PlanMap<PolynomialSize, PlanWrapper>;
pub(crate) static PLANS: OnceLock<PlanMap> = OnceLock::new();
fn plans() -> &'static PlanMap {
@@ -55,12 +55,12 @@ fn plans() -> &'static PlanMap {
impl Fft128 {
/// Real polynomial of size `size`.
pub fn new(size: PolynomialSize) -> Self {
pub fn new(polynomial_size: PolynomialSize) -> Self {
let global_plans = plans();
let n = size.0;
let plan = new_from_plan_map(global_plans, n, |n| PlanWrapper(Plan::new(n / 2)));
let plan = new_from_plan_map(global_plans, polynomial_size, |polynomial_size| {
PlanWrapper(Plan::new(polynomial_size.to_fourier_polynomial_size().0))
});
Self { plan }
}

23
tfhe/src/core_crypto/fft_impl/fft64/math/fft/mod.rs
View File

@@ -99,7 +99,7 @@ impl Fft {
}
}
type PlanMap = crate::core_crypto::commons::plan::PlanMap<usize, (Twisties, Plan)>;
type PlanMap = crate::core_crypto::commons::plan::PlanMap<PolynomialSize, (Twisties, Plan)>;
pub(crate) static PLANS: OnceLock<PlanMap> = OnceLock::new();
fn plans() -> &'static PlanMap {
@@ -148,35 +148,38 @@ pub(crate) fn id<From: 'static, To: 'static>(slice: &[From]) -> &[To] {
impl Fft {
/// Real polynomial of size `size`.
pub fn new(size: PolynomialSize) -> Self {
pub fn new(polynomial_size: PolynomialSize) -> Self {
let global_plans = plans();
let n = size.0;
let new = |polynomial_size: PolynomialSize| {
let fourier_polynomial_size = polynomial_size.to_fourier_polynomial_size();
let new = |n| {
#[cfg(not(feature = "experimental-force_fft_algo_dif4"))]
{
(
Twisties::new(n / 2),
Plan::new(n / 2, Method::Measure(Duration::from_millis(10))),
Twisties::new(fourier_polynomial_size.0),
Plan::new(
fourier_polynomial_size.0,
Method::Measure(Duration::from_millis(10)),
),
)
}
#[cfg(feature = "experimental-force_fft_algo_dif4")]
{
(
Twisties::new(n / 2),
Twisties::new(fourier_polynomial_size.0),
Plan::new(
n / 2,
fourier_polynomial_size.0,
Method::UserProvided {
base_algo: tfhe_fft::ordered::FftAlgo::Dif4,
base_n: n / 2,
base_n: fourier_polynomial_size.0,
},
),
)
}
};
let plan = new_from_plan_map(global_plans, n, new);
let plan = new_from_plan_map(global_plans, polynomial_size, new);
Self { plan }
}