From 827d257d0a2ad5914e99be27d58e9b5e5f5aaf3e Mon Sep 17 00:00:00 2001 From: Pro7ech Date: Thu, 16 Oct 2025 17:54:12 +0200 Subject: [PATCH] glwe trace --- poulpy-core/src/glwe_trace.rs | 256 ++++++++++++++++++---------------- 1 file changed, 137 insertions(+), 119 deletions(-) diff --git a/poulpy-core/src/glwe_trace.rs b/poulpy-core/src/glwe_trace.rs index dc5582c..48608d1 100644 --- a/poulpy-core/src/glwe_trace.rs +++ b/poulpy-core/src/glwe_trace.rs @@ -1,170 +1,188 @@ use std::collections::HashMap; use poulpy_hal::{ - api::{ - ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, - VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, - VecZnxRshInplace, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, - }, - layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx}, + api::ModuleLogN, + layouts::{Backend, DataMut, GaloisElement, Module, Scratch, VecZnx}, }; -use crate::layouts::{Base2K, GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWEInfos, prepared::AutomorphismKeyPrepared}; +use crate::{ + GLWEAutomorphism, GLWECopy, GLWEShift, ScratchTakeCore, + layouts::{ + Base2K, GGLWEInfos, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, LWEInfos, + prepared::{AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement}, + }, +}; impl GLWE> { - pub fn trace_galois_elements(module: &Module) -> Vec { - let mut gal_els: Vec = Vec::new(); - (0..module.log_n()).for_each(|i| { - if i == 0 { - gal_els.push(-1); - } else { - gal_els.push(module.galois_element(1 << (i - 1))); - } - }); - gal_els + pub fn trace_galois_elements(module: &M) -> Vec + where + M: GLWETrace, + { + module.glwe_trace_galois_elements() } - pub fn trace_tmp_bytes(module: &Module, out_infos: &OUT, in_infos: &IN, key_infos: &KEY) -> usize + pub fn trace_tmp_bytes(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize where - OUT: GLWEInfos, - IN: GLWEInfos, - KEY: GGLWEInfos, - Module: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes, + R: GLWEInfos, + A: GLWEInfos, + K: GGLWEInfos, + M: GLWETrace, { - let trace: usize = Self::automorphism_inplace_tmp_bytes(module, out_infos, key_infos); - if in_infos.base2k() != key_infos.base2k() { + module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos) + } +} + +impl GLWE { + pub fn trace( + &mut self, + module: &M, + start: usize, + end: usize, + a: &A, + keys: &HashMap, + scratch: &mut Scratch, + ) where + A: GLWEToRef, + K: AutomorphismKeyPreparedToRef + GGLWEInfos + GetAutomorphismGaloisElement, + Scratch: ScratchTakeCore, + M: GLWETrace, + { + module.glwe_trace(self, start, end, a, keys, scratch); + } + + pub fn trace_inplace( + &mut self, + module: &M, + start: usize, + end: usize, + keys: &HashMap, + scratch: &mut Scratch, + ) where + K: AutomorphismKeyPreparedToRef + GGLWEInfos + GetAutomorphismGaloisElement, + Scratch: ScratchTakeCore, + M: GLWETrace, + { + module.glwe_trace_inplace(self, start, end, keys, scratch); + } +} + +impl GLWETrace for Module where + Self: ModuleLogN + GaloisElement + GLWEAutomorphism + GLWEShift + GLWECopy +{ +} + +pub trait GLWETrace +where + Self: ModuleLogN + GaloisElement + GLWEAutomorphism + GLWEShift + GLWECopy, +{ + fn glwe_trace_galois_elements(&self) -> Vec { + (0..self.log_n()) + .map(|i| { + if i == 0 { + -1 + } else { + self.galois_element(1 << (i - 1)) + } + }) + .collect() + } + + fn glwe_trace_tmp_bytes(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize + where + R: GLWEInfos, + A: GLWEInfos, + K: GGLWEInfos, + { + let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos); + if a_infos.base2k() != key_infos.base2k() { let glwe_conv: usize = VecZnx::bytes_of( - module.n(), + self.n(), (key_infos.rank_out() + 1).into(), - out_infos.k().min(in_infos.k()).div_ceil(key_infos.base2k()) as usize, - ) + module.vec_znx_normalize_tmp_bytes(); + res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize, + ) + self.vec_znx_normalize_tmp_bytes(); return glwe_conv + trace; } trace } - pub fn trace_inplace_tmp_bytes(module: &Module, out_infos: &OUT, key_infos: &KEY) -> usize + fn glwe_trace(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &HashMap, scratch: &mut Scratch) where - OUT: GLWEInfos, - KEY: GGLWEInfos, - Module: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes, + R: GLWEToMut, + A: GLWEToRef, + K: AutomorphismKeyPreparedToRef + GGLWEInfos + GetAutomorphismGaloisElement, + Scratch: ScratchTakeCore, { - Self::trace_tmp_bytes(module, out_infos, out_infos, key_infos) - } -} - -impl GLWE { - pub fn trace( - &mut self, - module: &Module, - start: usize, - end: usize, - lhs: &GLWE, - auto_keys: &HashMap>, - scratch: &mut Scratch, - ) where - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxBigAutomorphismInplace - + VecZnxRshInplace - + VecZnxCopy - + VecZnxNormalizeTmpBytes - + VecZnxNormalize, - Scratch: ScratchAvailable, - { - self.copy(module, lhs); - self.trace_inplace(module, start, end, auto_keys, scratch); + self.glwe_copy(res, a); + self.glwe_trace_inplace(res, start, end, keys, scratch); } - pub fn trace_inplace( - &mut self, - module: &Module, - start: usize, - end: usize, - auto_keys: &HashMap>, - scratch: &mut Scratch, - ) where - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxBigAutomorphismInplace - + VecZnxRshInplace - + VecZnxNormalizeTmpBytes - + VecZnxNormalize, - Scratch: ScratchAvailable, + fn glwe_trace_inplace(&self, res: &mut R, start: usize, end: usize, keys: &HashMap, scratch: &mut Scratch) + where + R: GLWEToMut, + K: AutomorphismKeyPreparedToRef + GGLWEInfos + GetAutomorphismGaloisElement, + Scratch: ScratchTakeCore, { - let basek_ksk: Base2K = auto_keys - .get(auto_keys.keys().next().unwrap()) - .unwrap() - .base2k(); + let res: &mut GLWE<&mut [u8]> = &mut res.to_mut(); + + let basek_ksk: Base2K = keys.get(keys.keys().next().unwrap()).unwrap().base2k(); #[cfg(debug_assertions)] { - assert_eq!(self.n(), module.n() as u32); + assert_eq!(res.n(), self.n() as u32); assert!(start < end); - assert!(end <= module.log_n()); - for key in auto_keys.values() { - assert_eq!(key.n(), module.n() as u32); + assert!(end <= self.log_n()); + for key in keys.values() { + assert_eq!(key.n(), self.n() as u32); assert_eq!(key.base2k(), basek_ksk); - assert_eq!(key.rank_in(), self.rank()); - assert_eq!(key.rank_out(), self.rank()); + assert_eq!(key.rank_in(), res.rank()); + assert_eq!(key.rank_out(), res.rank()); } } - if self.base2k() != basek_ksk { - let (mut self_conv, scratch_1) = scratch.take_glwe_ct(&GLWELayout { - n: module.n().into(), - base2k: basek_ksk, - k: self.k(), - rank: self.rank(), - }); + if res.base2k() != basek_ksk { + let (mut self_conv, scratch_1) = scratch.take_glwe_ct( + self, + &GLWELayout { + n: self.n().into(), + base2k: basek_ksk, + k: res.k(), + rank: res.rank(), + }, + ); - for j in 0..(self.rank() + 1).into() { - module.vec_znx_normalize( + for j in 0..(res.rank() + 1).into() { + self.vec_znx_normalize( basek_ksk.into(), &mut self_conv.data, j, basek_ksk.into(), - &self.data, + res.data(), j, scratch_1, ); } for i in start..end { - self_conv.rsh(module, 1, scratch_1); + self.glwe_rsh(1, &mut self_conv, scratch_1); let p: i64 = if i == 0 { -1 } else { - module.galois_element(1 << (i - 1)) + self.galois_element(1 << (i - 1)) }; - if let Some(key) = auto_keys.get(&p) { - self_conv.automorphism_add_inplace(module, key, scratch_1); + if let Some(key) = keys.get(&p) { + self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1); } else { - panic!("auto_keys[{p}] is empty") + panic!("keys[{p}] is empty") } } - for j in 0..(self.rank() + 1).into() { - module.vec_znx_normalize( - self.base2k().into(), - &mut self.data, + for j in 0..(res.rank() + 1).into() { + self.vec_znx_normalize( + res.base2k().into(), + res.data_mut(), j, basek_ksk.into(), &self_conv.data, @@ -174,18 +192,18 @@ impl GLWE { } } else { for i in start..end { - self.rsh(module, 1, scratch); + self.glwe_rsh(1, res, scratch); let p: i64 = if i == 0 { -1 } else { - module.galois_element(1 << (i - 1)) + self.galois_element(1 << (i - 1)) }; - if let Some(key) = auto_keys.get(&p) { - self.automorphism_add_inplace(module, key, scratch); + if let Some(key) = keys.get(&p) { + self.glwe_automorphism_add_inplace(res, key, scratch); } else { - panic!("auto_keys[{p}] is empty") + panic!("keys[{p}] is empty") } } }