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 poly_fma_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)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
52/// element
53///
54/// - rlwe_in: Input ciphertext RLWE_{s(X)}(m(X)).
55/// - ksk: Auto key switching key with polynomials in evaluation domain
56/// - auto_map_index: If automorphism sends i^th coefficient of m(X) to j^th
57/// coefficient of m(X^k) then auto_map_index[i] = j
58/// - auto_sign_index: With a = m(X)[i], if m(X^k)[auto_map_index[i]] = -a, then
59/// auto_sign_index[i] = false, else auto_sign_index[i] = true
60/// - scratch_matrix: must have dimension at-least d+2 x ring_size. `d` rows to
61/// store decomposed polynomials nad 2 rows to store out RLWE temporarily.
62pub(crate) fn rlwe_auto<
63 MT: Matrix + IsTrivial + MatrixMut,
64 Mmut: MatrixMut<MatElement = MT::MatElement>,
65 ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
66 NttOp: Ntt<Element = MT::MatElement>,
67 D: Decomposer<Element = MT::MatElement>,
68>(
69 rlwe_in: &mut MT,
70 ksk: &Mmut,
71 scratch_matrix: &mut Mmut,
72 auto_map_index: &[usize],
73 auto_map_sign: &[bool],
74 mod_op: &ModOp,
75 ntt_op: &NttOp,
76 decomposer: &D,
77) where
78 <Mmut as Matrix>::R: RowMut,
79 <MT as Matrix>::R: RowMut,
80 MT::MatElement: Copy + Zero,
81{
82 let d = decomposer.decomposition_count();
83 let ring_size = rlwe_in.dimension().1;
84 assert!(rlwe_in.dimension().0 == 2);
85 assert!(scratch_matrix.fits(d + 2, ring_size));
86
87 // scratch matrix is guaranteed to have at-least d+2 rows but can have more than
88 // d+2 rows. We require to split them into sub-matrices of exact sizes one with
89 // d rows for storing decomposed polynomial and second with 2 rows to act
90 // tomperary space for RLWE ciphertext. Exact sizes is necessary to avoid any
91 // irrelevant extra FMA or NTT ops.
92 let (scratch_matrix_d_ring, other_half) = scratch_matrix.split_at_row_mut(d);
93 let (tmp_rlwe_out, _) = other_half.split_at_mut(2);
94
95 debug_assert!(tmp_rlwe_out.len() == 2);
96 debug_assert!(scratch_matrix_d_ring.len() == d);
97
98 if !rlwe_in.is_trivial() {
99 tmp_rlwe_out.iter_mut().for_each(|r| {
100 r.as_mut().fill(Mmut::MatElement::zero());
101 });
102
103 // send a(X) -> a(X^k) and decompose a(X^k)
104 izip!(
105 rlwe_in.get_row(0),
106 auto_map_index.iter(),
107 auto_map_sign.iter()
108 )
109 .for_each(|(el_in, to_index, sign)| {
110 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
111
112 decomposer
113 .decompose_iter(&el_out)
114 .enumerate()
115 .for_each(|(index, el)| {
116 scratch_matrix_d_ring[index].as_mut()[*to_index] = el;
117 });
118 });
119
120 // transform decomposed a(X^k) to evaluation domain
121 scratch_matrix_d_ring.iter_mut().for_each(|r| {
122 ntt_op.forward(r.as_mut());
123 });
124
125 // RLWE(m^k) = a', b'; RLWE(m) = a, b
126 // key switch: (a * RLWE'(s(X^k)))
127 let (ksk_a, ksk_b) = ksk.split_at_row(d);
128 // a' = decomp<a> * RLWE'_A(s(X^k))
129 poly_fma_routine(
130 tmp_rlwe_out[0].as_mut(),
131 scratch_matrix_d_ring,
132 ksk_a,
133 mod_op,
134 );
135
136 // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
137 poly_fma_routine(
138 tmp_rlwe_out[1].as_mut(),
139 scratch_matrix_d_ring,
140 ksk_b,
141 mod_op,
142 );
143
144 // transform RLWE(m^k) to coefficient domain
145 tmp_rlwe_out
146 .iter_mut()
147 .for_each(|r| ntt_op.backward(r.as_mut()));
148
149 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
150 izip!(
151 rlwe_in.get_row(1),
152 auto_map_index.iter(),
153 auto_map_sign.iter()
154 )
155 .for_each(|(el_in, to_index, sign)| {
156 let row = tmp_rlwe_out[1].as_mut();
157 if !*sign {
158 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
159 } else {
160 row[*to_index] = mod_op.add(&row[*to_index], el_in);
161 }
162 });
163
164 // copy over A; Leave B for later
165 rlwe_in
166 .get_row_mut(0)
167 .copy_from_slice(tmp_rlwe_out[0].as_ref());
168 } else {
169 // RLWE is trivial, a(X) is 0.
170 // send b(X) -> b(X^k)
171 izip!(
172 rlwe_in.get_row(1),
173 auto_map_index.iter(),
174 auto_map_sign.iter()
175 )
176 .for_each(|(el_in, to_index, sign)| {
177 if !*sign {
178 tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
179 } else {
180 tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
181 }
182 });
183 }
184
185 // Copy over B
186 rlwe_in
187 .get_row_mut(1)
188 .copy_from_slice(tmp_rlwe_out[1].as_ref());
189}
190
191/// Sends RLWE_{s(X)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
192/// element
193///
194/// This is same as `galois_auto` with the difference that alongside `ksk` with
195/// key switching polynomials in evaluation domain, shoup representation,
196/// `ksk_shoup`, of the polynomials in evaluation domain is also supplied.
197pub(crate) fn galois_auto_shoup<
198 Mmut: MatrixMut,
199 ModOp: ArithmeticOps<Element = Mmut::MatElement>
200 // + VectorOps<Element = MT::MatElement>
201 + ShoupMatrixFMA<Mmut::R>,
202 NttOp: Ntt<Element = Mmut::MatElement>,
203 D: Decomposer<Element = Mmut::MatElement>,
204>(
205 rlwe_in: &mut Mmut,
206 ksk: &Mmut,
207 ksk_shoup: &Mmut,
208 scratch_matrix: &mut Mmut,
209 auto_map_index: &[usize],
210 auto_map_sign: &[bool],
211 mod_op: &ModOp,
212 ntt_op: &NttOp,
213 decomposer: &D,
214 is_trivial: bool,
215) where
216 <Mmut as Matrix>::R: RowMut,
217 Mmut::MatElement: Copy + Zero,
218{
219 let d = decomposer.decomposition_count();
220 let ring_size = rlwe_in.dimension().1;
221 assert!(rlwe_in.dimension().0 == 2);
222 assert!(scratch_matrix.fits(d + 2, ring_size));
223
224 let (scratch_matrix_d_ring, other_half) = scratch_matrix.split_at_row_mut(d);
225 let (tmp_rlwe_out, _) = other_half.split_at_mut(2);
226
227 debug_assert!(tmp_rlwe_out.len() == 2);
228 debug_assert!(scratch_matrix_d_ring.len() == d);
229
230 if !is_trivial {
231 tmp_rlwe_out.iter_mut().for_each(|r| {
232 r.as_mut().fill(Mmut::MatElement::zero());
233 });
234
235 // send a(X) -> a(X^k) and decompose a(X^k)
236 izip!(
237 rlwe_in.get_row(0),
238 auto_map_index.iter(),
239 auto_map_sign.iter()
240 )
241 .for_each(|(el_in, to_index, sign)| {
242 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
243
244 decomposer
245 .decompose_iter(&el_out)
246 .enumerate()
247 .for_each(|(index, el)| {
248 scratch_matrix_d_ring[index].as_mut()[*to_index] = el;
249 });
250 });
251
252 // transform decomposed a(X^k) to evaluation domain
253 scratch_matrix_d_ring.iter_mut().for_each(|r| {
254 ntt_op.forward_lazy(r.as_mut());
255 });
256
257 // RLWE(m^k) = a', b'; RLWE(m) = a, b
258 // key switch: (a * RLWE'(s(X^k)))
259 let (ksk_a, ksk_b) = ksk.split_at_row(d);
260 let (ksk_a_shoup, ksk_b_shoup) = ksk_shoup.split_at_row(d);
261 // a' = decomp<a> * RLWE'_A(s(X^k))
262 mod_op.shoup_matrix_fma(
263 tmp_rlwe_out[0].as_mut(),
264 ksk_a,
265 ksk_a_shoup,
266 scratch_matrix_d_ring,
267 );
268
269 // b'= decomp<a(X^k)> * RLWE'_B(s(X^k))
270 mod_op.shoup_matrix_fma(
271 tmp_rlwe_out[1].as_mut(),
272 ksk_b,
273 ksk_b_shoup,
274 scratch_matrix_d_ring,
275 );
276
277 // transform RLWE(m^k) to coefficient domain
278 tmp_rlwe_out
279 .iter_mut()
280 .for_each(|r| ntt_op.backward(r.as_mut()));
281
282 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
283 izip!(
284 rlwe_in.get_row(1),
285 auto_map_index.iter(),
286 auto_map_sign.iter()
287 )
288 .for_each(|(el_in, to_index, sign)| {
289 let row = tmp_rlwe_out[1].as_mut();
290 if !*sign {
291 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
292 } else {
293 row[*to_index] = mod_op.add(&row[*to_index], el_in);
294 }
295 });
296
297 // copy over A; Leave B for later
298 rlwe_in
299 .get_row_mut(0)
300 .copy_from_slice(tmp_rlwe_out[0].as_ref());
301 } else {
302 // RLWE is trivial, a(X) is 0.
303 // send b(X) -> b(X^k)
304 izip!(
305 rlwe_in.get_row(1),
306 auto_map_index.iter(),
307 auto_map_sign.iter()
308 )
309 .for_each(|(el_in, to_index, sign)| {
310 if !*sign {
311 tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
312 } else {
313 tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
314 }
315 });
316 }
317
318 // Copy over B
319 rlwe_in
320 .get_row_mut(1)
321 .copy_from_slice(tmp_rlwe_out[1].as_ref());
322}
323
324/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
325///
326/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
327/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
328/// - scratch_matrix: with dimension (max(d_a, d_b) + 2) x ring_size columns.
329/// It's used to store decomposed polynomials and out RLWE temporarily
330pub(crate) fn rlwe_by_rgsw<
331 Mmut: MatrixMut,
332 MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
333 D: RlweDecomposer<Element = Mmut::MatElement>,
334 ModOp: VectorOps<Element = Mmut::MatElement>,
335 NttOp: Ntt<Element = Mmut::MatElement>,
336>(
337 rlwe_in: &mut MT,
338 rgsw_in: &Mmut,
339 scratch_matrix: &mut Mmut,
340 decomposer: &D,
341 ntt_op: &NttOp,
342 mod_op: &ModOp,
343) where
344 Mmut::MatElement: Copy + Zero,
345 <Mmut as Matrix>::R: RowMut,
346 <MT as Matrix>::R: RowMut,
347{
348 let decomposer_a = decomposer.a();
349 let decomposer_b = decomposer.b();
350 let d_a = decomposer_a.decomposition_count();
351 let d_b = decomposer_b.decomposition_count();
352 let max_d = std::cmp::max(d_a, d_b);
353 assert!(scratch_matrix.fits(max_d + 2, rlwe_in.dimension().1));
354 assert!(rgsw_in.dimension() == (d_a * 2 + d_b * 2, rlwe_in.dimension().1));
355
356 // decomposed RLWE x RGSW
357 let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_a * 2);
358 let (scratch_matrix_d_ring, rest) = scratch_matrix.split_at_row_mut(max_d);
359 let (scratch_rlwe_out, _) = rest.split_at_mut(2);
360
361 scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
362 scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
363
364 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
365 if !rlwe_in.is_trivial() {
366 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
367 // decomp<a_in>
368 decompose_r(
369 rlwe_in.get_row_slice(0),
370 &mut scratch_matrix_d_ring[..d_a],
371 decomposer_a,
372 );
373 scratch_matrix_d_ring
374 .iter_mut()
375 .take(d_a)
376 .for_each(|r| ntt_op.forward(r.as_mut()));
377 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
378 poly_fma_routine(
379 scratch_rlwe_out[0].as_mut(),
380 &scratch_matrix_d_ring[..d_a],
381 &rlwe_dash_nsm[..d_a],
382 mod_op,
383 );
384 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
385 poly_fma_routine(
386 scratch_rlwe_out[1].as_mut(),
387 &scratch_matrix_d_ring[..d_a],
388 &rlwe_dash_nsm[d_a..],
389 mod_op,
390 );
391 }
392 // decomp<b_in>
393 decompose_r(
394 rlwe_in.get_row_slice(1),
395 &mut scratch_matrix_d_ring[..d_b],
396 decomposer_b,
397 );
398 scratch_matrix_d_ring
399 .iter_mut()
400 .take(d_b)
401 .for_each(|r| ntt_op.forward(r.as_mut()));
402 // a_out += decomp<b_in> \cdot RLWE_A'(m)
403 poly_fma_routine(
404 scratch_rlwe_out[0].as_mut(),
405 &scratch_matrix_d_ring[..d_b],
406 &rlwe_dash_m[..d_b],
407 mod_op,
408 );
409 // b_out += decomp<b_in> \cdot RLWE_B'(m)
410 poly_fma_routine(
411 scratch_rlwe_out[1].as_mut(),
412 &scratch_matrix_d_ring[..d_b],
413 &rlwe_dash_m[d_b..],
414 mod_op,
415 );
416
417 // transform rlwe_out to coefficient domain
418 scratch_rlwe_out
419 .iter_mut()
420 .for_each(|r| ntt_op.backward(r.as_mut()));
421
422 rlwe_in
423 .get_row_mut(0)
424 .copy_from_slice(scratch_rlwe_out[0].as_mut());
425 rlwe_in
426 .get_row_mut(1)
427 .copy_from_slice(scratch_rlwe_out[1].as_mut());
428 rlwe_in.set_not_trivial();
429}
430
431/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
432///
433/// Same as `rlwe_by_rgsw` with the difference that alongside `rgsw_in` with
434/// polynomials in evaluation domain, shoup representation of polynomials in
435/// evaluation domain, `rgsw_in_shoup`, is also supplied.
436pub(crate) fn rlwe_by_rgsw_shoup<
437 Mmut: MatrixMut,
438 D: RlweDecomposer<Element = Mmut::MatElement>,
439 ModOp: ShoupMatrixFMA<Mmut::R>,
440 NttOp: Ntt<Element = Mmut::MatElement>,
441>(
442 rlwe_in: &mut Mmut,
443 rgsw_in: &Mmut,
444 rgsw_in_shoup: &Mmut,
445 scratch_matrix: &mut Mmut,
446 decomposer: &D,
447 ntt_op: &NttOp,
448 mod_op: &ModOp,
449 is_trivial: bool,
450) where
451 Mmut::MatElement: Copy + Zero,
452 <Mmut as Matrix>::R: RowMut,
453{
454 let decomposer_a = decomposer.a();
455 let decomposer_b = decomposer.b();
456 let d_a = decomposer_a.decomposition_count();
457 let d_b = decomposer_b.decomposition_count();
458 let max_d = std::cmp::max(d_a, d_b);
459 assert!(scratch_matrix.fits(max_d + 2, rlwe_in.dimension().1));
460 assert!(rgsw_in.dimension() == (d_a * 2 + d_b * 2, rlwe_in.dimension().1));
461 assert!(rgsw_in.dimension() == rgsw_in_shoup.dimension());
462
463 // decomposed RLWE x RGSW
464 let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_a * 2);
465 let (rlwe_dash_nsm_shoup, rlwe_dash_m_shoup) = rgsw_in_shoup.split_at_row(d_a * 2);
466 let (scratch_matrix_d_ring, rest) = scratch_matrix.split_at_row_mut(max_d);
467 let (scratch_rlwe_out, _) = rest.split_at_mut(2);
468
469 scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
470 scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
471
472 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
473 if !is_trivial {
474 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
475 // decomp<a_in>
476 decompose_r(
477 rlwe_in.get_row_slice(0),
478 &mut scratch_matrix_d_ring[..d_a],
479 decomposer_a,
480 );
481 scratch_matrix_d_ring
482 .iter_mut()
483 .take(d_a)
484 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
485
486 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
487 mod_op.shoup_matrix_fma(
488 scratch_rlwe_out[0].as_mut(),
489 &rlwe_dash_nsm[..d_a],
490 &rlwe_dash_nsm_shoup[..d_a],
491 &scratch_matrix_d_ring[..d_a],
492 );
493
494 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
495 mod_op.shoup_matrix_fma(
496 scratch_rlwe_out[1].as_mut(),
497 &rlwe_dash_nsm[d_a..],
498 &rlwe_dash_nsm_shoup[d_a..],
499 &scratch_matrix_d_ring[..d_a],
500 );
501 }
502 {
503 // decomp<b_in>
504 decompose_r(
505 rlwe_in.get_row_slice(1),
506 &mut scratch_matrix_d_ring[..d_b],
507 decomposer_b,
508 );
509 scratch_matrix_d_ring
510 .iter_mut()
511 .take(d_b)
512 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
513
514 // a_out += decomp<b_in> \cdot RLWE_A'(m)
515 mod_op.shoup_matrix_fma(
516 scratch_rlwe_out[0].as_mut(),
517 &rlwe_dash_m[..d_b],
518 &rlwe_dash_m_shoup[..d_b],
519 &scratch_matrix_d_ring[..d_b],
520 );
521
522 // b_out += decomp<b_in> \cdot RLWE_B'(m)
523 mod_op.shoup_matrix_fma(
524 scratch_rlwe_out[1].as_mut(),
525 &rlwe_dash_m[d_b..],
526 &rlwe_dash_m_shoup[d_b..],
527 &scratch_matrix_d_ring[..d_b],
528 );
529 }
530
531 // transform rlwe_out to coefficient domain
532 scratch_rlwe_out
533 .iter_mut()
534 .for_each(|r| ntt_op.backward(r.as_mut()));
535
536 rlwe_in
537 .get_row_mut(0)
538 .copy_from_slice(scratch_rlwe_out[0].as_mut());
539 rlwe_in
540 .get_row_mut(1)
541 .copy_from_slice(scratch_rlwe_out[1].as_mut());
542}
543
544/// Inplace mutates RGSW(m0) to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
545///
546/// RGSW x RGSW product requires multiple RLWE x RGSW products. For example,
547/// Define
548///
549/// RGSW(m0) = [RLWE(-sm), RLWE(\beta -sm), ..., RLWE(\beta^{d-1} -sm)
550/// RLWE(m), RLWE(\beta m), ..., RLWE(\beta^{d-1} m)]
551/// And RGSW(m1)
552///
553/// Then RGSW(m0) x RGSW(m1) equals:
554/// RGSW(m0m1) = [
555/// rlwe_x_rgsw(RLWE(-sm), RGSW(m1)),
556/// ...,
557/// rlwe_x_rgsw(RLWE(\beta^{d-1} -sm), RGSW(m1)),
558/// rlwe_x_rgsw(RLWE(m), RGSW(m1)),
559/// ...,
560/// rlwe_x_rgsw(RLWE(\beta^{d-1} m), RGSW(m1)),
561/// ]
562///
563/// Since noise growth in RLWE x RGSW depends on noise in RGSW ciphertext, it is
564/// clear to observe from above that noise in resulting RGSW(m0m1) equals noise
565/// accumulated in a single RLWE x RGSW and depends on noise in RGSW(m1) (i.e.
566/// rgsw_1_eval)
567///
568/// - rgsw_0: RGSW(m0) in coefficient domain
569/// - rgsw_1_eval: RGSW(m1) in evaluation domain
570pub(crate) fn rgsw_by_rgsw_inplace<
571 Mmut: MatrixMut,
572 D: RlweDecomposer<Element = Mmut::MatElement>,
573 ModOp: VectorOps<Element = Mmut::MatElement>,
574 NttOp: Ntt<Element = Mmut::MatElement>,
575>(
576 rgsw0: &mut Mmut,
577 rgsw0_da: usize,
578 rgsw0_db: usize,
579 rgsw1_eval: &Mmut,
580 decomposer: &D,
581 scratch_matrix: &mut Mmut,
582 ntt_op: &NttOp,
583 mod_op: &ModOp,
584) where
585 <Mmut as Matrix>::R: RowMut,
586 Mmut::MatElement: Copy + Zero,
587{
588 let decomposer_a = decomposer.a();
589 let decomposer_b = decomposer.b();
590 let d_a = decomposer_a.decomposition_count();
591 let d_b = decomposer_b.decomposition_count();
592 let max_d = std::cmp::max(d_a, d_b);
593 let rgsw1_rows = d_a * 2 + d_b * 2;
594 let rgsw0_rows = rgsw0_da * 2 + rgsw0_db * 2;
595 let ring_size = rgsw0.dimension().1;
596 assert!(rgsw0.dimension().0 == rgsw0_rows);
597 assert!(rgsw1_eval.dimension() == (rgsw1_rows, ring_size));
598 assert!(scratch_matrix.fits(max_d + rgsw0_rows, ring_size));
599
600 let (decomp_r_space, rgsw_space) = scratch_matrix.split_at_row_mut(max_d);
601
602 // zero rgsw_space
603 rgsw_space
604 .iter_mut()
605 .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
606 let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(rgsw0_da * 2);
607 let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
608 rlwe_dash_space_nsm.split_at_mut(rgsw0_da);
609 let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) =
610 rlwe_dash_space_m.split_at_mut(rgsw0_db);
611
612 let (rgsw0_nsm, rgsw0_m) = rgsw0.split_at_row(rgsw0_da * 2);
613 let (rgsw1_nsm, rgsw1_m) = rgsw1_eval.split_at_row(d_a * 2);
614
615 // RGSW x RGSW
616 izip!(
617 rgsw0_nsm
618 .iter()
619 .take(rgsw0_da)
620 .chain(rgsw0_m.iter().take(rgsw0_db)),
621 rgsw0_nsm
622 .iter()
623 .skip(rgsw0_da)
624 .chain(rgsw0_m.iter().skip(rgsw0_db)),
625 rlwe_dash_space_nsm_parta
626 .iter_mut()
627 .chain(rlwe_dash_space_m_parta.iter_mut()),
628 rlwe_dash_space_nsm_partb
629 .iter_mut()
630 .chain(rlwe_dash_space_m_partb.iter_mut()),
631 )
632 .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
633 // RLWE(m0) x RGSW(m1)
634
635 // Part A: Decomp<RLWE(m0)[A]> \cdot RLWE'(-sm1)
636 decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer_a);
637 decomp_r_space
638 .iter_mut()
639 .take(d_a)
640 .for_each(|ri| ntt_op.forward(ri.as_mut()));
641 poly_fma_routine(
642 rlwe_out_a.as_mut(),
643 &decomp_r_space[..d_a],
644 &rgsw1_nsm[..d_a],
645 mod_op,
646 );
647 poly_fma_routine(
648 rlwe_out_b.as_mut(),
649 &decomp_r_space[..d_a],
650 &rgsw1_nsm[d_a..],
651 mod_op,
652 );
653
654 // Part B: Decompose<RLWE(m0)[B]> \cdot RLWE'(m1)
655 decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer_b);
656 decomp_r_space
657 .iter_mut()
658 .take(d_b)
659 .for_each(|ri| ntt_op.forward(ri.as_mut()));
660 poly_fma_routine(
661 rlwe_out_a.as_mut(),
662 &decomp_r_space[..d_b],
663 &rgsw1_m[..d_b],
664 mod_op,
665 );
666 poly_fma_routine(
667 rlwe_out_b.as_mut(),
668 &decomp_r_space[..d_b],
669 &rgsw1_m[d_b..],
670 mod_op,
671 );
672 });
673
674 // copy over RGSW(m0m1) into RGSW(m0)
675 izip!(rgsw0.iter_rows_mut(), rgsw_space.iter())
676 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
677
678 // send back to coefficient domain
679 rgsw0
680 .iter_rows_mut()
681 .for_each(|ri| ntt_op.backward(ri.as_mut()));
682}
683
684/// Key switches input RLWE_{s'}(m) -> RLWE_{s}(m)
685///
686/// Let RLWE_{s'}(m) = [a, b] s.t. m+e = b - as'
687///
688/// Given key switchin key Ksk(s' -> s) = RLWE'_{s}(s') = [RLWE_{s}(beta^i s')]
689/// = [a, a*s + e + beta^i s'] for i \in [0,d), key switching computes:
690/// 1. RLWE_{s}(-s'a) = \sum signed_decompose(-a)[i] RLWE_{s}(beta^i s')
691/// 2. RLWE_{s}(m) = (b, 0) + RLWE_{s}(-s'a)
692///
693/// - rlwe_in: Input rlwe ciphertext
694/// - ksk: Key switching key Ksk(s' -> s) with polynomials in evaluation domain
695/// - ksk_shoup: Key switching key Ksk(s' -> s) with polynomials in evaluation
696/// domain in shoup representation
697/// - decomposer: Decomposer used for key switching
698pub(crate) fn rlwe_key_switch<
699 M: MatrixMut + MatrixEntity,
700 ModOp: GetModulus<Element = M::MatElement> + ShoupMatrixFMA<M::R> + VectorOps<Element = M::MatElement>,
701 NttOp: Ntt<Element = M::MatElement>,
702 D: Decomposer<Element = M::MatElement>,
703>(
704 rlwe_in: &M,
705 ksk: &M,
706 ksk_shoup: &M,
707 decomposer: &D,
708 ntt_op: &NttOp,
709 mod_op: &ModOp,
710) -> M
711where
712 <M as Matrix>::R: RowMut + RowEntity,
713 M::MatElement: Copy,
714{
715 let ring_size = rlwe_in.dimension().1;
716 assert!(rlwe_in.dimension().0 == 2);
717 assert!(ksk.dimension() == (decomposer.decomposition_count() * 2, ring_size));
718
719 let mut rlwe_out = M::zeros(2, ring_size);
720
721 let mut tmp = M::zeros(decomposer.decomposition_count(), ring_size);
722 let mut tmp_row = M::R::zeros(ring_size);
723
724 // key switch RLWE part -A
725 // negative A
726 tmp_row.as_mut().copy_from_slice(rlwe_in.get_row_slice(0));
727 mod_op.elwise_neg_mut(tmp_row.as_mut());
728 // decompose -A and send to evaluation domain
729 decompose_r(tmp_row.as_ref(), tmp.as_mut(), decomposer);
730 tmp.iter_rows_mut()
731 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
732
733 // RLWE_s(-A u) = B' + B, A' = (decomp(-A) * Ksk(u -> s)) + (B, 0)
734 let (ksk_part_a, ksk_part_b) = ksk.split_at_row(decomposer.decomposition_count());
735 let (ksk_part_a_shoup, ksk_part_b_shoup) =
736 ksk_shoup.split_at_row(decomposer.decomposition_count());
737 // Part A'
738 mod_op.shoup_matrix_fma(
739 rlwe_out.get_row_mut(0),
740 &ksk_part_a,
741 &ksk_part_a_shoup,
742 tmp.as_ref(),
743 );
744 // Part B'
745 mod_op.shoup_matrix_fma(
746 rlwe_out.get_row_mut(1),
747 &ksk_part_b,
748 &ksk_part_b_shoup,
749 tmp.as_ref(),
750 );
751 // back to coefficient domain
752 rlwe_out
753 .iter_rows_mut()
754 .for_each(|r| ntt_op.backward(r.as_mut()));
755
756 // B' + B
757 mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), rlwe_in.get_row_slice(1));
758
759 rlwe_out
760}