keyswitch tests

This commit is contained in:
Pro7ech
2025-10-20 15:32:52 +02:00
parent 0c894c19db
commit 252eda36fe
60 changed files with 918 additions and 945 deletions

View File

@@ -7,10 +7,8 @@ use crate::{
ScratchTakeCore,
automorphism::glwe_ct::GLWEAutomorphism,
layouts::{
AutomorphismKey, AutomorphismKeyToMut, AutomorphismKeyToRef, GGLWEInfos, GLWE, GLWEInfos,
prepared::{
AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement, SetAutomorphismGaloisElement,
},
AutomorphismKey, GGLWE, GGLWEInfos, GGLWEPreparedToRef, GGLWEToMut, GGLWEToRef, GLWE,
prepared::{GetAutomorphismGaloisElement, SetAutomorphismGaloisElement},
},
};
@@ -29,8 +27,8 @@ impl AutomorphismKey<Vec<u8>> {
impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
pub fn automorphism<A, K, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
where
A: AutomorphismKeyToRef + GetAutomorphismGaloisElement,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
A: GGLWEToRef + GetAutomorphismGaloisElement + GGLWEInfos,
K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
M: AutomorphismKeyAutomorphism<BE>,
{
@@ -39,7 +37,7 @@ impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
pub fn automorphism_inplace<K, M, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
where
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
M: AutomorphismKeyAutomorphism<BE>,
{
@@ -67,35 +65,34 @@ where
fn automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
where
R: AutomorphismKeyToMut + SetAutomorphismGaloisElement,
A: AutomorphismKeyToRef + GetAutomorphismGaloisElement,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
R: GGLWEToMut + SetAutomorphismGaloisElement + GGLWEInfos,
A: GGLWEToRef + GetAutomorphismGaloisElement + GGLWEInfos,
K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
assert!(
res.dnum().as_u32() <= a.dnum().as_u32(),
"res dnum: {} > a dnum: {}",
res.dnum(),
a.dnum()
);
assert_eq!(
res.dsize(),
a.dsize(),
"res dnum: {} != a dnum: {}",
res.dsize(),
a.dsize()
);
let cols_out: usize = (key.rank_out() + 1).into();
let p: i64 = a.p();
let p_inv: i64 = self.galois_element_inv(p);
{
let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
let a: &AutomorphismKey<&[u8]> = &a.to_ref();
let key: &AutomorphismKeyPrepared<&[u8], _> = &key.to_ref();
assert!(
res.dnum().as_u32() <= a.dnum().as_u32(),
"res dnum: {} > a dnum: {}",
res.dnum(),
a.dnum()
);
assert_eq!(
res.dsize(),
a.dsize(),
"res dnum: {} != a dnum: {}",
res.dsize(),
a.dsize()
);
let cols_out: usize = (key.rank_out() + 1).into();
let p: i64 = a.p();
let p_inv: i64 = self.galois_element_inv(p);
let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
let a: &GGLWE<&[u8]> = &a.to_ref();
for row in 0..res.dnum().as_usize() {
for col in 0..cols_out {
@@ -104,11 +101,11 @@ where
// Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
for i in 0..cols_out {
self.vec_znx_automorphism(a.p(), res_tmp.data_mut(), i, &a_ct.data, i);
self.vec_znx_automorphism(p, res_tmp.data_mut(), i, &a_ct.data, i);
}
// Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
self.glwe_keyswitch_inplace(&mut res_tmp, &key.key, scratch);
self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
// Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
(0..cols_out).for_each(|i| {
@@ -118,32 +115,29 @@ where
}
}
res.set_p((a.p() * key.p()) % (self.cyclotomic_order() as i64));
res.set_p((p * key.p()) % (self.cyclotomic_order() as i64));
}
fn automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
where
R: AutomorphismKeyToMut + SetAutomorphismGaloisElement + GetAutomorphismGaloisElement,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
R: GGLWEToMut + SetAutomorphismGaloisElement + GetAutomorphismGaloisElement + GGLWEInfos,
K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
assert_eq!(
res.rank(),
key.rank(),
"key rank: {} != key rank: {}",
res.rank(),
key.rank()
);
let cols_out: usize = (key.rank_out() + 1).into();
let p: i64 = res.p();
let p_inv: i64 = self.galois_element_inv(p);
{
let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
let key: &AutomorphismKeyPrepared<&[u8], _> = &key.to_ref();
assert_eq!(
res.rank(),
key.rank(),
"key rank: {} != key rank: {}",
res.rank(),
key.rank()
);
let cols_out: usize = (key.rank_out() + 1).into();
let p: i64 = res.p();
let p_inv: i64 = self.galois_element_inv(p);
let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
for row in 0..res.dnum().as_usize() {
for col in 0..cols_out {
let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
@@ -154,7 +148,7 @@ where
}
// Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
self.glwe_keyswitch_inplace(&mut res_tmp, &key.key, scratch);
self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
// Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
for i in 0..cols_out {

View File

@@ -7,8 +7,8 @@ use crate::{
GGSWExpandRows, ScratchTakeCore,
automorphism::glwe_ct::GLWEAutomorphism,
layouts::{
GGLWEInfos, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos,
prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, TensorKeyPrepared, TensorKeyPreparedToRef},
GGLWEInfos, GGLWEPreparedToRef, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GetAutomorphismGaloisElement,
prepared::{TensorKeyPrepared, TensorKeyPreparedToRef},
},
};
@@ -35,7 +35,7 @@ impl<D: DataMut> GGSW<D> {
pub fn automorphism<A, K, T, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
where
A: GGSWToRef,
K: AutomorphismKeyPreparedToRef<BE>,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
T: TensorKeyPreparedToRef<BE>,
Scratch<BE>: ScratchTakeCore<BE>,
M: GGSWAutomorphism<BE>,
@@ -45,7 +45,7 @@ impl<D: DataMut> GGSW<D> {
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>,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
T: TensorKeyPreparedToRef<BE>,
Scratch<BE>: ScratchTakeCore<BE>,
M: GGSWAutomorphism<BE>,
@@ -78,18 +78,16 @@ where
where
R: GGSWToMut,
A: GGSWToRef,
K: AutomorphismKeyPreparedToRef<BE>,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
T: TensorKeyPreparedToRef<BE>,
Scratch<BE>: ScratchTakeCore<BE>,
{
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_eq!(res.dsize(), a.dsize());
assert!(res.dnum() <= a.dnum());
assert!(scratch.available() >= self.ggsw_automorphism_tmp_bytes(res, a, key, tsk));
// Keyswitch the j-th row of the col 0
@@ -105,12 +103,11 @@ where
fn ggsw_automorphism_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
where
R: GGSWToMut,
K: AutomorphismKeyPreparedToRef<BE>,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
T: TensorKeyPreparedToRef<BE>,
Scratch<BE>: ScratchTakeCore<BE>,
{
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

View File

@@ -9,8 +9,7 @@ use poulpy_hal::{
use crate::{
GLWEKeyswitch, ScratchTakeCore, keyswitch_internal,
layouts::{
GGLWEInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos,
prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement},
GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::GetAutomorphismGaloisElement,
},
};
@@ -31,7 +30,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
where
M: GLWEAutomorphism<BE>,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism(self, a, key, scratch);
@@ -41,7 +40,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
where
M: GLWEAutomorphism<BE>,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_add(self, a, key, scratch);
@@ -51,7 +50,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
where
M: GLWEAutomorphism<BE>,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_sub(self, a, key, scratch);
@@ -61,7 +60,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
where
M: GLWEAutomorphism<BE>,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_sub_negate(self, a, key, scratch);
@@ -70,7 +69,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
pub fn automorphism_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
where
M: GLWEAutomorphism<BE>,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_inplace(self, key, scratch);
@@ -79,7 +78,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
pub fn automorphism_add_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
where
M: GLWEAutomorphism<BE>,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_add_inplace(self, key, scratch);
@@ -88,7 +87,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
pub fn automorphism_sub_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
where
M: GLWEAutomorphism<BE>,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_sub_inplace(self, key, scratch);
@@ -97,7 +96,7 @@ impl<DataSelf: DataMut> GLWE<DataSelf> {
pub fn automorphism_sub_negate_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
where
M: GLWEAutomorphism<BE>,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.glwe_automorphism_sub_negate_inplace(self, key, scratch);
@@ -125,10 +124,10 @@ where
where
R: GLWEToMut,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
self.glwe_keyswitch(res, a, &key.to_ref().key, scratch);
self.glwe_keyswitch(res, a, key, scratch);
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
@@ -140,10 +139,10 @@ where
fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
self.glwe_keyswitch_inplace(res, &key.to_ref().key, scratch);
self.glwe_keyswitch_inplace(res, key, scratch);
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
@@ -156,15 +155,14 @@ where
where
R: GLWEToMut,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let a: &GLWE<&[u8]> = &a.to_ref();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
@@ -184,14 +182,13 @@ where
fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
@@ -212,15 +209,14 @@ where
where
R: GLWEToMut,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let a: &GLWE<&[u8]> = &a.to_ref();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
@@ -241,15 +237,14 @@ where
where
R: GLWEToMut,
A: GLWEToRef,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let a: &GLWE<&[u8]> = &a.to_ref();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
@@ -269,14 +264,13 @@ where
fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
@@ -296,14 +290,13 @@ where
fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
K: GetAutomorphismGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
for i in 0..res.rank().as_usize() + 1 {
self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);