Add cross-basek normalization (#90)

* added cross_basek_normalization

* updated method signatures to take layouts

* fixed cross-base normalization

fix #91
fix #93
This commit is contained in:
Jean-Philippe Bossuat
2025-09-30 14:40:10 +02:00
committed by GitHub
parent 4da790ea6a
commit 37e13b965c
216 changed files with 12481 additions and 7745 deletions

View File

@@ -1,43 +1,42 @@
use poulpy_hal::{
api::{
ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume,
VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
},
layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
};
use crate::layouts::{GGLWEAutomorphismKey, GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
use crate::layouts::{GGLWEAutomorphismKey, GGLWELayoutInfos, GLWECiphertext, prepared::GGLWEAutomorphismKeyPrepared};
impl GGLWEAutomorphismKey<Vec<u8>> {
#[allow(clippy::too_many_arguments)]
pub fn automorphism_scratch_space<B: Backend>(
pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY>(
module: &Module<B>,
basek: usize,
k_out: usize,
k_in: usize,
k_ksk: usize,
digits: usize,
rank: usize,
out_infos: &OUT,
in_infos: &IN,
key_infos: &KEY,
) -> usize
where
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
OUT: GGLWELayoutInfos,
IN: GGLWELayoutInfos,
KEY: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
{
GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank, rank)
GLWECiphertext::keyswitch_scratch_space(
module,
&out_infos.glwe_layout(),
&in_infos.glwe_layout(),
key_infos,
)
}
pub fn automorphism_inplace_scratch_space<B: Backend>(
module: &Module<B>,
basek: usize,
k_out: usize,
k_ksk: usize,
digits: usize,
rank: usize,
) -> usize
pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
where
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
OUT: GGLWELayoutInfos,
KEY: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
{
GGLWEAutomorphismKey::automorphism_scratch_space(module, basek, k_out, k_out, k_ksk, digits, rank)
GGLWEAutomorphismKey::automorphism_scratch_space(module, out_infos, out_infos, key_infos)
}
}
@@ -59,11 +58,15 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
+ VecZnxBigAddSmallInplace<B>
+ VecZnxBigNormalize<B>
+ VecZnxAutomorphism
+ VecZnxAutomorphismInplace<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B>,
+ VecZnxAutomorphismInplace<B>
+ VecZnxNormalize<B>
+ VecZnxNormalizeTmpBytes,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
{
#[cfg(debug_assertions)]
{
use crate::layouts::LWEInfos;
assert_eq!(
self.rank_in(),
lhs.rank_in(),
@@ -93,13 +96,13 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
)
}
let cols_out: usize = rhs.rank_out() + 1;
let cols_out: usize = (rhs.rank_out() + 1).into();
let p: i64 = lhs.p();
let p_inv = module.galois_element_inv(p);
let p_inv: i64 = module.galois_element_inv(p);
(0..self.rank_in()).for_each(|col_i| {
(0..self.rows()).for_each(|row_j| {
(0..self.rank_in().into()).for_each(|col_i| {
(0..self.rows().into()).for_each(|row_j| {
let mut res_ct: GLWECiphertext<&mut [u8]> = self.at_mut(row_j, col_i);
let lhs_ct: GLWECiphertext<&[u8]> = lhs.at(row_j, col_i);
@@ -118,8 +121,8 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
});
});
(self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
(0..self.rank_in()).for_each(|col_j| {
(self.rows().min(lhs.rows()).into()..self.rows().into()).for_each(|row_i| {
(0..self.rank_in().into()).for_each(|col_j| {
self.at_mut(row_i, col_j).data.zero();
});
});
@@ -143,8 +146,10 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
+ VecZnxBigAddSmallInplace<B>
+ VecZnxBigNormalize<B>
+ VecZnxAutomorphism
+ VecZnxAutomorphismInplace<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B>,
+ VecZnxAutomorphismInplace<B>
+ VecZnxNormalize<B>
+ VecZnxNormalizeTmpBytes,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
{
#[cfg(debug_assertions)]
{
@@ -164,13 +169,13 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
);
}
let cols_out: usize = rhs.rank_out() + 1;
let cols_out: usize = (rhs.rank_out() + 1).into();
let p: i64 = self.p();
let p_inv = module.galois_element_inv(p);
(0..self.rank_in()).for_each(|col_i| {
(0..self.rows()).for_each(|row_j| {
(0..self.rank_in().into()).for_each(|col_i| {
(0..self.rows().into()).for_each(|row_j| {
let mut res_ct: GLWECiphertext<&mut [u8]> = self.at_mut(row_j, col_i);
// Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)