Skip to main content

poulpy_ckks/default/
imag.rs

1use anyhow::Result;
2use poulpy_core::{
3    GLWECopy, GLWENegate, GLWERotate, GLWEShift,
4    layouts::{GLWEInfos, GLWEToBackendMut, LWEInfos},
5};
6use poulpy_hal::{
7    api::ModuleN,
8    layouts::{Backend, ScratchArena},
9};
10
11use crate::GLWEToBackendRef;
12use crate::{CKKSInfos, SetCKKSInfos, checked_log_budget_sub, ckks_offset_unary};
13
14pub trait CKKSImagDefault<BE: Backend> {
15    fn ckks_mul_i_tmp_bytes_default(&self) -> usize
16    where
17        Self: GLWERotate<BE> + GLWEShift<BE>,
18    {
19        self.glwe_rotate_tmp_bytes().max(self.glwe_shift_tmp_bytes())
20    }
21
22    fn ckks_div_i_tmp_bytes_default(&self) -> usize
23    where
24        Self: GLWERotate<BE> + GLWEShift<BE>,
25    {
26        self.glwe_rotate_tmp_bytes().max(self.glwe_shift_tmp_bytes())
27    }
28
29    fn ckks_mul_i_into_default<Dst, Src>(&self, dst: &mut Dst, src: &Src, scratch: &mut ScratchArena<'_, BE>) -> Result<()>
30    where
31        Self: GLWERotate<BE> + GLWEShift<BE> + ModuleN,
32        Dst: GLWEToBackendMut<BE> + LWEInfos + CKKSInfos + SetCKKSInfos,
33        Src: GLWEToBackendRef<BE> + GLWEInfos + LWEInfos + CKKSInfos,
34    {
35        let offset = ckks_offset_unary(dst, src);
36        let k = (self.n() / 2) as i64;
37        if offset == 0 {
38            self.glwe_rotate(k, dst, src);
39        } else {
40            self.glwe_lsh(dst, src, offset, scratch);
41            self.glwe_rotate_assign(k, dst, scratch);
42        }
43        dst.set_meta(src.meta());
44        dst.set_log_budget(checked_log_budget_sub("mul_i", src.log_budget(), offset)?);
45        Ok(())
46    }
47
48    fn ckks_mul_i_assign_default<Dst>(&self, dst: &mut Dst, scratch: &mut ScratchArena<'_, BE>) -> Result<()>
49    where
50        Self: GLWERotate<BE> + ModuleN,
51        Dst: GLWEToBackendMut<BE> + LWEInfos + CKKSInfos + SetCKKSInfos,
52    {
53        self.glwe_rotate_assign((self.n() / 2) as i64, dst, scratch);
54        Ok(())
55    }
56
57    fn ckks_div_i_into_default<Dst, Src>(&self, dst: &mut Dst, src: &Src, scratch: &mut ScratchArena<'_, BE>) -> Result<()>
58    where
59        Self: GLWECopy<BE> + GLWENegate<BE> + GLWERotate<BE> + GLWEShift<BE> + ModuleN,
60        Dst: GLWEToBackendMut<BE> + LWEInfos + CKKSInfos + SetCKKSInfos,
61        Src: GLWEToBackendRef<BE> + GLWEInfos + LWEInfos + CKKSInfos,
62    {
63        self.ckks_mul_i_into_default(dst, src, scratch)?;
64        self.glwe_negate_assign(dst);
65        Ok(())
66    }
67
68    fn ckks_div_i_assign_default<Dst>(&self, dst: &mut Dst, scratch: &mut ScratchArena<'_, BE>) -> Result<()>
69    where
70        Self: GLWENegate<BE> + GLWERotate<BE> + ModuleN,
71        Dst: GLWEToBackendMut<BE> + LWEInfos + CKKSInfos + SetCKKSInfos,
72    {
73        self.ckks_mul_i_assign_default(dst, scratch)?;
74        self.glwe_negate_assign(dst);
75        Ok(())
76    }
77}