mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 05:06:44 +01:00
automorphism ggsw
This commit is contained in:
Jean-Philippe Bossuat
@@ -1,165 +1,127 @@
|
|||||||
use poulpy_hal::{
|
use poulpy_hal::{
|
||||||
api::{
|
api::ScratchAvailable,
|
||||||
ScratchAvailable, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize,
|
layouts::{Backend, DataMut, Module, Scratch},
|
||||||
VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume,
|
|
||||||
VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
|
|
||||||
VmpApplyDftToDftTmpBytes,
|
|
||||||
},
|
|
||||||
layouts::{Backend, DataMut, DataRef, Module, Scratch},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::layouts::{
|
use crate::{
|
||||||
GGLWEInfos, GGSW, GGSWInfos, GLWE,
|
GGSWExpandRows, ScratchTakeCore,
|
||||||
prepared::{AutomorphismKeyPrepared, TensorKeyPrepared},
|
automorphism::glwe_ct::GLWEAutomorphism,
|
||||||
|
layouts::{
|
||||||
|
GGLWEInfos, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos,
|
||||||
|
prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, TensorKeyPrepared, TensorKeyPreparedToRef},
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
impl GGSW<Vec<u8>> {
|
impl GGSW<Vec<u8>> {
|
||||||
pub fn automorphism_tmp_bytes<B: Backend, OUT, IN, KEY, TSK>(
|
pub fn automorphism_tmp_bytes<R, A, K, T, M, BE: Backend>(
|
||||||
module: &Module<B>,
|
module: &M,
|
||||||
out_infos: &OUT,
|
res_infos: &R,
|
||||||
in_infos: &IN,
|
a_infos: &A,
|
||||||
key_infos: &KEY,
|
key_infos: &K,
|
||||||
tsk_infos: &TSK,
|
tsk_infos: &T,
|
||||||
) -> usize
|
) -> usize
|
||||||
where
|
where
|
||||||
OUT: GGSWInfos,
|
R: GGSWInfos,
|
||||||
IN: GGSWInfos,
|
A: GGSWInfos,
|
||||||
KEY: GGLWEInfos,
|
K: GGLWEInfos,
|
||||||
TSK: GGLWEInfos,
|
T: GGLWEInfos,
|
||||||
Module<B>:
|
M: GGSWAutomorphism<BE>,
|
||||||
VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
|
|
||||||
{
|
{
|
||||||
let out_size: usize = out_infos.size();
|
module.ggsw_automorphism_tmp_bytes(res_infos, a_infos, key_infos, tsk_infos)
|
||||||
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<B: Backend, OUT, KEY, TSK>(
|
impl<D: DataMut> GGSW<D> {
|
||||||
module: &Module<B>,
|
pub fn automorphism<A, K, T, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
|
||||||
out_infos: &OUT,
|
|
||||||
key_infos: &KEY,
|
|
||||||
tsk_infos: &TSK,
|
|
||||||
) -> usize
|
|
||||||
where
|
where
|
||||||
OUT: GGSWInfos,
|
A: GGSWToRef,
|
||||||
KEY: GGLWEInfos,
|
K: AutomorphismKeyPreparedToRef<BE>,
|
||||||
TSK: GGLWEInfos,
|
T: TensorKeyPreparedToRef<BE>,
|
||||||
Module<B>:
|
Scratch<BE>: ScratchTakeCore<BE>,
|
||||||
VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
|
M: GGSWAutomorphism<BE>,
|
||||||
{
|
{
|
||||||
GGSW::automorphism_tmp_bytes(module, out_infos, out_infos, key_infos, tsk_infos)
|
module.ggsw_automorphism(self, a, key, tsk, scratch);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn automorphism_inplace<K, T, M, BE: Backend>(&mut self, module: &M, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
|
||||||
|
where
|
||||||
|
K: AutomorphismKeyPreparedToRef<BE>,
|
||||||
|
T: TensorKeyPreparedToRef<BE>,
|
||||||
|
Scratch<BE>: ScratchTakeCore<BE>,
|
||||||
|
M: GGSWAutomorphism<BE>,
|
||||||
|
{
|
||||||
|
module.ggsw_automorphism_inplace(self, key, tsk, scratch);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<DataSelf: DataMut> GGSW<DataSelf> {
|
impl<BE: Backend> GGSWAutomorphism<BE> for Module<BE> where Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE> {}
|
||||||
pub fn automorphism<DataLhs: DataRef, DataAk: DataRef, DataTsk: DataRef, B: Backend>(
|
|
||||||
&mut self,
|
|
||||||
module: &Module<B>,
|
|
||||||
lhs: &GGSW<DataLhs>,
|
|
||||||
auto_key: &AutomorphismKeyPrepared<DataAk, B>,
|
|
||||||
tensor_key: &TensorKeyPrepared<DataTsk, B>,
|
|
||||||
scratch: &mut Scratch<B>,
|
|
||||||
) where
|
|
||||||
Module<B>: VecZnxDftBytesOf
|
|
||||||
+ VmpApplyDftToDftTmpBytes
|
|
||||||
+ VecZnxBigNormalizeTmpBytes
|
|
||||||
+ VmpApplyDftToDft<B>
|
|
||||||
+ VmpApplyDftToDftAdd<B>
|
|
||||||
+ VecZnxDftApply<B>
|
|
||||||
+ VecZnxIdftApplyConsume<B>
|
|
||||||
+ VecZnxBigAddSmallInplace<B>
|
|
||||||
+ VecZnxBigNormalize<B>
|
|
||||||
+ VecZnxAutomorphismInplace<B>
|
|
||||||
+ VecZnxBigBytesOf
|
|
||||||
+ VecZnxNormalizeTmpBytes
|
|
||||||
+ VecZnxDftCopy<B>
|
|
||||||
+ VecZnxDftAddInplace<B>
|
|
||||||
+ VecZnxIdftApplyTmpA<B>
|
|
||||||
+ VecZnxNormalize<B>,
|
|
||||||
Scratch<B>: ScratchAvailable,
|
|
||||||
{
|
|
||||||
#[cfg(debug_assertions)]
|
|
||||||
{
|
|
||||||
use crate::layouts::{GLWEInfos, LWEInfos};
|
|
||||||
|
|
||||||
assert_eq!(self.n(), module.n() as u32);
|
pub trait GGSWAutomorphism<BE: Backend>
|
||||||
assert_eq!(lhs.n(), module.n() as u32);
|
where
|
||||||
assert_eq!(auto_key.n(), module.n() as u32);
|
Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE>,
|
||||||
assert_eq!(tensor_key.n(), module.n() as u32);
|
{
|
||||||
|
fn ggsw_automorphism_tmp_bytes<R, A, K, T>(&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))
|
||||||
|
}
|
||||||
|
|
||||||
assert_eq!(
|
fn ggsw_automorphism<R, A, K, T>(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
|
||||||
self.rank(),
|
where
|
||||||
lhs.rank(),
|
R: GGSWToMut,
|
||||||
"ggsw_out rank: {} != ggsw_in rank: {}",
|
A: GGSWToRef,
|
||||||
self.rank(),
|
K: AutomorphismKeyPreparedToRef<BE>,
|
||||||
lhs.rank()
|
T: TensorKeyPreparedToRef<BE>,
|
||||||
);
|
Scratch<BE>: ScratchTakeCore<BE>,
|
||||||
assert_eq!(
|
{
|
||||||
self.rank(),
|
let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
|
||||||
auto_key.rank_out(),
|
let a: &GGSW<&[u8]> = &a.to_ref();
|
||||||
"ggsw_in rank: {} != auto_key rank: {}",
|
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
|
||||||
self.rank(),
|
let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
|
||||||
auto_key.rank_out()
|
|
||||||
);
|
assert_eq!(res.ggsw_layout(), a.ggsw_layout());
|
||||||
assert_eq!(
|
assert_eq!(res.glwe_layout(), a.glwe_layout());
|
||||||
self.rank(),
|
assert_eq!(res.lwe_layout(), a.lwe_layout());
|
||||||
tensor_key.rank_out(),
|
assert!(scratch.available() >= self.ggsw_automorphism_tmp_bytes(res, a, key, tsk));
|
||||||
"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
|
// Keyswitch the j-th row of the col 0
|
||||||
(0..lhs.dnum().into()).for_each(|row_i| {
|
for row in 0..res.dnum().as_usize() {
|
||||||
// Key-switch column 0, i.e.
|
// 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)
|
// 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)
|
self.glwe_automorphism(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch);
|
||||||
.automorphism(module, &lhs.at(row_i, 0), auto_key, scratch);
|
|
||||||
});
|
|
||||||
self.expand_row(module, tensor_key, scratch);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn automorphism_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
|
self.ggsw_expand_row(res, tsk, scratch);
|
||||||
&mut self,
|
}
|
||||||
module: &Module<B>,
|
|
||||||
auto_key: &AutomorphismKeyPrepared<DataKsk, B>,
|
fn ggsw_automorphism_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
|
||||||
tensor_key: &TensorKeyPrepared<DataTsk, B>,
|
where
|
||||||
scratch: &mut Scratch<B>,
|
R: GGSWToMut,
|
||||||
) where
|
K: AutomorphismKeyPreparedToRef<BE>,
|
||||||
Module<B>: VecZnxDftBytesOf
|
T: TensorKeyPreparedToRef<BE>,
|
||||||
+ VmpApplyDftToDftTmpBytes
|
Scratch<BE>: ScratchTakeCore<BE>,
|
||||||
+ VecZnxBigNormalizeTmpBytes
|
|
||||||
+ VmpApplyDftToDft<B>
|
|
||||||
+ VmpApplyDftToDftAdd<B>
|
|
||||||
+ VecZnxDftApply<B>
|
|
||||||
+ VecZnxIdftApplyConsume<B>
|
|
||||||
+ VecZnxBigAddSmallInplace<B>
|
|
||||||
+ VecZnxBigNormalize<B>
|
|
||||||
+ VecZnxAutomorphismInplace<B>
|
|
||||||
+ VecZnxBigBytesOf
|
|
||||||
+ VecZnxNormalizeTmpBytes
|
|
||||||
+ VecZnxDftCopy<B>
|
|
||||||
+ VecZnxDftAddInplace<B>
|
|
||||||
+ VecZnxIdftApplyTmpA<B>
|
|
||||||
+ VecZnxNormalize<B>,
|
|
||||||
Scratch<B>: ScratchAvailable,
|
|
||||||
{
|
{
|
||||||
|
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
|
// Keyswitch the j-th row of the col 0
|
||||||
(0..self.dnum().into()).for_each(|row_i| {
|
for row in 0..res.dnum().as_usize() {
|
||||||
// Key-switch column 0, i.e.
|
// 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)
|
// 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)
|
self.glwe_automorphism_inplace(&mut res.at_mut(row, 0), key, scratch);
|
||||||
.automorphism_inplace(module, auto_key, scratch);
|
}
|
||||||
});
|
|
||||||
self.expand_row(module, tensor_key, scratch);
|
self.ggsw_expand_row(res, tsk, scratch);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl<DataSelf: DataMut> GGSW<DataSelf> {}
|
||||||
|
|||||||
Reference in New Issue
Block a user