phantom-zone/src/rgsw/runtime.rs
src/rgsw/runtime.rsRaw
1use itertools::izip;
2use num_traits::Zero;
3
4use crate::{
5 backend::{ArithmeticOps, GetModulus, Modulus, VectorOps},
6 decomposer::{Decomposer, RlweDecomposer},
7 ntt::Ntt,
8 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
9};
10
11use super::IsTrivial;
12
13pub(crate) fn routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
14 write_to_row: &mut [R::Element],
15 matrix_a: &[R],
16 matrix_b: &[R],
17 mod_op: &ModOp,
18) {
19 izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
20 mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
21 });
22}
23
24/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
25/// output matrix decomp_r
26///
27/// Note that decomposition of r(X) requires decomposition of each of
28/// coefficients.
29///
30/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
31/// will be stored at i^th row.
32pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
33 r: &[R::Element],
34 decomp_r: &mut [R],
35 decomposer: &D,
36) where
37 R::Element: Copy,
38{
39 let ring_size = r.len();
40
41 for ri in 0..ring_size {
42 decomposer
43 .decompose_iter(&r[ri])
44 .enumerate()
45 .for_each(|(index, el)| {
46 decomp_r[index].as_mut()[ri] = el;
47 });
48 }
49}
50
51/// Sends RLWE_{s}(X) -> RLWE_{s}(X^k) where k is some galois element
52///
53/// - scratch_matrix: must have dimension at-least d+2 x ring_size. d rows to
54/// store decomposed polynomials and 2 for rlwe
55pub(crate) fn galois_auto<
56 MT: Matrix + IsTrivial + MatrixMut,
57 Mmut: MatrixMut<MatElement = MT::MatElement>,
58 ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
59 NttOp: Ntt<Element = MT::MatElement>,
60 D: Decomposer<Element = MT::MatElement>,
61>(
62 rlwe_in: &mut MT,
63 ksk: &Mmut,
64 scratch_matrix: &mut Mmut,
65 auto_map_index: &[usize],
66 auto_map_sign: &[bool],
67 mod_op: &ModOp,
68 ntt_op: &NttOp,
69 decomposer: &D,
70) where
71 <Mmut as Matrix>::R: RowMut,
72 <MT as Matrix>::R: RowMut,
73 MT::MatElement: Copy + Zero,
74{
75 let d = decomposer.decomposition_count();
76 let ring_size = rlwe_in.dimension().1;
77 assert!(rlwe_in.dimension().0 == 2);
78 assert!(scratch_matrix.fits(d + 2, ring_size));
79
80 // scratch matrix is guaranteed to have at-least d+2 rows but can have more than
81 // d+2 rows. We require to split them into sub-matrices of exact sizes one with
82 // d rows for storing decomposed polynomial and second with 2 rows to act
83 // tomperary space for RLWE ciphertext. Exact sizes is necessary to avoid any
84 // irrelevant extra FMA or NTT ops.
85 let (scratch_matrix_d_ring, other_half) = scratch_matrix.split_at_row_mut(d);
86 let (tmp_rlwe_out, _) = other_half.split_at_mut(2);
87
88 debug_assert!(tmp_rlwe_out.len() == 2);
89 debug_assert!(scratch_matrix_d_ring.len() == d);
90
91 if !rlwe_in.is_trivial() {
92 tmp_rlwe_out.iter_mut().for_each(|r| {
93 r.as_mut().fill(Mmut::MatElement::zero());
94 });
95
96 // send a(X) -> a(X^k) and decompose a(X^k)
97 izip!(
98 rlwe_in.get_row(0),
99 auto_map_index.iter(),
100 auto_map_sign.iter()
101 )
102 .for_each(|(el_in, to_index, sign)| {
103 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
104
105 decomposer
106 .decompose_iter(&el_out)
107 .enumerate()
108 .for_each(|(index, el)| {
109 scratch_matrix_d_ring[index].as_mut()[*to_index] = el;
110 });
111 });
112
113 // transform decomposed a(X^k) to evaluation domain
114 scratch_matrix_d_ring.iter_mut().for_each(|r| {
115 ntt_op.forward(r.as_mut());
116 });
117
118 // RLWE(m^k) = a', b'; RLWE(m) = a, b
119 // key switch: (a * RLWE'(s(X^k)))
120 let (ksk_a, ksk_b) = ksk.split_at_row(d);
121 // a' = decomp<a> * RLWE'_A(s(X^k))
122 routine(
123 tmp_rlwe_out[0].as_mut(),
124 scratch_matrix_d_ring,
125 ksk_a,
126 mod_op,
127 );
128
129 // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
130 routine(
131 tmp_rlwe_out[1].as_mut(),
132 scratch_matrix_d_ring,
133 ksk_b,
134 mod_op,
135 );
136
137 // transform RLWE(m^k) to coefficient domain
138 tmp_rlwe_out
139 .iter_mut()
140 .for_each(|r| ntt_op.backward(r.as_mut()));
141
142 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
143 izip!(
144 rlwe_in.get_row(1),
145 auto_map_index.iter(),
146 auto_map_sign.iter()
147 )
148 .for_each(|(el_in, to_index, sign)| {
149 let row = tmp_rlwe_out[1].as_mut();
150 if !*sign {
151 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
152 } else {
153 row[*to_index] = mod_op.add(&row[*to_index], el_in);
154 }
155 });
156
157 // copy over A; Leave B for later
158 rlwe_in
159 .get_row_mut(0)
160 .copy_from_slice(tmp_rlwe_out[0].as_ref());
161 } else {
162 // RLWE is trivial, a(X) is 0.
163 // send b(X) -> b(X^k)
164 izip!(
165 rlwe_in.get_row(1),
166 auto_map_index.iter(),
167 auto_map_sign.iter()
168 )
169 .for_each(|(el_in, to_index, sign)| {
170 if !*sign {
171 tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
172 } else {
173 tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
174 }
175 });
176 }
177
178 // Copy over B
179 rlwe_in
180 .get_row_mut(1)
181 .copy_from_slice(tmp_rlwe_out[1].as_ref());
182}
183
184/// Returns RLWE(m0m1) = RLWE(m0) x RGSW(m1). Mutates rlwe_in inplace to equal
185/// RLWE(m0m1)
186///
187/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
188/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
189/// - scratch_matrix_d_ring: is a matrix with atleast max(d_a, d_b) rows and
190/// ring_size columns. It's used to store decomposed polynomials and out RLWE
191/// temoporarily
192pub(crate) fn rlwe_by_rgsw<
193 Mmut: MatrixMut,
194 MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
195 D: RlweDecomposer<Element = Mmut::MatElement>,
196 ModOp: VectorOps<Element = Mmut::MatElement>,
197 NttOp: Ntt<Element = Mmut::MatElement>,
198>(
199 rlwe_in: &mut MT,
200 rgsw_in: &Mmut,
201 scratch_matrix: &mut Mmut,
202 decomposer: &D,
203 ntt_op: &NttOp,
204 mod_op: &ModOp,
205) where
206 Mmut::MatElement: Copy + Zero,
207 <Mmut as Matrix>::R: RowMut,
208 <MT as Matrix>::R: RowMut,
209{
210 let decomposer_a = decomposer.a();
211 let decomposer_b = decomposer.b();
212 let d_a = decomposer_a.decomposition_count();
213 let d_b = decomposer_b.decomposition_count();
214 let max_d = std::cmp::max(d_a, d_b);
215 assert!(scratch_matrix.fits(max_d + 2, rlwe_in.dimension().1));
216 assert!(rgsw_in.dimension() == (d_a * 2 + d_b * 2, rlwe_in.dimension().1));
217
218 // decomposed RLWE x RGSW
219 let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_a * 2);
220 let (scratch_matrix_d_ring, scratch_rlwe_out) = scratch_matrix.split_at_row_mut(max_d);
221 scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
222 scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
223 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
224 if !rlwe_in.is_trivial() {
225 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
226 // decomp<a_in>
227 decompose_r(
228 rlwe_in.get_row_slice(0),
229 &mut scratch_matrix_d_ring[..d_a],
230 decomposer_a,
231 );
232 scratch_matrix_d_ring
233 .iter_mut()
234 .take(d_a)
235 .for_each(|r| ntt_op.forward(r.as_mut()));
236 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
237 routine(
238 scratch_rlwe_out[0].as_mut(),
239 scratch_matrix_d_ring.as_ref(),
240 &rlwe_dash_nsm[..d_a],
241 mod_op,
242 );
243 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
244 routine(
245 scratch_rlwe_out[1].as_mut(),
246 scratch_matrix_d_ring.as_ref(),
247 &rlwe_dash_nsm[d_a..],
248 mod_op,
249 );
250 }
251 // decomp<b_in>
252 decompose_r(
253 rlwe_in.get_row_slice(1),
254 &mut scratch_matrix_d_ring[..d_b],
255 decomposer_b,
256 );
257 scratch_matrix_d_ring
258 .iter_mut()
259 .take(d_b)
260 .for_each(|r| ntt_op.forward(r.as_mut()));
261 // a_out += decomp<b_in> \cdot RLWE_A'(m)
262 routine(
263 scratch_rlwe_out[0].as_mut(),
264 scratch_matrix_d_ring.as_ref(),
265 &rlwe_dash_m[..d_b],
266 mod_op,
267 );
268 // b_out += decomp<b_in> \cdot RLWE_B'(m)
269 routine(
270 scratch_rlwe_out[1].as_mut(),
271 scratch_matrix_d_ring.as_ref(),
272 &rlwe_dash_m[d_b..],
273 mod_op,
274 );
275
276 // transform rlwe_out to coefficient domain
277 scratch_rlwe_out
278 .iter_mut()
279 .for_each(|r| ntt_op.backward(r.as_mut()));
280
281 rlwe_in
282 .get_row_mut(0)
283 .copy_from_slice(scratch_rlwe_out[0].as_mut());
284 rlwe_in
285 .get_row_mut(1)
286 .copy_from_slice(scratch_rlwe_out[1].as_mut());
287 rlwe_in.set_not_trivial();
288}
289
290/// Inplace mutates rlwe_0 to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
291/// in evaluation domain
292///
293/// Warning -
294/// Pass a fresh RGSW ciphertext as the second operand, i.e. as `rgsw_1`.
295/// This is to assure minimal error growth in the resulting RGSW ciphertext.
296/// RGSWxRGSW boils down to d_rgsw*2 RLWExRGSW multiplications. Hence, the noise
297/// growth in resulting ciphertext depends on the norm of second RGSW
298/// ciphertext, not the first. This is useful in cases where one is accumulating
299/// multiple RGSW ciphertexts into 1. In which case, pass the accumulating RGSW
300/// ciphertext as rlwe_0 (the one with higher noise) and subsequent RGSW
301/// ciphertexts, with lower noise, to be accumulated as second
302/// operand.
303///
304/// - rgsw_0: RGSW(m0)
305/// - rgsw_1_eval: RGSW(m1) in Evaluation domain
306/// - scratch_matrix_d_plus_rgsw_by_ring: scratch space matrix with rows
307/// (max(d_a, d_b) + d_a*2+d_b*2) and columns ring_size
308pub(crate) fn rgsw_by_rgsw_inplace<
309 Mmut: MatrixMut,
310 D: RlweDecomposer<Element = Mmut::MatElement>,
311 ModOp: VectorOps<Element = Mmut::MatElement>,
312 NttOp: Ntt<Element = Mmut::MatElement>,
313>(
314 rgsw_0: &mut Mmut,
315 rgsw_1_eval: &Mmut,
316 decomposer: &D,
317 scratch_matrix: &mut Mmut,
318 ntt_op: &NttOp,
319 mod_op: &ModOp,
320) where
321 <Mmut as Matrix>::R: RowMut,
322 Mmut::MatElement: Copy + Zero,
323{
324 let decomposer_a = decomposer.a();
325 let decomposer_b = decomposer.b();
326 let d_a = decomposer_a.decomposition_count();
327 let d_b = decomposer_b.decomposition_count();
328 let max_d = std::cmp::max(d_a, d_b);
329 let rgsw_rows = d_a * 2 + d_b * 2;
330 assert!(rgsw_0.dimension().0 == rgsw_rows);
331 let ring_size = rgsw_0.dimension().1;
332 assert!(rgsw_1_eval.dimension() == (rgsw_rows, ring_size));
333 assert!(scratch_matrix.fits(max_d + rgsw_rows, ring_size));
334
335 let (decomp_r_space, rgsw_space) = scratch_matrix.split_at_row_mut(max_d);
336
337 // zero rgsw_space
338 rgsw_space
339 .iter_mut()
340 .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
341 let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(d_a * 2);
342 let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
343 rlwe_dash_space_nsm.split_at_mut(d_a);
344 let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) = rlwe_dash_space_m.split_at_mut(d_b);
345
346 let (rgsw0_nsm, rgsw0_m) = rgsw_0.split_at_row(d_a * 2);
347 let (rgsw1_nsm, rgsw1_m) = rgsw_1_eval.split_at_row(d_a * 2);
348
349 // RGSW x RGSW
350 izip!(
351 rgsw0_nsm.iter().take(d_a).chain(rgsw0_m.iter().take(d_b)),
352 rgsw0_nsm.iter().skip(d_a).chain(rgsw0_m.iter().skip(d_b)),
353 rlwe_dash_space_nsm_parta
354 .iter_mut()
355 .chain(rlwe_dash_space_m_parta.iter_mut()),
356 rlwe_dash_space_nsm_partb
357 .iter_mut()
358 .chain(rlwe_dash_space_m_partb.iter_mut()),
359 )
360 .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
361 // Part A
362 decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer_a);
363 decomp_r_space
364 .iter_mut()
365 .take(d_a)
366 .for_each(|ri| ntt_op.forward(ri.as_mut()));
367 routine(
368 rlwe_out_a.as_mut(),
369 &decomp_r_space[..d_a],
370 &rgsw1_nsm[..d_a],
371 mod_op,
372 );
373 routine(
374 rlwe_out_b.as_mut(),
375 &decomp_r_space[..d_a],
376 &rgsw1_nsm[d_a..],
377 mod_op,
378 );
379
380 // Part B
381 decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer_b);
382 decomp_r_space
383 .iter_mut()
384 .take(d_b)
385 .for_each(|ri| ntt_op.forward(ri.as_mut()));
386 routine(
387 rlwe_out_a.as_mut(),
388 &decomp_r_space[..d_b],
389 &rgsw1_m[..d_b],
390 mod_op,
391 );
392 routine(
393 rlwe_out_b.as_mut(),
394 &decomp_r_space[..d_b],
395 &rgsw1_m[d_b..],
396 mod_op,
397 );
398 });
399
400 // copy over RGSW(m0m1) into RGSW(m0)
401 izip!(rgsw_0.iter_rows_mut(), rgsw_space.iter())
402 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
403
404 // send back to coefficient domain
405 rgsw_0
406 .iter_rows_mut()
407 .for_each(|ri| ntt_op.backward(ri.as_mut()));
408}