mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 05:06:44 +01:00
wip
This commit is contained in:
Pro7ech
committed by
Jean-Philippe Bossuat
Jean-Philippe Bossuat
parent
f72363cc4b
commit
2b2b994f7d
@@ -1,224 +1,205 @@
|
||||
use poulpy_hal::{
|
||||
api::{
|
||||
ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
|
||||
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
|
||||
VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
|
||||
use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch};
|
||||
|
||||
use crate::{
|
||||
ScratchTakeCore,
|
||||
keyswitching::glwe_ct::GLWEKeySwitch,
|
||||
layouts::{
|
||||
AutomorphismKey, AutomorphismKeyToRef, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWESwitchingKey,
|
||||
GLWESwitchingKeyToRef,
|
||||
prepared::{GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToRef},
|
||||
},
|
||||
layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
|
||||
};
|
||||
|
||||
use crate::layouts::{
|
||||
GGLWEAutomorphismKey, GGLWEInfos, GGLWESwitchingKey, GLWECiphertext, GLWEInfos,
|
||||
prepared::{GGLWEAutomorphismKeyPrepared, GGLWESwitchingKeyPrepared},
|
||||
};
|
||||
|
||||
impl GGLWEAutomorphismKey<Vec<u8>> {
|
||||
pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
|
||||
module: &Module<B>,
|
||||
out_infos: &OUT,
|
||||
in_infos: &IN,
|
||||
key_infos: &KEY,
|
||||
) -> usize
|
||||
impl AutomorphismKey<Vec<u8>> {
|
||||
pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
|
||||
where
|
||||
OUT: GGLWEInfos,
|
||||
IN: GGLWEInfos,
|
||||
KEY: GGLWEInfos,
|
||||
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
|
||||
R: GGLWEInfos,
|
||||
A: GGLWEInfos,
|
||||
K: GGLWEInfos,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
GGLWESwitchingKey::keyswitch_scratch_space(module, out_infos, in_infos, key_infos)
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
|
||||
where
|
||||
OUT: GGLWEInfos,
|
||||
KEY: GGLWEInfos,
|
||||
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
|
||||
{
|
||||
GGLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_infos, key_infos)
|
||||
module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
|
||||
pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
|
||||
&mut self,
|
||||
module: &Module<B>,
|
||||
lhs: &GGLWEAutomorphismKey<DataLhs>,
|
||||
rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
|
||||
scratch: &mut Scratch<B>,
|
||||
) where
|
||||
Module<B>: VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<B>
|
||||
+ VmpApplyDftToDftAdd<B>
|
||||
+ VecZnxDftApply<B>
|
||||
+ VecZnxIdftApplyConsume<B>
|
||||
+ VecZnxBigAddSmallInplace<B>
|
||||
+ VecZnxBigNormalize<B>
|
||||
+ VecZnxNormalize<B>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
|
||||
impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
|
||||
pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: AutomorphismKeyToRef,
|
||||
B: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
self.key.keyswitch(module, &lhs.key, rhs, scratch);
|
||||
module.gglwe_keyswitch(&mut self.key.key, &a.to_ref().key.key, b, scratch);
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
|
||||
&mut self,
|
||||
module: &Module<B>,
|
||||
rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
|
||||
scratch: &mut Scratch<B>,
|
||||
) where
|
||||
Module<B>: VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<B>
|
||||
+ VmpApplyDftToDftAdd<B>
|
||||
+ VecZnxDftApply<B>
|
||||
+ VecZnxIdftApplyConsume<B>
|
||||
+ VecZnxBigAddSmallInplace<B>
|
||||
+ VecZnxBigNormalize<B>
|
||||
+ VecZnxNormalize<B>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
|
||||
pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
self.key.keyswitch_inplace(module, &rhs.key, scratch);
|
||||
module.gglwe_keyswitch_inplace(&mut self.key.key, a, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
impl GGLWESwitchingKey<Vec<u8>> {
|
||||
pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
|
||||
module: &Module<B>,
|
||||
out_infos: &OUT,
|
||||
in_infos: &IN,
|
||||
key_apply: &KEY,
|
||||
) -> usize
|
||||
impl GLWESwitchingKey<Vec<u8>> {
|
||||
pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
|
||||
where
|
||||
OUT: GGLWEInfos,
|
||||
IN: GGLWEInfos,
|
||||
KEY: GGLWEInfos,
|
||||
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
|
||||
R: GGLWEInfos,
|
||||
A: GGLWEInfos,
|
||||
K: GGLWEInfos,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
GLWECiphertext::keyswitch_scratch_space(module, out_infos, in_infos, key_apply)
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_apply: &KEY) -> usize
|
||||
where
|
||||
OUT: GGLWEInfos + GLWEInfos,
|
||||
KEY: GGLWEInfos + GLWEInfos,
|
||||
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
|
||||
{
|
||||
GLWECiphertext::keyswitch_inplace_scratch_space(module, out_infos, key_apply)
|
||||
module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf: DataMut> GGLWESwitchingKey<DataSelf> {
|
||||
pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
|
||||
&mut self,
|
||||
module: &Module<B>,
|
||||
lhs: &GGLWESwitchingKey<DataLhs>,
|
||||
rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
|
||||
scratch: &mut Scratch<B>,
|
||||
) where
|
||||
Module<B>: VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<B>
|
||||
+ VmpApplyDftToDftAdd<B>
|
||||
+ VecZnxDftApply<B>
|
||||
+ VecZnxIdftApplyConsume<B>
|
||||
+ VecZnxBigAddSmallInplace<B>
|
||||
+ VecZnxBigNormalize<B>
|
||||
+ VecZnxNormalize<B>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
|
||||
impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
|
||||
pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: GLWESwitchingKeyToRef,
|
||||
B: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(
|
||||
self.rank_in(),
|
||||
lhs.rank_in(),
|
||||
"ksk_out input rank: {} != ksk_in input rank: {}",
|
||||
self.rank_in(),
|
||||
lhs.rank_in()
|
||||
);
|
||||
assert_eq!(
|
||||
lhs.rank_out(),
|
||||
rhs.rank_in(),
|
||||
"ksk_in output rank: {} != ksk_apply input rank: {}",
|
||||
self.rank_out(),
|
||||
rhs.rank_in()
|
||||
);
|
||||
assert_eq!(
|
||||
self.rank_out(),
|
||||
rhs.rank_out(),
|
||||
"ksk_out output rank: {} != ksk_apply output rank: {}",
|
||||
self.rank_out(),
|
||||
rhs.rank_out()
|
||||
);
|
||||
assert!(
|
||||
self.dnum() <= lhs.dnum(),
|
||||
"self.dnum()={} > lhs.dnum()={}",
|
||||
self.dnum(),
|
||||
lhs.dnum()
|
||||
);
|
||||
assert_eq!(
|
||||
self.dsize(),
|
||||
lhs.dsize(),
|
||||
"ksk_out dsize: {} != ksk_in dsize: {}",
|
||||
self.dsize(),
|
||||
lhs.dsize()
|
||||
)
|
||||
module.gglwe_keyswitch(&mut self.key, &a.to_ref().key, b, scratch);
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
module.gglwe_keyswitch_inplace(&mut self.key, a, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
impl GGLWE<Vec<u8>> {
|
||||
pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
|
||||
where
|
||||
R: GGLWEInfos,
|
||||
A: GGLWEInfos,
|
||||
K: GGLWEInfos,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf: DataMut> GGLWE<DataSelf> {
|
||||
pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: GGLWEToRef,
|
||||
B: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
module.gglwe_keyswitch(self, a, b, scratch);
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
A: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
M: GGLWEKeySwitch<BE>,
|
||||
{
|
||||
module.gglwe_keyswitch_inplace(self, a, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
impl<BE: Backend> GGLWEKeySwitch<BE> for Module<BE> where Self: GLWEKeySwitch<BE> {}
|
||||
|
||||
pub trait GGLWEKeySwitch<BE: Backend>
|
||||
where
|
||||
Self: GLWEKeySwitch<BE>,
|
||||
{
|
||||
fn gglwe_keyswitch_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
|
||||
where
|
||||
R: GGLWEInfos,
|
||||
A: GGLWEInfos,
|
||||
K: GGLWEInfos,
|
||||
{
|
||||
self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
|
||||
}
|
||||
|
||||
fn gglwe_keyswitch<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
R: GGLWEToMut,
|
||||
A: GGLWEToRef,
|
||||
B: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
{
|
||||
let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
|
||||
let a: &GGLWE<&[u8]> = &a.to_ref();
|
||||
let b: &GLWESwitchingKeyPrepared<&[u8], BE> = &b.to_ref();
|
||||
|
||||
assert_eq!(
|
||||
res.rank_in(),
|
||||
a.rank_in(),
|
||||
"res input rank: {} != a input rank: {}",
|
||||
res.rank_in(),
|
||||
a.rank_in()
|
||||
);
|
||||
assert_eq!(
|
||||
a.rank_out(),
|
||||
b.rank_in(),
|
||||
"res output rank: {} != b input rank: {}",
|
||||
a.rank_out(),
|
||||
b.rank_in()
|
||||
);
|
||||
assert_eq!(
|
||||
res.rank_out(),
|
||||
b.rank_out(),
|
||||
"res output rank: {} != b output rank: {}",
|
||||
res.rank_out(),
|
||||
b.rank_out()
|
||||
);
|
||||
assert!(
|
||||
res.dnum() <= a.dnum(),
|
||||
"res.dnum()={} > a.dnum()={}",
|
||||
res.dnum(),
|
||||
a.dnum()
|
||||
);
|
||||
assert_eq!(
|
||||
res.dsize(),
|
||||
a.dsize(),
|
||||
"res dsize: {} != a dsize: {}",
|
||||
res.dsize(),
|
||||
a.dsize()
|
||||
);
|
||||
|
||||
for row in 0..res.dnum().into() {
|
||||
for col in 0..res.rank_in().into() {
|
||||
self.glwe_keyswitch(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
(0..self.rank_in().into()).for_each(|col_i| {
|
||||
(0..self.dnum().into()).for_each(|row_j| {
|
||||
self.at_mut(row_j, col_i)
|
||||
.keyswitch(module, &lhs.at(row_j, col_i), rhs, scratch);
|
||||
});
|
||||
});
|
||||
|
||||
(self.dnum().min(lhs.dnum()).into()..self.dnum().into()).for_each(|row_i| {
|
||||
(0..self.rank_in().into()).for_each(|col_j| {
|
||||
self.at_mut(row_i, col_j).data.zero();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
|
||||
&mut self,
|
||||
module: &Module<B>,
|
||||
rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
|
||||
scratch: &mut Scratch<B>,
|
||||
) where
|
||||
Module<B>: VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<B>
|
||||
+ VmpApplyDftToDftAdd<B>
|
||||
+ VecZnxDftApply<B>
|
||||
+ VecZnxIdftApplyConsume<B>
|
||||
+ VecZnxBigAddSmallInplace<B>
|
||||
+ VecZnxBigNormalize<B>
|
||||
+ VecZnxNormalize<B>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
|
||||
fn gglwe_keyswitch_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
|
||||
where
|
||||
R: GGLWEToMut,
|
||||
A: GLWESwitchingKeyPreparedToRef<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
{
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(
|
||||
self.rank_out(),
|
||||
rhs.rank_out(),
|
||||
"ksk_out output rank: {} != ksk_apply output rank: {}",
|
||||
self.rank_out(),
|
||||
rhs.rank_out()
|
||||
);
|
||||
}
|
||||
let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
|
||||
let a: &GLWESwitchingKeyPrepared<&[u8], BE> = &a.to_ref();
|
||||
|
||||
(0..self.rank_in().into()).for_each(|col_i| {
|
||||
(0..self.dnum().into()).for_each(|row_j| {
|
||||
self.at_mut(row_j, col_i)
|
||||
.keyswitch_inplace(module, rhs, scratch)
|
||||
});
|
||||
});
|
||||
assert_eq!(
|
||||
res.rank_out(),
|
||||
a.rank_out(),
|
||||
"res output rank: {} != a output rank: {}",
|
||||
res.rank_out(),
|
||||
a.rank_out()
|
||||
);
|
||||
|
||||
for row in 0..res.dnum().into() {
|
||||
for col in 0..res.rank_in().into() {
|
||||
self.glwe_keyswitch_inplace(&mut res.at_mut(row, col), a, scratch);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {}
|
||||
|
||||
Reference in New Issue
Block a user