From 93e0671edad6583e0a5be1322139981fcc3248b1 Mon Sep 17 00:00:00 2001 From: Pro7ech Date: Thu, 16 Oct 2025 22:17:22 +0200 Subject: [PATCH 1/3] glwe packing --- poulpy-core/src/glwe_packing.rs | 246 +++++++++++++++-------------- poulpy-core/src/operations/glwe.rs | 11 +- 2 files changed, 139 insertions(+), 118 deletions(-) diff --git a/poulpy-core/src/glwe_packing.rs b/poulpy-core/src/glwe_packing.rs index 5304df8..4b7d5b1 100644 --- a/poulpy-core/src/glwe_packing.rs +++ b/poulpy-core/src/glwe_packing.rs @@ -1,12 +1,13 @@ use std::collections::HashMap; use poulpy_hal::{ - api::{ModuleLogN, VecZnxCopy, VecZnxRotateInplace}, - layouts::{Backend, DataMut, DataRef, GaloisElement, Module, Scratch}, + api::ModuleLogN, + layouts::{Backend, GaloisElement, Module, Scratch}, }; use crate::{ - GLWEAdd, GLWEAutomorphism, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore, + GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore, + glwe_trace::GLWETrace, layouts::{ GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::{AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement}, @@ -40,10 +41,10 @@ impl Accumulator { /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation. /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus. /// * `rank`: rank of the GLWE ciphertext. - pub fn alloc(module: &M, infos: &A) -> Self + pub fn alloc(module: &M, infos: &A) -> Self where A: GLWEInfos, - M: GLWEAlloc, + M: GLWEPacking, { Self { data: GLWE::alloc_from_infos(module, infos), @@ -64,15 +65,15 @@ impl GLWEPacker { /// and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients /// which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts /// can be packed. - pub fn new(module: &M, infos: &A, log_batch: usize) -> Self + pub fn alloc(module: &M, infos: &A, log_batch: usize) -> Self where A: GLWEInfos, - M: GLWEAlloc, + M: GLWEPacking, { let mut accumulators: Vec = Vec::::new(); let log_n: usize = infos.n().log2(); (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(module, infos))); - Self { + GLWEPacker { accumulators, log_batch, counter: 0, @@ -93,13 +94,19 @@ impl GLWEPacker { where R: GLWEInfos, K: GGLWEInfos, - M: GLWEAlloc + GLWEAutomorphism, + M: GLWEPacking, { - pack_core_tmp_bytes(module, res_infos, key_infos) + module.bytes_of_glwe_from_infos(res_infos) + + module + .glwe_rsh_tmp_byte() + .max(module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos)) } - pub fn galois_elements(module: &Module) -> Vec { - GLWE::trace_galois_elements(module) + pub fn galois_elements(module: &M) -> Vec + where + M: GLWETrace, + { + module.glwe_trace_galois_elements() } /// Adds a GLWE ciphertext to the [GLWEPacker]. @@ -111,11 +118,11 @@ impl GLWEPacker { /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext. /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s. /// * `scratch`: scratch space of size at least [Self::tmp_bytes]. - pub fn add(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap, scratch: &mut Scratch) + pub fn add(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap, scratch: &mut Scratch) where - A: GLWEToRef, - K: AutomorphismKeyPreparedToRef, - M: GLWEAutomorphism, + A: GLWEToRef + GLWEInfos, + K: AutomorphismKeyPreparedToRef + GetAutomorphismGaloisElement, + M: GLWEPacking, Scratch: ScratchTakeCore, { assert!( @@ -136,14 +143,15 @@ impl GLWEPacker { } /// Flush result to`res`. - pub fn flush(&mut self, module: &Module, res: &mut GLWE) + pub fn flush(&mut self, module: &M, res: &mut R) where - Module: VecZnxCopy, + R: GLWEToMut, + M: GLWEPacking, { assert!(self.counter as u32 == self.accumulators[0].data.n()); // Copy result GLWE into res GLWE - res.copy( - module, + module.glwe_copy( + res, &self.accumulators[module.log_n() - self.log_batch - 1].data, ); @@ -151,13 +159,76 @@ impl GLWEPacker { } } -fn pack_core_tmp_bytes(module: &M, res_infos: &R, key_infos: &K) -> usize -where - R: GLWEInfos, - K: GGLWEInfos, - M: GLWEAlloc + GLWEAutomorphism, +impl GLWEPacking for Module where + Self: GLWEAutomorphism + + GaloisElement + + ModuleLogN + + GLWERotate + + GLWESub + + GLWEShift + + GLWEAdd + + GLWENormalize + + GLWECopy + + GLWEAlloc { - combine_tmp_bytes(module, res_infos, key_infos) +} + +pub trait GLWEPacking +where + Self: GLWEAutomorphism + + GaloisElement + + ModuleLogN + + GLWERotate + + GLWESub + + GLWEShift + + GLWEAdd + + GLWENormalize + + GLWECopy + + GLWEAlloc, +{ + /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)] + /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)] + fn glwe_pack( + &self, + cts: &mut HashMap, + log_gap_out: usize, + keys: &HashMap, + scratch: &mut Scratch, + ) where + R: GLWEToMut + GLWEToRef + GLWEInfos, + K: AutomorphismKeyPreparedToRef + GetAutomorphismGaloisElement, + Scratch: ScratchTakeCore, + { + #[cfg(debug_assertions)] + { + assert!(*cts.keys().max().unwrap() < self.n()) + } + + let log_n: usize = self.log_n(); + + for i in 0..(log_n - log_gap_out) { + let t: usize = (1 << log_n).min(1 << (log_n - 1 - i)); + + let key: &K = if i == 0 { + keys.get(&-1).unwrap() + } else { + keys.get(&self.galois_element(1 << (i - 1))).unwrap() + }; + + for j in 0..t { + let mut a: Option<&mut R> = cts.remove(&j); + let mut b: Option<&mut R> = cts.remove(&(j + t)); + + pack_internal(self, &mut a, &mut b, i, key, scratch); + + if let Some(a) = a { + cts.insert(j, a); + } else if let Some(b) = b { + cts.insert(j, b); + } + } + } + } } fn pack_core( @@ -169,8 +240,16 @@ fn pack_core( scratch: &mut Scratch, ) where A: GLWEToRef + GLWEInfos, - K: AutomorphismKeyPreparedToRef, - M: GLWEAutomorphism + ModuleLogN + VecZnxCopy, + K: AutomorphismKeyPreparedToRef + GetAutomorphismGaloisElement, + M: ModuleLogN + + GLWEAutomorphism + + GaloisElement + + GLWERotate + + GLWESub + + GLWEShift + + GLWEAdd + + GLWENormalize + + GLWECopy, Scratch: ScratchTakeCore, { let log_n: usize = module.log_n(); @@ -188,7 +267,7 @@ fn pack_core( // No previous value -> copies and sets flags accordingly if let Some(a_ref) = a { - acc_mut_ref.data.copy(module, a_ref); + module.glwe_copy(&mut acc_mut_ref.data, a_ref); acc_mut_ref.value = true } else { acc_mut_ref.value = false @@ -222,16 +301,6 @@ fn pack_core( } } -fn combine_tmp_bytes(module: &M, res_infos: &R, key_infos: &K) -> usize -where - R: GLWEInfos, - K: GGLWEInfos, - M: GLWEAlloc + GLWEAutomorphism, -{ - GLWE::bytes_of_from_infos(module, res_infos) - + (GLWE::rsh_tmp_bytes(module.n()) | module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos)) -} - /// [combine] merges two ciphertexts together. fn combine( module: &M, @@ -242,8 +311,9 @@ fn combine( scratch: &mut Scratch, ) where B: GLWEToRef + GLWEInfos, - K: AutomorphismKeyPreparedToRef, - M: GLWEAutomorphism + GaloisElement + VecZnxRotateInplace, + M: GLWEAutomorphism + GaloisElement + GLWERotate + GLWESub + GLWEShift + GLWEAdd + GLWENormalize, + B: GLWEToRef + GLWEInfos, + K: AutomorphismKeyPreparedToRef + GetAutomorphismGaloisElement, Scratch: ScratchTakeCore, { let log_n: usize = acc.data.n().log2(); @@ -272,50 +342,50 @@ fn combine( let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a); // a = a * X^-t - a.rotate_inplace(module, -t, scratch_1); + module.glwe_rotate_inplace(-t, a, scratch_1); // tmp_b = a * X^-t - b - tmp_b.sub(module, a, b); - tmp_b.rsh(module, 1, scratch_1); + module.glwe_sub(&mut tmp_b, a, b); + module.glwe_rsh(1, &mut tmp_b, scratch_1); // a = a * X^-t + b - a.add_inplace(module, b); - a.rsh(module, 1, scratch_1); + module.glwe_add_inplace(a, b); + module.glwe_rsh(1, a, scratch_1); - tmp_b.normalize_inplace(module, scratch_1); + module.glwe_normalize_inplace(&mut tmp_b, scratch_1); // tmp_b = phi(a * X^-t - b) - if let Some(key) = auto_keys.get(&gal_el) { - tmp_b.automorphism_inplace(module, key, scratch_1); + if let Some(auto_key) = auto_keys.get(&gal_el) { + module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1); } else { panic!("auto_key[{gal_el}] not found"); } // a = a * X^-t + b - phi(a * X^-t - b) - a.sub_inplace_ab(module, &tmp_b); - a.normalize_inplace(module, scratch_1); + module.glwe_sub_inplace(a, &tmp_b); + module.glwe_normalize_inplace(a, scratch_1); // a = a + b * X^t - phi(a * X^-t - b) * X^t // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t) // = a + b * X^t + phi(a - b * X^t) - a.rotate_inplace(module, t, scratch_1); + module.glwe_rotate_inplace(t, a, scratch_1); } else { - a.rsh(module, 1, scratch); + module.glwe_rsh(1, a, scratch); // a = a + phi(a) - if let Some(key) = auto_keys.get(&gal_el) { - a.automorphism_add_inplace(module, key, scratch); + if let Some(auto_key) = auto_keys.get(&gal_el) { + module.glwe_automorphism_add_inplace(a, auto_key, scratch); } else { panic!("auto_key[{gal_el}] not found"); } } } else if let Some(b) = b { - let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a); - tmp_b.rotate(module, 1 << (log_n - i - 1), b); - tmp_b.rsh(module, 1, scratch_1); + let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a); + module.glwe_rotate(t, &mut tmp_b, b); + module.glwe_rsh(1, &mut tmp_b, scratch_1); // a = (b* X^t - phi(b* X^t)) - if let Some(key) = auto_keys.get(&gal_el) { - a.automorphism_sub_negate(module, &tmp_b, key, scratch_1); + if let Some(auto_key) = auto_keys.get(&gal_el) { + module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1); } else { panic!("auto_key[{gal_el}] not found"); } @@ -324,62 +394,6 @@ fn combine( } } -pub trait GLWEPacking -where - Self: GLWEAutomorphism - + GaloisElement - + ModuleLogN - + GLWERotate - + GLWESub - + GLWEShift - + GLWEAdd - + GLWENormalize, -{ - /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)] - /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)] - fn glwe_pack( - &self, - cts: &mut HashMap, - log_gap_out: usize, - keys: &HashMap, - scratch: &mut Scratch, - ) where - R: GLWEToMut + GLWEToRef + GLWEInfos, - K: AutomorphismKeyPreparedToRef + GetAutomorphismGaloisElement, - Scratch: ScratchTakeCore, - { - #[cfg(debug_assertions)] - { - assert!(*cts.keys().max().unwrap() < self.n()) - } - - let log_n: usize = self.log_n(); - - for i in 0..(log_n - log_gap_out){ - let t: usize = (1 << log_n).min(1 << (log_n - 1 - i)); - - let key: &K = if i == 0 { - keys.get(&-1).unwrap() - } else { - keys.get(&self.galois_element(1 << (i - 1))).unwrap() - }; - - for j in 0..t{ - let mut a: Option<&mut R> = cts.remove(&j); - let mut b: Option<&mut R> = cts.remove(&(j + t)); - - pack_internal(self, &mut a, &mut b, i, key, scratch); - - if let Some(a) = a { - cts.insert(j, a); - } else if let Some(b) = b { - cts.insert(j, b); - } - }; - }; - } -} - #[allow(clippy::too_many_arguments)] fn pack_internal( module: &M, diff --git a/poulpy-core/src/operations/glwe.rs b/poulpy-core/src/operations/glwe.rs index 3e507e2..27021e7 100644 --- a/poulpy-core/src/operations/glwe.rs +++ b/poulpy-core/src/operations/glwe.rs @@ -277,6 +277,10 @@ pub trait GLWEShift where Self: ModuleN + VecZnxRshInplace, { + fn glwe_rsh_tmp_byte(&self) -> usize { + VecZnx::rsh_tmp_bytes(self.n()) + } + fn glwe_rsh(&self, k: usize, res: &mut R, scratch: &mut Scratch) where R: GLWEToMut, @@ -291,8 +295,11 @@ where } impl GLWE> { - pub fn rsh_tmp_bytes(n: usize) -> usize { - VecZnx::rsh_tmp_bytes(n) + pub fn rsh_tmp_bytes(module: &M) -> usize + where + M: GLWEShift, + { + module.glwe_rsh_tmp_byte() } } From 920b9252162ee5704592b75b252d1dfd8424457f Mon Sep 17 00:00:00 2001 From: Pro7ech Date: Thu, 16 Oct 2025 22:40:23 +0200 Subject: [PATCH 2/3] conversions --- poulpy-core/src/conversion/glwe_to_lwe.rs | 183 ++++++++++++---------- poulpy-core/src/conversion/lwe_to_glwe.rs | 134 ++++++++-------- 2 files changed, 177 insertions(+), 140 deletions(-) diff --git a/poulpy-core/src/conversion/glwe_to_lwe.rs b/poulpy-core/src/conversion/glwe_to_lwe.rs index b6c6ed1..49744ba 100644 --- a/poulpy-core/src/conversion/glwe_to_lwe.rs +++ b/poulpy-core/src/conversion/glwe_to_lwe.rs @@ -1,100 +1,125 @@ use poulpy_hal::{ - api::{ - VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf, - VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, - VmpApplyDftToDftTmpBytes, - }, - layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero}, + api::ModuleN, + layouts::{Backend, DataMut, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero}, }; -use crate::layouts::{GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWE, LWEInfos, Rank, prepared::GLWEToLWESwitchingKeyPrepared}; +use crate::{ + GLWEKeyswitch, ScratchTakeCore, + layouts::{ + GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWELayout, GLWEToRef, LWE, LWEInfos, LWEToMut, Rank, + prepared::{LWEToGLWESwitchingKeyPrepared, LWEToGLWESwitchingKeyPreparedToRef}, + }, +}; -impl LWE> { - pub fn from_glwe_tmp_bytes( - module: &Module, - lwe_infos: &OUT, - glwe_infos: &IN, - key_infos: &KEY, - ) -> usize +pub trait LWESampleExtract +where + Self: ModuleN, +{ + fn lwe_sample_extract(&self, res: &mut R, a: &A) where - OUT: LWEInfos, - IN: GLWEInfos, - KEY: GGLWEInfos, - Module: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes, + R: LWEToMut, + A: GLWEToRef, { - let glwe_layout: GLWELayout = GLWELayout { - n: module.n().into(), + let res: &mut LWE<&mut [u8]> = &mut res.to_mut(); + let a: &GLWE<&[u8]> = &a.to_ref(); + + assert!(res.n() <= a.n()); + assert_eq!(a.n(), self.n() as u32); + assert!(res.base2k() == a.base2k()); + + let min_size: usize = res.size().min(a.size()); + let n: usize = res.n().into(); + + res.data.zero(); + (0..min_size).for_each(|i| { + let data_lwe: &mut [i64] = res.data.at_mut(0, i); + data_lwe[0] = a.data.at(0, i)[0]; + data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]); + }); + } +} + +impl LWESampleExtract for Module where Self: ModuleN {} +impl LWEFromGLWE for Module where Self: GLWEKeyswitch + GLWEAlloc + LWESampleExtract {} + +pub trait LWEFromGLWE +where + Self: GLWEKeyswitch + GLWEAlloc + LWESampleExtract, +{ + fn lwe_from_glwe_tmp_bytes(&self, lwe_infos: &R, glwe_infos: &A, key_infos: &K) -> usize + where + R: LWEInfos, + A: GLWEInfos, + K: GGLWEInfos, + { + let res_infos: GLWELayout = GLWELayout { + n: self.n().into(), base2k: lwe_infos.base2k(), k: lwe_infos.k(), rank: Rank(1), }; - GLWE::bytes_of( - module.n().into(), - lwe_infos.base2k(), - lwe_infos.k(), - 1u32.into(), - ) + GLWE::keyswitch_tmp_bytes(module, &glwe_layout, glwe_infos, key_infos) - } -} - -impl LWE { - pub fn sample_extract(&mut self, a: &GLWE) { - #[cfg(debug_assertions)] - { - assert!(self.n() <= a.n()); - assert!(self.base2k() == a.base2k()); - } - - let min_size: usize = self.size().min(a.size()); - let n: usize = self.n().into(); - - self.data.zero(); - (0..min_size).for_each(|i| { - let data_lwe: &mut [i64] = self.data.at_mut(0, i); - data_lwe[0] = a.data.at(0, i)[0]; - data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]); - }); + self.bytes_of_glwe(lwe_infos.base2k(), lwe_infos.k(), 1u32.into()) + + self.glwe_keyswitch_tmp_bytes(&res_infos, glwe_infos, key_infos) } - pub fn from_glwe( - &mut self, - module: &Module, - a: &GLWE, - ks: &GLWEToLWESwitchingKeyPrepared, - scratch: &mut Scratch, - ) where - DGlwe: DataRef, - DKs: DataRef, - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxNormalize - + VecZnxNormalizeTmpBytes, - Scratch:, + fn lwe_from_glwe(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch) + where + R: LWEToMut, + A: GLWEToRef, + K: LWEToGLWESwitchingKeyPreparedToRef + GGLWEInfos, + Scratch: ScratchTakeCore, { - #[cfg(debug_assertions)] - { - assert_eq!(a.n(), module.n() as u32); - assert_eq!(ks.n(), module.n() as u32); - assert!(self.n() <= module.n() as u32); - } + let res: &mut LWE<&mut [u8]> = &mut res.to_mut(); + let a: &GLWE<&[u8]> = &a.to_ref(); + let key: &LWEToGLWESwitchingKeyPrepared<&[u8], BE> = &key.to_ref(); + + assert_eq!(a.n(), self.n() as u32); + assert_eq!(key.n(), self.n() as u32); + assert!(res.n() <= self.n() as u32); let glwe_layout: GLWELayout = GLWELayout { - n: module.n().into(), - base2k: self.base2k(), - k: self.k(), + n: self.n().into(), + base2k: res.base2k(), + k: res.k(), rank: Rank(1), }; - let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(&glwe_layout); - tmp_glwe.keyswitch(module, a, &ks.0, scratch_1); - self.sample_extract(&tmp_glwe); + let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(self, &glwe_layout); + self.glwe_keyswitch(&mut tmp_glwe, a, &key.0, scratch_1); + self.lwe_sample_extract(res, &tmp_glwe); + } +} + +impl LWE> { + pub fn from_glwe_tmp_bytes(module: &M, lwe_infos: &R, glwe_infos: &A, key_infos: &K) -> usize + where + R: LWEInfos, + A: GLWEInfos, + K: GGLWEInfos, + M: LWEFromGLWE, + { + module.lwe_from_glwe_tmp_bytes(lwe_infos, glwe_infos, key_infos) + } +} + +impl LWE { + pub fn sample_extract(&mut self, module: &M, a: &A) + where + A: GLWEToRef, + M: LWESampleExtract, + { + module.lwe_sample_extract(self, a); + } + + pub fn from_glwe(&self, module: &M, res: &mut R, a: &A, key: &K, scratch: &mut Scratch) + where + R: LWEToMut, + A: GLWEToRef, + K: LWEToGLWESwitchingKeyPreparedToRef + GGLWEInfos, + M: LWEFromGLWE, + Scratch: ScratchTakeCore, + { + module.lwe_from_glwe(res, a, key, scratch); } } diff --git a/poulpy-core/src/conversion/lwe_to_glwe.rs b/poulpy-core/src/conversion/lwe_to_glwe.rs index c4a3b88..7be6574 100644 --- a/poulpy-core/src/conversion/lwe_to_glwe.rs +++ b/poulpy-core/src/conversion/lwe_to_glwe.rs @@ -1,79 +1,67 @@ use poulpy_hal::{ - api::{ - ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, - VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, - VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, - }, - layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, ZnxView, ZnxViewMut, ZnxZero}, + api::ScratchTakeBasic, + layouts::{Backend, DataMut, Module, Scratch, VecZnx, ZnxView, ZnxViewMut, ZnxZero}, }; -use crate::layouts::{GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWE, LWEInfos, prepared::LWEToGLWESwitchingKeyPrepared}; +use crate::{ + GLWEKeyswitch, ScratchTakeCore, + layouts::{ + GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWELayout, GLWEToMut, LWE, LWEInfos, LWEToRef, + prepared::{LWEToGLWESwitchingKeyPrepared, LWEToGLWESwitchingKeyPreparedToRef}, + }, +}; -impl GLWE> { - pub fn from_lwe_tmp_bytes( - module: &Module, - glwe_infos: &OUT, - lwe_infos: &IN, - key_infos: &KEY, - ) -> usize +impl GLWEFromLWE for Module where Self: GLWEKeyswitch + GLWEAlloc {} + +pub trait GLWEFromLWE +where + Self: GLWEKeyswitch + GLWEAlloc, +{ + fn glwe_from_lwe_tmp_bytes(&self, glwe_infos: &R, lwe_infos: &A, key_infos: &K) -> usize where - OUT: GLWEInfos, - IN: LWEInfos, - KEY: GGLWEInfos, - Module: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes, + R: GLWEInfos, + A: LWEInfos, + K: GGLWEInfos, { - let ct: usize = GLWE::bytes_of( - module.n().into(), + let ct: usize = self.bytes_of_glwe( key_infos.base2k(), lwe_infos.k().max(glwe_infos.k()), 1u32.into(), ); - let ks: usize = GLWE::keyswitch_inplace_tmp_bytes(module, glwe_infos, key_infos); + + let ks: usize = self.glwe_keyswitch_tmp_bytes(glwe_infos, glwe_infos, key_infos); if lwe_infos.base2k() == key_infos.base2k() { ct + ks } else { - let a_conv = VecZnx::bytes_of(module.n(), 1, lwe_infos.size()) + module.vec_znx_normalize_tmp_bytes(); + let a_conv = VecZnx::bytes_of(self.n(), 1, lwe_infos.size()) + self.vec_znx_normalize_tmp_bytes(); ct + a_conv + ks } } -} -impl GLWE { - pub fn from_lwe( - &mut self, - module: &Module, - lwe: &LWE, - ksk: &LWEToGLWESwitchingKeyPrepared, - scratch: &mut Scratch, - ) where - DLwe: DataRef, - DKsk: DataRef, - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxNormalize - + VecZnxNormalizeTmpBytes, - Scratch: ScratchAvailable, + fn glwe_from_lwe(&self, res: &mut R, lwe: &A, ksk: &K, scratch: &mut Scratch) + where + R: GLWEToMut, + A: LWEToRef, + K: LWEToGLWESwitchingKeyPreparedToRef, + Scratch: ScratchTakeCore, { - #[cfg(debug_assertions)] - { - assert_eq!(self.n(), module.n() as u32); - assert_eq!(ksk.n(), module.n() as u32); - assert!(lwe.n() <= module.n() as u32); - } + let res: &mut GLWE<&mut [u8]> = &mut res.to_mut(); + let lwe: &LWE<&[u8]> = &lwe.to_ref(); + let ksk: &LWEToGLWESwitchingKeyPrepared<&[u8], BE> = &ksk.to_ref(); - let (mut glwe, scratch_1) = scratch.take_glwe_ct(&GLWELayout { - n: ksk.n(), - base2k: ksk.base2k(), - k: lwe.k(), - rank: 1u32.into(), - }); + assert_eq!(res.n(), self.n() as u32); + assert_eq!(ksk.n(), self.n() as u32); + assert!(lwe.n() <= self.n() as u32); + + let (mut glwe, scratch_1) = scratch.take_glwe_ct( + self, + &GLWELayout { + n: ksk.n(), + base2k: ksk.base2k(), + k: lwe.k(), + rank: 1u32.into(), + }, + ); glwe.data.zero(); let n_lwe: usize = lwe.n().into(); @@ -85,14 +73,14 @@ impl GLWE { glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]); } } else { - let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(module.n(), 1, lwe.size()); + let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self, 1, lwe.size()); a_conv.zero(); for j in 0..lwe.size() { let data_lwe: &[i64] = lwe.data.at(0, j); a_conv.at_mut(0, j)[0] = data_lwe[0] } - module.vec_znx_normalize( + self.vec_znx_normalize( ksk.base2k().into(), &mut glwe.data, 0, @@ -108,7 +96,7 @@ impl GLWE { a_conv.at_mut(0, j)[..n_lwe].copy_from_slice(&data_lwe[1..]); } - module.vec_znx_normalize( + self.vec_znx_normalize( ksk.base2k().into(), &mut glwe.data, 1, @@ -119,6 +107,30 @@ impl GLWE { ); } - self.keyswitch(module, &glwe, &ksk.0, scratch_1); + self.glwe_keyswitch(res, &glwe, &ksk.0, scratch_1); + } +} + +impl GLWE> { + pub fn from_lwe_tmp_bytes(module: &M, glwe_infos: &R, lwe_infos: &A, key_infos: &K) -> usize + where + R: GLWEInfos, + A: LWEInfos, + K: GGLWEInfos, + M: GLWEFromLWE, + { + module.glwe_from_lwe_tmp_bytes(glwe_infos, lwe_infos, key_infos) + } +} + +impl GLWE { + pub fn from_lwe(&mut self, module: &M, lwe: &A, ksk: &K, scratch: &mut Scratch) + where + M: GLWEFromLWE, + A: LWEToRef, + K: LWEToGLWESwitchingKeyPreparedToRef, + Scratch: ScratchTakeCore, + { + module.glwe_from_lwe(self, lwe, ksk, scratch); } } From 90a25638c248750ece4eb1182fc4f359a085cfd5 Mon Sep 17 00:00:00 2001 From: Pro7ech Date: Thu, 16 Oct 2025 22:41:40 +0200 Subject: [PATCH 3/3] fix keyswitching --- poulpy-core/src/conversion/mod.rs | 2 ++ poulpy-core/src/keyswitching/lwe_ct.rs | 6 +++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/poulpy-core/src/conversion/mod.rs b/poulpy-core/src/conversion/mod.rs index 9771531..8be980a 100644 --- a/poulpy-core/src/conversion/mod.rs +++ b/poulpy-core/src/conversion/mod.rs @@ -3,3 +3,5 @@ mod glwe_to_lwe; mod lwe_to_glwe; pub use gglwe_to_ggsw::*; +pub use glwe_to_lwe::*; +pub use lwe_to_glwe::*; diff --git a/poulpy-core/src/keyswitching/lwe_ct.rs b/poulpy-core/src/keyswitching/lwe_ct.rs index ff9fb5f..7aa5321 100644 --- a/poulpy-core/src/keyswitching/lwe_ct.rs +++ b/poulpy-core/src/keyswitching/lwe_ct.rs @@ -4,7 +4,7 @@ use poulpy_hal::{ }; use crate::{ - ScratchTakeCore, + LWESampleExtract, ScratchTakeCore, keyswitching::glwe_ct::GLWEKeyswitch, layouts::{ GGLWEInfos, GLWE, GLWEAlloc, GLWELayout, LWE, LWEInfos, LWEToMut, LWEToRef, Rank, TorusPrecision, @@ -40,7 +40,7 @@ impl LWEKeySwitch for Module where Self: LWEKeySwitch { pub trait LWEKeySwitch where - Self: GLWEKeyswitch + GLWEAlloc, + Self: GLWEKeyswitch + GLWEAlloc + LWESampleExtract, { fn lwe_keyswitch_tmp_bytes(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize where @@ -121,6 +121,6 @@ where } self.glwe_keyswitch(&mut glwe_out, &glwe_in, &ksk.0, scratch_1); - res.sample_extract(&glwe_out); + self.lwe_sample_extract(res, &glwe_out); } }