poulpy_core/external_product/
gglwe.rs1use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch, ZnxZero};
2
3use crate::{
4 GLWEExternalProduct, ScratchTakeCore,
5 layouts::{
6 GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GGSWInfos, GGSWPrepared, GLWEAutomorphismKey, GLWEInfos, GLWESwitchingKey,
7 prepared::GGSWPreparedToRef,
8 },
9};
10
11impl GLWEAutomorphismKey<Vec<u8>> {
12 pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(
13 &self,
14 module: &M,
15 res_infos: &R,
16 a_infos: &A,
17 b_infos: &B,
18 ) -> usize
19 where
20 R: GGLWEInfos,
21 A: GGLWEInfos,
22 B: GGSWInfos,
23 M: GGLWEExternalProduct<BE>,
24 {
25 module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
26 }
27}
28
29impl<DataSelf: DataMut> GLWEAutomorphismKey<DataSelf> {
30 pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
31 where
32 M: GGLWEExternalProduct<BE>,
33 A: GGLWEToRef + GGLWEInfos,
34 B: GGSWPreparedToRef<BE> + GGSWInfos,
35 Scratch<BE>: ScratchTakeCore<BE>,
36 {
37 module.gglwe_external_product(self, a, b, scratch);
38 }
39
40 pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
41 where
42 M: GGLWEExternalProduct<BE>,
43 A: GGSWPreparedToRef<BE>,
44 Scratch<BE>: ScratchTakeCore<BE>,
45 {
46 module.gglwe_external_product_inplace(self, a, scratch);
47 }
48}
49
50pub trait GGLWEExternalProduct<BE: Backend>
51where
52 Self: GLWEExternalProduct<BE>,
53{
54 fn gglwe_external_product_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
55 where
56 R: GGLWEInfos,
57 A: GGLWEInfos,
58 B: GGSWInfos,
59 {
60 self.glwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
61 }
62
63 fn gglwe_external_product<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
64 where
65 R: GGLWEToMut + GGLWEInfos,
66 A: GGLWEToRef + GGLWEInfos,
67 B: GGSWPreparedToRef<BE> + GGSWInfos,
68 Scratch<BE>: ScratchTakeCore<BE>,
69 {
70 assert_eq!(
71 res.rank_in(),
72 a.rank_in(),
73 "res input rank_in: {} != a input rank_in: {}",
74 res.rank_in(),
75 a.rank_in()
76 );
77 assert_eq!(
78 a.rank_out(),
79 b.rank(),
80 "a output rank_out: {} != b rank: {}",
81 a.rank_out(),
82 b.rank()
83 );
84 assert_eq!(
85 res.rank_out(),
86 b.rank(),
87 "res output rank_out: {} != b rank: {}",
88 res.rank_out(),
89 b.rank()
90 );
91 assert_eq!(res.base2k(), a.base2k());
92
93 let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
94 let a: &GGLWE<&[u8]> = &a.to_ref();
95 let b: &GGSWPrepared<&[u8], BE> = &b.to_ref();
96
97 for row in 0..res.dnum().into() {
98 for col in 0..res.rank_in().into() {
99 self.glwe_external_product(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
100 }
101 }
102
103 for row in res.dnum().min(a.dnum()).into()..res.dnum().into() {
104 for col in 0..res.rank_in().into() {
105 res.at_mut(row, col).data_mut().zero();
106 }
107 }
108 }
109
110 fn gglwe_external_product_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
111 where
112 R: GGLWEToMut,
113 A: GGSWPreparedToRef<BE>,
114 Scratch<BE>: ScratchTakeCore<BE>,
115 {
116 let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
117 let a: &GGSWPrepared<&[u8], BE> = &a.to_ref();
118
119 assert_eq!(
120 res.rank_out(),
121 a.rank(),
122 "res output rank: {} != a rank: {}",
123 res.rank_out(),
124 a.rank()
125 );
126
127 for row in 0..res.dnum().into() {
128 for col in 0..res.rank_in().into() {
129 self.glwe_external_product_inplace(&mut res.at_mut(row, col), a, scratch);
130 }
131 }
132 }
133}
134
135impl<BE: Backend> GGLWEExternalProduct<BE> for Module<BE> where Self: GLWEExternalProduct<BE> {}
136
137impl GLWESwitchingKey<Vec<u8>> {
138 pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
139 where
140 R: GGLWEInfos,
141 A: GGLWEInfos,
142 B: GGSWInfos,
143 M: GGLWEExternalProduct<BE>,
144 {
145 module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
146 }
147}
148
149impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
150 pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
151 where
152 M: GGLWEExternalProduct<BE>,
153 A: GGLWEToRef + GGLWEInfos,
154 B: GGSWPreparedToRef<BE> + GGSWInfos,
155 Scratch<BE>: ScratchTakeCore<BE>,
156 {
157 module.gglwe_external_product(self, a, b, scratch);
158 }
159
160 pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
161 where
162 M: GGLWEExternalProduct<BE>,
163 A: GGSWPreparedToRef<BE>,
164 Scratch<BE>: ScratchTakeCore<BE>,
165 {
166 module.gglwe_external_product_inplace(self, a, scratch);
167 }
168}