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