diff --git a/poulpy-core/src/automorphism/ggsw_ct.rs b/poulpy-core/src/automorphism/ggsw_ct.rs index a3cef86..8035063 100644 --- a/poulpy-core/src/automorphism/ggsw_ct.rs +++ b/poulpy-core/src/automorphism/ggsw_ct.rs @@ -1,165 +1,127 @@ use poulpy_hal::{ - api::{ - ScratchAvailable, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize, - VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume, - VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, - VmpApplyDftToDftTmpBytes, - }, - layouts::{Backend, DataMut, DataRef, Module, Scratch}, + api::ScratchAvailable, + layouts::{Backend, DataMut, Module, Scratch}, }; -use crate::layouts::{ - GGLWEInfos, GGSW, GGSWInfos, GLWE, - prepared::{AutomorphismKeyPrepared, TensorKeyPrepared}, +use crate::{ + GGSWExpandRows, ScratchTakeCore, + automorphism::glwe_ct::GLWEAutomorphism, + layouts::{ + GGLWEInfos, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos, + prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, TensorKeyPrepared, TensorKeyPreparedToRef}, + }, }; impl GGSW> { - pub fn automorphism_tmp_bytes( - module: &Module, - out_infos: &OUT, - in_infos: &IN, - key_infos: &KEY, - tsk_infos: &TSK, + pub fn automorphism_tmp_bytes( + module: &M, + res_infos: &R, + a_infos: &A, + key_infos: &K, + tsk_infos: &T, ) -> usize where - OUT: GGSWInfos, - IN: GGSWInfos, - KEY: GGLWEInfos, - TSK: GGLWEInfos, - Module: - VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes, + R: GGSWInfos, + A: GGSWInfos, + K: GGLWEInfos, + T: GGLWEInfos, + M: GGSWAutomorphism, { - let out_size: usize = out_infos.size(); - let ci_dft: usize = module.bytes_of_vec_znx_dft((key_infos.rank_out() + 1).into(), out_size); - let ks_internal: usize = GLWE::keyswitch_tmp_bytes( - module, - &out_infos.glwe_layout(), - &in_infos.glwe_layout(), - key_infos, - ); - let expand: usize = GGSW::expand_row_tmp_bytes(module, out_infos, tsk_infos); - ci_dft + (ks_internal | expand) - } - - pub fn automorphism_inplace_tmp_bytes( - module: &Module, - out_infos: &OUT, - key_infos: &KEY, - tsk_infos: &TSK, - ) -> usize - where - OUT: GGSWInfos, - KEY: GGLWEInfos, - TSK: GGLWEInfos, - Module: - VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes, - { - GGSW::automorphism_tmp_bytes(module, out_infos, out_infos, key_infos, tsk_infos) + module.ggsw_automorphism_tmp_bytes(res_infos, a_infos, key_infos, tsk_infos) } } -impl GGSW { - pub fn automorphism( - &mut self, - module: &Module, - lhs: &GGSW, - auto_key: &AutomorphismKeyPrepared, - tensor_key: &TensorKeyPrepared, - scratch: &mut Scratch, - ) where - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxAutomorphismInplace - + VecZnxBigBytesOf - + VecZnxNormalizeTmpBytes - + VecZnxDftCopy - + VecZnxDftAddInplace - + VecZnxIdftApplyTmpA - + VecZnxNormalize, - Scratch: ScratchAvailable, +impl GGSW { + pub fn automorphism(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch) + where + A: GGSWToRef, + K: AutomorphismKeyPreparedToRef, + T: TensorKeyPreparedToRef, + Scratch: ScratchTakeCore, + M: GGSWAutomorphism, { - #[cfg(debug_assertions)] - { - use crate::layouts::{GLWEInfos, LWEInfos}; - - assert_eq!(self.n(), module.n() as u32); - assert_eq!(lhs.n(), module.n() as u32); - assert_eq!(auto_key.n(), module.n() as u32); - assert_eq!(tensor_key.n(), module.n() as u32); - - assert_eq!( - self.rank(), - lhs.rank(), - "ggsw_out rank: {} != ggsw_in rank: {}", - self.rank(), - lhs.rank() - ); - assert_eq!( - self.rank(), - auto_key.rank_out(), - "ggsw_in rank: {} != auto_key rank: {}", - self.rank(), - auto_key.rank_out() - ); - assert_eq!( - self.rank(), - tensor_key.rank_out(), - "ggsw_in rank: {} != tensor_key rank: {}", - self.rank(), - tensor_key.rank_out() - ); - assert!(scratch.available() >= GGSW::automorphism_tmp_bytes(module, self, lhs, auto_key, tensor_key)) - }; - - // Keyswitch the j-th row of the col 0 - (0..lhs.dnum().into()).for_each(|row_i| { - // Key-switch column 0, i.e. - // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2) - self.at_mut(row_i, 0) - .automorphism(module, &lhs.at(row_i, 0), auto_key, scratch); - }); - self.expand_row(module, tensor_key, scratch); + module.ggsw_automorphism(self, a, key, tsk, scratch); } - pub fn automorphism_inplace( - &mut self, - module: &Module, - auto_key: &AutomorphismKeyPrepared, - tensor_key: &TensorKeyPrepared, - scratch: &mut Scratch, - ) where - Module: VecZnxDftBytesOf - + VmpApplyDftToDftTmpBytes - + VecZnxBigNormalizeTmpBytes - + VmpApplyDftToDft - + VmpApplyDftToDftAdd - + VecZnxDftApply - + VecZnxIdftApplyConsume - + VecZnxBigAddSmallInplace - + VecZnxBigNormalize - + VecZnxAutomorphismInplace - + VecZnxBigBytesOf - + VecZnxNormalizeTmpBytes - + VecZnxDftCopy - + VecZnxDftAddInplace - + VecZnxIdftApplyTmpA - + VecZnxNormalize, - Scratch: ScratchAvailable, + pub fn automorphism_inplace(&mut self, module: &M, key: &K, tsk: &T, scratch: &mut Scratch) + where + K: AutomorphismKeyPreparedToRef, + T: TensorKeyPreparedToRef, + Scratch: ScratchTakeCore, + M: GGSWAutomorphism, { - // Keyswitch the j-th row of the col 0 - (0..self.dnum().into()).for_each(|row_i| { - // Key-switch column 0, i.e. - // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2) - self.at_mut(row_i, 0) - .automorphism_inplace(module, auto_key, scratch); - }); - self.expand_row(module, tensor_key, scratch); + module.ggsw_automorphism_inplace(self, key, tsk, scratch); } } + +impl GGSWAutomorphism for Module where Self: GLWEAutomorphism + GGSWExpandRows {} + +pub trait GGSWAutomorphism +where + Self: GLWEAutomorphism + GGSWExpandRows, +{ + fn ggsw_automorphism_tmp_bytes(&self, res_infos: &R, a_infos: &A, key_infos: &K, tsk_infos: &T) -> usize + where + R: GGSWInfos, + A: GGSWInfos, + K: GGLWEInfos, + T: GGLWEInfos, + { + let out_size: usize = res_infos.size(); + let ci_dft: usize = self.bytes_of_vec_znx_dft((key_infos.rank_out() + 1).into(), out_size); + let ks_internal: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos); + let expand: usize = self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos); + ci_dft + (ks_internal.max(expand)) + } + + fn ggsw_automorphism(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch) + where + R: GGSWToMut, + A: GGSWToRef, + K: AutomorphismKeyPreparedToRef, + T: TensorKeyPreparedToRef, + Scratch: ScratchTakeCore, + { + let res: &mut GGSW<&mut [u8]> = &mut res.to_mut(); + let a: &GGSW<&[u8]> = &a.to_ref(); + let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref(); + let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref(); + + assert_eq!(res.ggsw_layout(), a.ggsw_layout()); + assert_eq!(res.glwe_layout(), a.glwe_layout()); + assert_eq!(res.lwe_layout(), a.lwe_layout()); + assert!(scratch.available() >= self.ggsw_automorphism_tmp_bytes(res, a, key, tsk)); + + // Keyswitch the j-th row of the col 0 + for row in 0..res.dnum().as_usize() { + // Key-switch column 0, i.e. + // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2) + self.glwe_automorphism(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch); + } + + self.ggsw_expand_row(res, tsk, scratch); + } + + fn ggsw_automorphism_inplace(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch) + where + R: GGSWToMut, + K: AutomorphismKeyPreparedToRef, + T: TensorKeyPreparedToRef, + Scratch: ScratchTakeCore, + { + let res: &mut GGSW<&mut [u8]> = &mut res.to_mut(); + let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref(); + let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref(); + + // Keyswitch the j-th row of the col 0 + for row in 0..res.dnum().as_usize() { + // Key-switch column 0, i.e. + // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2) + self.glwe_automorphism_inplace(&mut res.at_mut(row, 0), key, scratch); + } + + self.ggsw_expand_row(res, tsk, scratch); + } +} + +impl GGSW {}