core refactoring (#69)

bin_fhe/circuit_bootstrapping/circuit_bootstrapping.rs

@@ -0,0 +1,514 @@
+use std::{collections::HashMap, time::Instant, usize};
+
+use backend::hal::{
+    api::{
+        ScratchAvailable, TakeMatZnx, TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, TakeVecZnxDftSlice,
+        TakeVecZnxSlice, VecZnxAddInplace, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
+        VecZnxBigAutomorphismInplace, VecZnxBigSubSmallBInplace, VecZnxCopy, VecZnxDftCopy, VecZnxDftToVecZnxBigTmpA,
+        VecZnxNegateInplace, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
+        VecZnxSubABInplace, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
+    },
+    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
+    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
+};
+use sampling::source::Source;
+
+use crate::{
+    BlindRotationKeyCGGI, BlindRotationKeyCGGIExec, BlindRotationKeyCGGIExecLayoutFamily, CCGIBlindRotationFamily, GLWEOps,
+    Infos, LookUpTable, LookUpTableRotationDirection, TakeGGLWE, TakeGLWECt, cggi_blind_rotate,
+    layouts::{
+        AutomorphismKey, GGSWCiphertext, GLWECiphertext, GLWESecret, GLWETensorKey, LWECiphertext, LWESecret,
+        prepared::{AutomorphismKeyExec, GLWESecretExec, GLWETensorKeyExec},
+    },
+};
+
+use crate::trait_families::{
+    AutomorphismKeyEncryptSkFamily, GGSWEncryptSkFamily, GLWESecretExecModuleFamily, GLWETensorKeyEncryptSkFamily,
+    GLWETraceModuleFamily,
+};
+
+pub struct CircuitBootstrappingKeyCGGI<D: Data> {
+    pub(crate) brk: BlindRotationKeyCGGI<D>,
+    pub(crate) tsk: GLWETensorKey<Vec<u8>>,
+    pub(crate) atk: HashMap<i64, AutomorphismKey<Vec<u8>>>,
+}
+
+impl CircuitBootstrappingKeyCGGI<Vec<u8>> {
+    pub fn generate<DLwe, DGlwe, B: Backend>(
+        module: &Module<B>,
+        basek: usize,
+        sk_lwe: &LWESecret<DLwe>,
+        sk_glwe: &GLWESecret<DGlwe>,
+        k_brk: usize,
+        rows_brk: usize,
+        k_trace: usize,
+        rows_trace: usize,
+        k_tsk: usize,
+        rows_tsk: usize,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        sigma: f64,
+        scratch: &mut Scratch<B>,
+    ) -> Self
+    where
+        Module<B>: GGSWEncryptSkFamily<B>
+            + GLWESecretExecModuleFamily<B>
+            + VecZnxAddScalarInplace
+            + AutomorphismKeyEncryptSkFamily<B>
+            + VecZnxAutomorphism
+            + VecZnxSwithcDegree
+            + GLWETensorKeyEncryptSkFamily<B>,
+        DLwe: DataRef,
+        DGlwe: DataRef,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
+    {
+        let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>>> = HashMap::new();
+        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(&module);
+        gal_els.iter().for_each(|gal_el| {
+            let mut key: AutomorphismKey<Vec<u8>> =
+                AutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
+            key.encrypt_sk(
+                &module, *gal_el, &sk_glwe, source_xa, source_xe, sigma, scratch,
+            );
+            auto_keys.insert(*gal_el, key);
+        });
+
+        let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
+
+        let mut brk: BlindRotationKeyCGGI<Vec<u8>> = BlindRotationKeyCGGI::alloc(
+            sk_glwe.n(),
+            sk_lwe.n(),
+            basek,
+            k_brk,
+            rows_brk,
+            sk_glwe.rank(),
+        );
+        brk.generate_from_sk(
+            module,
+            &sk_glwe_exec,
+            sk_lwe,
+            source_xa,
+            source_xe,
+            sigma,
+            scratch,
+        );
+
+        let mut tsk: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
+        tsk.encrypt_sk(module, &sk_glwe, source_xa, source_xe, sigma, scratch);
+
+        Self {
+            brk,
+            atk: auto_keys,
+            tsk,
+        }
+    }
+}
+
+pub struct CircuitBootstrappingKeyCGGIExec<D: Data, B: Backend> {
+    pub(crate) brk: BlindRotationKeyCGGIExec<D, B>,
+    pub(crate) tsk: GLWETensorKeyExec<Vec<u8>, B>,
+    pub(crate) atk: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
+}
+
+impl<B: Backend> CircuitBootstrappingKeyCGGIExec<Vec<u8>, B> {
+    pub fn from<DataOther: DataRef>(
+        module: &Module<B>,
+        other: &CircuitBootstrappingKeyCGGI<DataOther>,
+        scratch: &mut Scratch<B>,
+    ) -> CircuitBootstrappingKeyCGGIExec<Vec<u8>, B>
+    where
+        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + VmpPMatAlloc<B> + VmpPMatPrepare<B>,
+    {
+        let brk: BlindRotationKeyCGGIExec<Vec<u8>, B> = BlindRotationKeyCGGIExec::from(module, &other.brk, scratch);
+        let tsk: GLWETensorKeyExec<Vec<u8>, B> = GLWETensorKeyExec::from(module, &other.tsk, scratch);
+        let mut atk: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>> = HashMap::new();
+        for (key, value) in &other.atk {
+            atk.insert(*key, AutomorphismKeyExec::from(module, value, scratch));
+        }
+        CircuitBootstrappingKeyCGGIExec { brk, tsk, atk }
+    }
+}
+
+pub trait CGGICircuitBootstrapFamily<B: Backend> = VecZnxRotateInplace
+    + VecZnxNormalizeInplace<B>
+    + VecZnxNormalizeTmpBytes
+    + CCGIBlindRotationFamily<B>
+    + VecZnxSwithcDegree
+    + VecZnxBigAutomorphismInplace<B>
+    + VecZnxRshInplace
+    + VecZnxDftCopy<B>
+    + VecZnxDftToVecZnxBigTmpA<B>
+    + VecZnxSub
+    + VecZnxAddInplace
+    + VecZnxNegateInplace
+    + VecZnxCopy
+    + VecZnxSubABInplace
+    + GLWETraceModuleFamily<B>
+    + VecZnxRotateInplace
+    + VecZnxAutomorphismInplace
+    + VecZnxBigSubSmallBInplace<B>;
+
+pub fn circuit_bootstrap_to_constant_cggi<DRes, DLwe, DBrk, B: Backend>(
+    module: &Module<B>,
+    res: &mut GGSWCiphertext<DRes>,
+    lwe: &LWECiphertext<DLwe>,
+    log_domain: usize,
+    extension_factor: usize,
+    key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
+    scratch: &mut Scratch<B>,
+) where
+    DRes: DataMut,
+    DLwe: DataRef,
+    DBrk: DataRef,
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
+    Scratch<B>: TakeVecZnx
+        + TakeVecZnxDftSlice<B>
+        + TakeVecZnxBig<B>
+        + TakeVecZnxDft<B>
+        + TakeMatZnx
+        + ScratchAvailable
+        + TakeVecZnxSlice,
+{
+    circuit_bootstrap_core_cggi(
+        false,
+        module,
+        0,
+        res,
+        lwe,
+        log_domain,
+        extension_factor,
+        key,
+        scratch,
+    );
+}
+
+pub fn circuit_bootstrap_to_exponent_cggi<DRes, DLwe, DBrk, B: Backend>(
+    module: &Module<B>,
+    log_gap_out: usize,
+    res: &mut GGSWCiphertext<DRes>,
+    lwe: &LWECiphertext<DLwe>,
+    log_domain: usize,
+    extension_factor: usize,
+    key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
+    scratch: &mut Scratch<B>,
+) where
+    DRes: DataMut,
+    DLwe: DataRef,
+    DBrk: DataRef,
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
+    Scratch<B>: TakeVecZnx
+        + TakeVecZnxDftSlice<B>
+        + TakeVecZnxBig<B>
+        + TakeVecZnxDft<B>
+        + TakeMatZnx
+        + ScratchAvailable
+        + TakeVecZnxSlice,
+{
+    circuit_bootstrap_core_cggi(
+        true,
+        module,
+        log_gap_out,
+        res,
+        lwe,
+        log_domain,
+        extension_factor,
+        key,
+        scratch,
+    );
+}
+
+pub fn circuit_bootstrap_core_cggi<DRes, DLwe, DBrk, B: Backend>(
+    to_exponent: bool,
+    module: &Module<B>,
+    log_gap_out: usize,
+    res: &mut GGSWCiphertext<DRes>,
+    lwe: &LWECiphertext<DLwe>,
+    log_domain: usize,
+    extension_factor: usize,
+    key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
+    scratch: &mut Scratch<B>,
+) where
+    DRes: DataMut,
+    DLwe: DataRef,
+    DBrk: DataRef,
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
+    Scratch<B>: TakeGGLWE<B>
+        + TakeVecZnxDftSlice<B>
+        + TakeVecZnxBig<B>
+        + TakeVecZnxDft<B>
+        + TakeVecZnx
+        + ScratchAvailable
+        + TakeVecZnxSlice,
+{
+    #[cfg(debug_assertions)]
+    {
+        use crate::Infos;
+
+        assert_eq!(res.n(), key.brk.n());
+        assert_eq!(lwe.basek(), key.brk.basek());
+        assert_eq!(res.basek(), key.brk.basek());
+    }
+
+    let n: usize = res.n();
+    let basek: usize = res.basek();
+    let rows: usize = res.rows();
+    let rank: usize = res.rank();
+    let k: usize = res.k();
+
+    let alpha: usize = rows.next_power_of_two();
+
+    let mut f: Vec<i64> = vec![0i64; (1 << log_domain) * alpha];
+
+    if to_exponent {
+        (0..rows).for_each(|i| {
+            f[i] = 1 << (basek * (rows - 1 - i));
+        });
+    } else {
+        (0..1 << log_domain).for_each(|j| {
+            (0..rows).for_each(|i| {
+                f[j * alpha + i] = j as i64 * (1 << (basek * (rows - 1 - i)));
+            });
+        });
+    }
+
+    // Lut precision, basically must be able to hold the decomposition power basis of the GGSW
+    let mut lut: LookUpTable = LookUpTable::alloc(n, basek, basek * rows, extension_factor);
+    lut.set(module, &f, basek * rows);
+
+    if to_exponent {
+        lut.set_rotation_direction(LookUpTableRotationDirection::Right);
+    }
+
+    // TODO: separate GGSW k from output of blind rotation k
+    let (mut res_glwe, scratch1) = scratch.take_glwe_ct(n, basek, k, rank);
+    let (mut tmp_gglwe, scratch2) = scratch1.take_gglwe(n, basek, k, rows, 1, rank, rank);
+
+    let now: Instant = Instant::now();
+    cggi_blind_rotate(module, &mut res_glwe, &lwe, &lut, &key.brk, scratch2);
+    println!("cggi_blind_rotate: {} ms", now.elapsed().as_millis());
+
+    let gap: usize = 2 * lut.drift / lut.extension_factor();
+
+    let log_gap_in: usize = (usize::BITS - (gap * alpha - 1).leading_zeros()) as _;
+
+    (0..rows).for_each(|i| {
+        let mut tmp_glwe: GLWECiphertext<&mut [u8]> = tmp_gglwe.at_mut(i, 0);
+
+        if to_exponent {
+            let now: Instant = Instant::now();
+
+            // Isolates i-th LUT and moves coefficients according to requested gap.
+            post_process(
+                module,
+                &mut tmp_glwe,
+                &res_glwe,
+                log_gap_in,
+                log_gap_out,
+                log_domain,
+                &key.atk,
+                scratch2,
+            );
+            println!("post_process: {} ms", now.elapsed().as_millis());
+        } else {
+            tmp_glwe.trace(module, 0, module.log_n(), &res_glwe, &key.atk, scratch2);
+        }
+
+        if i < rows {
+            res_glwe.rotate_inplace(module, -(gap as i64));
+        }
+    });
+
+    // Expands GGLWE to GGSW using GGLWE(s^2)
+    res.from_gglwe(module, &tmp_gglwe, &key.tsk, scratch2);
+}
+
+fn post_process<DataRes, DataA, B: Backend>(
+    module: &Module<B>,
+    res: &mut GLWECiphertext<DataRes>,
+    a: &GLWECiphertext<DataA>,
+    log_gap_in: usize,
+    log_gap_out: usize,
+    log_domain: usize,
+    auto_keys: &HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
+    scratch: &mut Scratch<B>,
+) where
+    DataRes: DataMut,
+    DataA: DataRef,
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
+{
+    let log_n: usize = module.log_n();
+
+    let mut cts: HashMap<usize, GLWECiphertext<Vec<u8>>> = HashMap::new();
+
+    // First partial trace, vanishes all coefficients which are not multiples of gap_in
+    // [1, 1, 1, 1, 0, 0, 0, ..., 0, 0, -1, -1, -1, -1] -> [1, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0, 0, 0]
+    res.trace(
+        module,
+        module.log_n() - log_gap_in as usize + 1,
+        log_n,
+        &a,
+        auto_keys,
+        scratch,
+    );
+
+    // TODO: optimize with packing and final partial trace
+    // If gap_out < gap_in, then we need to repack, i.e. reduce the cap between coefficients.
+    if log_gap_in != log_gap_out {
+        let steps: i32 = 1 << log_domain;
+        (0..steps).for_each(|i| {
+            if i != 0 {
+                res.rotate_inplace(module, -(1 << log_gap_in));
+            }
+            cts.insert(i as usize * (1 << log_gap_out), res.clone());
+        });
+
+        let now: Instant = Instant::now();
+        pack(module, &mut cts, log_gap_out, auto_keys, scratch);
+        println!("pack: {} ms", now.elapsed().as_millis());
+        let packed: GLWECiphertext<Vec<u8>> = cts.remove(&0).unwrap();
+        res.trace(
+            module,
+            log_n - log_gap_out,
+            log_n,
+            &packed,
+            auto_keys,
+            scratch,
+        );
+    }
+}
+
+pub fn pack<D: DataMut, B: Backend>(
+    module: &Module<B>,
+    cts: &mut HashMap<usize, GLWECiphertext<D>>,
+    log_gap_out: usize,
+    auto_keys: &HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
+    scratch: &mut Scratch<B>,
+) where
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
+{
+    let log_n: usize = module.log_n();
+
+    let basek: usize = cts.get(&0).unwrap().basek();
+    let k: usize = cts.get(&0).unwrap().k();
+    let rank: usize = cts.get(&0).unwrap().rank();
+
+    (0..log_n - log_gap_out).for_each(|i| {
+        let now: Instant = Instant::now();
+
+        let t = 16.min(1 << (log_n - 1 - i));
+
+        let auto_key: &AutomorphismKeyExec<Vec<u8>, B>;
+        if i == 0 {
+            auto_key = auto_keys.get(&-1).unwrap()
+        } else {
+            auto_key = auto_keys.get(&module.galois_element(1 << (i - 1))).unwrap();
+        }
+
+        (0..t).for_each(|j| {
+            let mut a: Option<GLWECiphertext<D>> = cts.remove(&j);
+            let mut b: Option<GLWECiphertext<D>> = cts.remove(&(j + t));
+
+            combine(
+                module,
+                basek,
+                k,
+                rank,
+                a.as_mut(),
+                b.as_mut(),
+                i,
+                auto_key,
+                scratch,
+            );
+
+            if let Some(a) = a {
+                cts.insert(j, a);
+            } else if let Some(b) = b {
+                cts.insert(j, b);
+            }
+        });
+
+        println!("combine: {} us", now.elapsed().as_micros());
+    });
+}
+
+fn combine<A: DataMut, D: DataMut, DataAK: DataRef, B: Backend>(
+    module: &Module<B>,
+    basek: usize,
+    k: usize,
+    rank: usize,
+    a: Option<&mut GLWECiphertext<A>>,
+    b: Option<&mut GLWECiphertext<D>>,
+    i: usize,
+    auto_key: &AutomorphismKeyExec<DataAK, B>,
+    scratch: &mut Scratch<B>,
+) where
+    Module<B>: CGGICircuitBootstrapFamily<B>,
+    Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
+{
+    // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
+    // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
+    // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
+    // Different cases for wether a and/or b are zero.
+    //
+    // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
+    // Necessary so that the scaling of the plaintext remains constant.
+    // It however is ok to do so here because coefficients are eventually
+    // either mapped to garbage or twice their value which vanishes I(X)
+    // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
+    if let Some(a) = a {
+        let n: usize = a.n();
+        let log_n: usize = (u64::BITS - (n - 1).leading_zeros()) as _;
+        let t: i64 = 1 << (log_n - i - 1);
+
+        if let Some(b) = b {
+            let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
+
+            // a = a * X^-t
+            a.rotate_inplace(module, -t);
+
+            // tmp_b = a * X^-t - b
+            tmp_b.sub(module, a, b);
+            tmp_b.rsh(module, 1);
+
+            // a = a * X^-t + b
+            a.add_inplace(module, b);
+            a.rsh(module, 1);
+
+            tmp_b.normalize_inplace(module, scratch_1);
+
+            // tmp_b = phi(a * X^-t - b)
+            tmp_b.automorphism_inplace(module, auto_key, scratch_1);
+
+            // a = a * X^-t + b - phi(a * X^-t - b)
+            a.sub_inplace_ab(module, &tmp_b);
+            a.normalize_inplace(module, scratch_1);
+
+            // a = a + b * X^t - phi(a * X^-t - b) * X^t
+            //   = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
+            //   = a + b * X^t + phi(a - b * X^t)
+            a.rotate_inplace(module, t);
+        } else {
+            a.rsh(module, 1);
+            // a = a + phi(a)
+            a.automorphism_add_inplace(module, auto_key, scratch);
+        }
+    } else {
+        if let Some(b) = b {
+            let n: usize = b.n();
+            let log_n: usize = (u64::BITS - (n - 1).leading_zeros()) as _;
+            let t: i64 = 1 << (log_n - i - 1);
+
+            let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
+            tmp_b.rotate(module, t, b);
+            tmp_b.rsh(module, 1);
+
+            // a = (b* X^t - phi(b* X^t))
+            b.automorphism_sub_ba(module, &tmp_b, auto_key, scratch_1);
+        }
+    }
+}