phantom-zone/src/rgsw/runtime.rs
src/rgsw/runtime.rsRaw
1use itertools::izip;
2use num_traits::Zero;
3
4use crate::{
5 backend::{ArithmeticOps, GetModulus, ShoupMatrixFMA, VectorOps},
6 decomposer::{Decomposer, RlweDecomposer},
7 ntt::Ntt,
8 Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut,
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
184pub(crate) fn galois_auto_shoup<
185 MT: Matrix + IsTrivial + MatrixMut,
186 Mmut: MatrixMut<MatElement = MT::MatElement>,
187 ModOp: ArithmeticOps<Element = MT::MatElement>
188 // + VectorOps<Element = MT::MatElement>
189 + ShoupMatrixFMA<Mmut::R>,
190 NttOp: Ntt<Element = MT::MatElement>,
191 D: Decomposer<Element = MT::MatElement>,
192>(
193 rlwe_in: &mut MT,
194 ksk: &Mmut,
195 ksk_shoup: &Mmut,
196 scratch_matrix: &mut Mmut,
197 auto_map_index: &[usize],
198 auto_map_sign: &[bool],
199 mod_op: &ModOp,
200 ntt_op: &NttOp,
201 decomposer: &D,
202) where
203 <Mmut as Matrix>::R: RowMut,
204 <MT as Matrix>::R: RowMut,
205 MT::MatElement: Copy + Zero,
206{
207 let d = decomposer.decomposition_count();
208 let ring_size = rlwe_in.dimension().1;
209 assert!(rlwe_in.dimension().0 == 2);
210 assert!(scratch_matrix.fits(d + 2, ring_size));
211
212 let (scratch_matrix_d_ring, other_half) = scratch_matrix.split_at_row_mut(d);
213 let (tmp_rlwe_out, _) = other_half.split_at_mut(2);
214
215 debug_assert!(tmp_rlwe_out.len() == 2);
216 debug_assert!(scratch_matrix_d_ring.len() == d);
217
218 if !rlwe_in.is_trivial() {
219 tmp_rlwe_out.iter_mut().for_each(|r| {
220 r.as_mut().fill(Mmut::MatElement::zero());
221 });
222
223 // send a(X) -> a(X^k) and decompose a(X^k)
224 izip!(
225 rlwe_in.get_row(0),
226 auto_map_index.iter(),
227 auto_map_sign.iter()
228 )
229 .for_each(|(el_in, to_index, sign)| {
230 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
231
232 decomposer
233 .decompose_iter(&el_out)
234 .enumerate()
235 .for_each(|(index, el)| {
236 scratch_matrix_d_ring[index].as_mut()[*to_index] = el;
237 });
238 });
239
240 // transform decomposed a(X^k) to evaluation domain
241 scratch_matrix_d_ring.iter_mut().for_each(|r| {
242 ntt_op.forward_lazy(r.as_mut());
243 });
244
245 // RLWE(m^k) = a', b'; RLWE(m) = a, b
246 // key switch: (a * RLWE'(s(X^k)))
247 let (ksk_a, ksk_b) = ksk.split_at_row(d);
248 let (ksk_a_shoup, ksk_b_shoup) = ksk_shoup.split_at_row(d);
249 // a' = decomp<a> * RLWE'_A(s(X^k))
250 mod_op.shoup_matrix_fma(
251 tmp_rlwe_out[0].as_mut(),
252 ksk_a,
253 ksk_a_shoup,
254 scratch_matrix_d_ring,
255 );
256
257 // b'= decomp<a(X^k)> * RLWE'_B(s(X^k))
258 mod_op.shoup_matrix_fma(
259 tmp_rlwe_out[1].as_mut(),
260 ksk_b,
261 ksk_b_shoup,
262 scratch_matrix_d_ring,
263 );
264
265 // transform RLWE(m^k) to coefficient domain
266 tmp_rlwe_out
267 .iter_mut()
268 .for_each(|r| ntt_op.backward(r.as_mut()));
269
270 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
271 izip!(
272 rlwe_in.get_row(1),
273 auto_map_index.iter(),
274 auto_map_sign.iter()
275 )
276 .for_each(|(el_in, to_index, sign)| {
277 let row = tmp_rlwe_out[1].as_mut();
278 if !*sign {
279 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
280 } else {
281 row[*to_index] = mod_op.add(&row[*to_index], el_in);
282 }
283 });
284
285 // copy over A; Leave B for later
286 rlwe_in
287 .get_row_mut(0)
288 .copy_from_slice(tmp_rlwe_out[0].as_ref());
289 } else {
290 // RLWE is trivial, a(X) is 0.
291 // send b(X) -> b(X^k)
292 izip!(
293 rlwe_in.get_row(1),
294 auto_map_index.iter(),
295 auto_map_sign.iter()
296 )
297 .for_each(|(el_in, to_index, sign)| {
298 if !*sign {
299 tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
300 } else {
301 tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
302 }
303 });
304 }
305
306 // Copy over B
307 rlwe_in
308 .get_row_mut(1)
309 .copy_from_slice(tmp_rlwe_out[1].as_ref());
310}
311
312/// Returns RLWE(m0m1) = RLWE(m0) x RGSW(m1). Mutates rlwe_in inplace to equal
313/// RLWE(m0m1)
314///
315/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
316/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
317/// - scratch_matrix_d_ring: is a matrix with atleast max(d_a, d_b) rows and
318/// ring_size columns. It's used to store decomposed polynomials and out RLWE
319/// temoporarily
320pub(crate) fn rlwe_by_rgsw<
321 Mmut: MatrixMut,
322 MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
323 D: RlweDecomposer<Element = Mmut::MatElement>,
324 ModOp: VectorOps<Element = Mmut::MatElement>,
325 NttOp: Ntt<Element = Mmut::MatElement>,
326>(
327 rlwe_in: &mut MT,
328 rgsw_in: &Mmut,
329 scratch_matrix: &mut Mmut,
330 decomposer: &D,
331 ntt_op: &NttOp,
332 mod_op: &ModOp,
333) where
334 Mmut::MatElement: Copy + Zero,
335 <Mmut as Matrix>::R: RowMut,
336 <MT as Matrix>::R: RowMut,
337{
338 let decomposer_a = decomposer.a();
339 let decomposer_b = decomposer.b();
340 let d_a = decomposer_a.decomposition_count();
341 let d_b = decomposer_b.decomposition_count();
342 let max_d = std::cmp::max(d_a, d_b);
343 assert!(scratch_matrix.fits(max_d + 2, rlwe_in.dimension().1));
344 assert!(rgsw_in.dimension() == (d_a * 2 + d_b * 2, rlwe_in.dimension().1));
345
346 // decomposed RLWE x RGSW
347 let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_a * 2);
348 let (scratch_matrix_d_ring, rest) = scratch_matrix.split_at_row_mut(max_d);
349 let (scratch_rlwe_out, _) = rest.split_at_mut(2);
350
351 scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
352 scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
353
354 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
355 if !rlwe_in.is_trivial() {
356 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
357 // decomp<a_in>
358 decompose_r(
359 rlwe_in.get_row_slice(0),
360 &mut scratch_matrix_d_ring[..d_a],
361 decomposer_a,
362 );
363 scratch_matrix_d_ring
364 .iter_mut()
365 .take(d_a)
366 .for_each(|r| ntt_op.forward(r.as_mut()));
367 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
368 routine(
369 scratch_rlwe_out[0].as_mut(),
370 &scratch_matrix_d_ring[..d_a],
371 &rlwe_dash_nsm[..d_a],
372 mod_op,
373 );
374 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
375 routine(
376 scratch_rlwe_out[1].as_mut(),
377 &scratch_matrix_d_ring[..d_a],
378 &rlwe_dash_nsm[d_a..],
379 mod_op,
380 );
381 }
382 // decomp<b_in>
383 decompose_r(
384 rlwe_in.get_row_slice(1),
385 &mut scratch_matrix_d_ring[..d_b],
386 decomposer_b,
387 );
388 scratch_matrix_d_ring
389 .iter_mut()
390 .take(d_b)
391 .for_each(|r| ntt_op.forward(r.as_mut()));
392 // a_out += decomp<b_in> \cdot RLWE_A'(m)
393 routine(
394 scratch_rlwe_out[0].as_mut(),
395 &scratch_matrix_d_ring[..d_b],
396 &rlwe_dash_m[..d_b],
397 mod_op,
398 );
399 // b_out += decomp<b_in> \cdot RLWE_B'(m)
400 routine(
401 scratch_rlwe_out[1].as_mut(),
402 &scratch_matrix_d_ring[..d_b],
403 &rlwe_dash_m[d_b..],
404 mod_op,
405 );
406
407 // transform rlwe_out to coefficient domain
408 scratch_rlwe_out
409 .iter_mut()
410 .for_each(|r| ntt_op.backward(r.as_mut()));
411
412 rlwe_in
413 .get_row_mut(0)
414 .copy_from_slice(scratch_rlwe_out[0].as_mut());
415 rlwe_in
416 .get_row_mut(1)
417 .copy_from_slice(scratch_rlwe_out[1].as_mut());
418 rlwe_in.set_not_trivial();
419}
420
421pub(crate) fn rlwe_by_rgsw_shoup<
422 Mmut: MatrixMut,
423 MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
424 D: RlweDecomposer<Element = Mmut::MatElement>,
425 ModOp: ShoupMatrixFMA<Mmut::R>,
426 NttOp: Ntt<Element = Mmut::MatElement>,
427>(
428 rlwe_in: &mut MT,
429 rgsw_in: &Mmut,
430 rgsw_in_shoup: &Mmut,
431 scratch_matrix: &mut Mmut,
432 decomposer: &D,
433 ntt_op: &NttOp,
434 mod_op: &ModOp,
435) where
436 Mmut::MatElement: Copy + Zero,
437 <Mmut as Matrix>::R: RowMut,
438 <MT as Matrix>::R: RowMut,
439{
440 let decomposer_a = decomposer.a();
441 let decomposer_b = decomposer.b();
442 let d_a = decomposer_a.decomposition_count();
443 let d_b = decomposer_b.decomposition_count();
444 let max_d = std::cmp::max(d_a, d_b);
445 assert!(scratch_matrix.fits(max_d + 2, rlwe_in.dimension().1));
446 assert!(rgsw_in.dimension() == (d_a * 2 + d_b * 2, rlwe_in.dimension().1));
447 assert!(rgsw_in.dimension() == rgsw_in_shoup.dimension());
448
449 // decomposed RLWE x RGSW
450 let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_a * 2);
451 let (rlwe_dash_nsm_shoup, rlwe_dash_m_shoup) = rgsw_in_shoup.split_at_row(d_a * 2);
452 let (scratch_matrix_d_ring, rest) = scratch_matrix.split_at_row_mut(max_d);
453 let (scratch_rlwe_out, _) = rest.split_at_mut(2);
454
455 scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
456 scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
457
458 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
459 if !rlwe_in.is_trivial() {
460 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
461 // decomp<a_in>
462 decompose_r(
463 rlwe_in.get_row_slice(0),
464 &mut scratch_matrix_d_ring[..d_a],
465 decomposer_a,
466 );
467 scratch_matrix_d_ring
468 .iter_mut()
469 .take(d_a)
470 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
471
472 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
473 mod_op.shoup_matrix_fma(
474 scratch_rlwe_out[0].as_mut(),
475 &rlwe_dash_nsm[..d_a],
476 &rlwe_dash_nsm_shoup[..d_a],
477 &scratch_matrix_d_ring[..d_a],
478 );
479
480 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
481 mod_op.shoup_matrix_fma(
482 scratch_rlwe_out[1].as_mut(),
483 &rlwe_dash_nsm[d_a..],
484 &rlwe_dash_nsm_shoup[d_a..],
485 &scratch_matrix_d_ring[..d_a],
486 );
487 }
488 {
489 // decomp<b_in>
490 decompose_r(
491 rlwe_in.get_row_slice(1),
492 &mut scratch_matrix_d_ring[..d_b],
493 decomposer_b,
494 );
495 scratch_matrix_d_ring
496 .iter_mut()
497 .take(d_b)
498 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
499
500 // a_out += decomp<b_in> \cdot RLWE_A'(m)
501 mod_op.shoup_matrix_fma(
502 scratch_rlwe_out[0].as_mut(),
503 &rlwe_dash_m[..d_b],
504 &rlwe_dash_m_shoup[..d_b],
505 &scratch_matrix_d_ring[..d_b],
506 );
507
508 // b_out += decomp<b_in> \cdot RLWE_B'(m)
509 mod_op.shoup_matrix_fma(
510 scratch_rlwe_out[1].as_mut(),
511 &rlwe_dash_m[d_b..],
512 &rlwe_dash_m_shoup[d_b..],
513 &scratch_matrix_d_ring[..d_b],
514 );
515 }
516
517 // transform rlwe_out to coefficient domain
518 scratch_rlwe_out
519 .iter_mut()
520 .for_each(|r| ntt_op.backward(r.as_mut()));
521
522 rlwe_in
523 .get_row_mut(0)
524 .copy_from_slice(scratch_rlwe_out[0].as_mut());
525 rlwe_in
526 .get_row_mut(1)
527 .copy_from_slice(scratch_rlwe_out[1].as_mut());
528 rlwe_in.set_not_trivial();
529}
530
531/// Inplace mutates rlwe_0 to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
532/// in evaluation domain
533///
534/// Warning -
535/// Pass a fresh RGSW ciphertext as the second operand, i.e. as `rgsw_1`.
536/// This is to assure minimal error growth in the resulting RGSW ciphertext.
537/// RGSWxRGSW boils down to d_rgsw*2 RLWExRGSW multiplications. Hence, the noise
538/// growth in resulting ciphertext depends on the norm of second RGSW
539/// ciphertext, not the first. This is useful in cases where one is accumulating
540/// multiple RGSW ciphertexts into 1. In which case, pass the accumulating RGSW
541/// ciphertext as rlwe_0 (the one with higher noise) and subsequent RGSW
542/// ciphertexts, with lower noise, to be accumulated as second
543/// operand.
544///
545/// - rgsw_0: RGSW(m0)
546/// - rgsw_1_eval: RGSW(m1) in Evaluation domain
547/// - scratch_matrix_d_plus_rgsw_by_ring: scratch space matrix with rows
548/// (max(d_a, d_b) + d_a*2+d_b*2) and columns ring_size
549///
550/// ## Note:
551/// - We treat RGSW x RGSW as multiple RLWE x RGSW multiplications. .
552pub(crate) fn rgsw_by_rgsw_inplace<
553 Mmut: MatrixMut,
554 D: RlweDecomposer<Element = Mmut::MatElement>,
555 ModOp: VectorOps<Element = Mmut::MatElement>,
556 NttOp: Ntt<Element = Mmut::MatElement>,
557>(
558 rgsw0: &mut Mmut,
559 rgsw0_da: usize,
560 rgsw0_db: usize,
561 rgsw1_eval: &Mmut,
562 decomposer: &D,
563 scratch_matrix: &mut Mmut,
564 ntt_op: &NttOp,
565 mod_op: &ModOp,
566) where
567 <Mmut as Matrix>::R: RowMut,
568 Mmut::MatElement: Copy + Zero,
569{
570 let decomposer_a = decomposer.a();
571 let decomposer_b = decomposer.b();
572 let d_a = decomposer_a.decomposition_count();
573 let d_b = decomposer_b.decomposition_count();
574 let max_d = std::cmp::max(d_a, d_b);
575 let rgsw1_rows = d_a * 2 + d_b * 2;
576 let rgsw0_rows = rgsw0_da * 2 + rgsw0_db * 2;
577 let ring_size = rgsw0.dimension().1;
578 assert!(rgsw0.dimension().0 == rgsw0_rows);
579 assert!(rgsw1_eval.dimension() == (rgsw1_rows, ring_size));
580 assert!(scratch_matrix.fits(max_d + rgsw0_rows, ring_size));
581
582 let (decomp_r_space, rgsw_space) = scratch_matrix.split_at_row_mut(max_d);
583
584 // zero rgsw_space
585 rgsw_space
586 .iter_mut()
587 .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
588 let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(rgsw0_da * 2);
589 let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
590 rlwe_dash_space_nsm.split_at_mut(rgsw0_da);
591 let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) =
592 rlwe_dash_space_m.split_at_mut(rgsw0_db);
593
594 let (rgsw0_nsm, rgsw0_m) = rgsw0.split_at_row(rgsw0_da * 2);
595 let (rgsw1_nsm, rgsw1_m) = rgsw1_eval.split_at_row(d_a * 2);
596
597 // RGSW x RGSW
598 izip!(
599 rgsw0_nsm
600 .iter()
601 .take(rgsw0_da)
602 .chain(rgsw0_m.iter().take(rgsw0_db)),
603 rgsw0_nsm
604 .iter()
605 .skip(rgsw0_da)
606 .chain(rgsw0_m.iter().skip(rgsw0_db)),
607 rlwe_dash_space_nsm_parta
608 .iter_mut()
609 .chain(rlwe_dash_space_m_parta.iter_mut()),
610 rlwe_dash_space_nsm_partb
611 .iter_mut()
612 .chain(rlwe_dash_space_m_partb.iter_mut()),
613 )
614 .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
615 // RLWE(m0) x RGSW(m1)
616
617 // Part A: Decomp<RLWE(m0)[A]> \cdot RLWE'(-sm1)
618 decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer_a);
619 decomp_r_space
620 .iter_mut()
621 .take(d_a)
622 .for_each(|ri| ntt_op.forward(ri.as_mut()));
623 routine(
624 rlwe_out_a.as_mut(),
625 &decomp_r_space[..d_a],
626 &rgsw1_nsm[..d_a],
627 mod_op,
628 );
629 routine(
630 rlwe_out_b.as_mut(),
631 &decomp_r_space[..d_a],
632 &rgsw1_nsm[d_a..],
633 mod_op,
634 );
635
636 // Part B: Decompose<RLWE(m0)[B]> \cdot RLWE'(m1)
637 decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer_b);
638 decomp_r_space
639 .iter_mut()
640 .take(d_b)
641 .for_each(|ri| ntt_op.forward(ri.as_mut()));
642 routine(
643 rlwe_out_a.as_mut(),
644 &decomp_r_space[..d_b],
645 &rgsw1_m[..d_b],
646 mod_op,
647 );
648 routine(
649 rlwe_out_b.as_mut(),
650 &decomp_r_space[..d_b],
651 &rgsw1_m[d_b..],
652 mod_op,
653 );
654 });
655
656 // copy over RGSW(m0m1) into RGSW(m0)
657 izip!(rgsw0.iter_rows_mut(), rgsw_space.iter())
658 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
659
660 // send back to coefficient domain
661 rgsw0
662 .iter_rows_mut()
663 .for_each(|ri| ntt_op.backward(ri.as_mut()));
664}
665
666pub(crate) fn key_switch<
667 M: MatrixMut + MatrixEntity,
668 ModOp: GetModulus<Element = M::MatElement> + ShoupMatrixFMA<M::R> + VectorOps<Element = M::MatElement>,
669 NttOp: Ntt<Element = M::MatElement>,
670 D: Decomposer<Element = M::MatElement>,
671>(
672 rlwe_in: &M,
673 ksk: &M,
674 ksk_shoup: &M,
675 decomposer: &D,
676 ntt_op: &NttOp,
677 mod_op: &ModOp,
678) -> M
679where
680 <M as Matrix>::R: RowMut + RowEntity,
681 M::MatElement: Copy,
682{
683 let ring_size = rlwe_in.dimension().1;
684 assert!(rlwe_in.dimension().0 == 2);
685 assert!(ksk.dimension() == (decomposer.decomposition_count() * 2, ring_size));
686
687 let mut rlwe_out = M::zeros(2, ring_size);
688
689 let mut tmp = M::zeros(decomposer.decomposition_count(), ring_size);
690 let mut tmp_row = M::R::zeros(ring_size);
691
692 // key switch RLWE part -A
693 // negative A
694 tmp_row.as_mut().copy_from_slice(rlwe_in.get_row_slice(0));
695 mod_op.elwise_neg_mut(tmp_row.as_mut());
696 // decompose -A and send to evaluation domain
697 decompose_r(tmp_row.as_ref(), tmp.as_mut(), decomposer);
698 tmp.iter_rows_mut()
699 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
700
701 // RLWE_s(-A u) = B' + B, A' = (decomp(-A) * Ksk(u -> s)) + (B, 0)
702 let (ksk_part_a, ksk_part_b) = ksk.split_at_row(decomposer.decomposition_count());
703 let (ksk_part_a_shoup, ksk_part_b_shoup) =
704 ksk_shoup.split_at_row(decomposer.decomposition_count());
705 // Part A'
706 mod_op.shoup_matrix_fma(
707 rlwe_out.get_row_mut(0),
708 &ksk_part_a,
709 &ksk_part_a_shoup,
710 tmp.as_ref(),
711 );
712 // Part B'
713 mod_op.shoup_matrix_fma(
714 rlwe_out.get_row_mut(1),
715 &ksk_part_b,
716 &ksk_part_b_shoup,
717 tmp.as_ref(),
718 );
719 // back to coefficient domain
720 rlwe_out
721 .iter_rows_mut()
722 .for_each(|r| ntt_op.backward(r.as_mut()));
723
724 // B' + B
725 mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), rlwe_in.get_row_slice(1));
726
727 rlwe_out
728}