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 parameters::{DecompositionCount, DoubleDecomposerParams, SingleDecomposerParams},
9 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut,
10};
11
12/// Degree 1 RLWE ciphertext.
13///
14/// RLWE(m) = [a, b] s.t. m+e = b - as
15pub(crate) trait RlweCiphertext {
16 type R: RowMut;
17 /// Returns polynomial `a` of RLWE ciphertext as slice of elements
18 fn part_a(&self) -> &[<Self::R as Row>::Element];
19 /// Returns polynomial `a` of RLWE ciphertext as mutable slice of elements
20 fn part_a_mut(&mut self) -> &mut [<Self::R as Row>::Element];
21 /// Returns polynomial `b` of RLWE ciphertext as slice of elements
22 fn part_b(&self) -> &[<Self::R as Row>::Element];
23 /// Returns polynomial `b` of RLWE ciphertext as mut slice of elements
24 fn part_b_mut(&mut self) -> &mut [<Self::R as Row>::Element];
25 /// Returns ring size of polynomials
26 fn ring_size(&self) -> usize;
27}
28
29/// RGSW ciphertext
30///
31/// RGSW is a collection of RLWE' ciphertext which are collection degree 1 of
32/// RLWE ciphertexts
33///
34/// RGSW = [RLWE'(-sm) || RLWE'(m)]
35///
36/// As usual we refer to decomposition count for RLWE_A in RLWE x RGSW
37/// multiplicaiton as `d_a` and decomposition count for RLWE_B in RLWE x RGDW
38/// multiplication as `d_b`.
39pub(crate) trait RgswCiphertext {
40 type R: Row;
41
42 /// Splits RGSW ciphertext and returns references:
43 /// (RLWE'_A(-sm), RLWE'_B(-sm)), (RLWE'_A(m), RLWE'_B(m))
44 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R]));
45}
46
47pub(crate) trait RgswCiphertextMut: RgswCiphertext {
48 /// Splits RGSW ciphertext and returns mutable references:
49 /// (RLWE'_A(-sm), RLWE'_B(-sm)), (RLWE'_A(m), RLWE'_B(m))
50 fn split_mut(
51 &mut self,
52 ) -> (
53 (&mut [Self::R], &mut [Self::R]),
54 (&mut [Self::R], &mut [Self::R]),
55 );
56}
57
58/// RLWE Key switching Key
59///
60/// Key switching key from s' -> s consists of multiple RLWE cipheretxts.
61/// For gadget vector: [1, beta, ..., beta^{d-1}]
62/// RLWE'_{s}(-s'm) = [RWLE_{s}(-s'm), ..., RLWE_{s}(beta^{d-1} -s'm)]
63pub(crate) trait RlweKsk {
64 type R: Row;
65 /// Returns reference to RLWE'_A(-s'm) polynomials
66 fn ksk_part_a(&self) -> &[Self::R];
67 /// Returns reference to RLWE'_B(-s'm) polynomials
68 fn ksk_part_b(&self) -> &[Self::R];
69}
70
71/// Scratch matrix used in several rlwe/rgsw runtime operations
72pub(crate) trait RuntimeScratchMatrix {
73 type R: RowMut;
74 type Rgsw: RgswCiphertext;
75
76 /// Returns scratch matrix for RLWE automorphism (not trivial case)
77 ///
78 /// RLWE auto requires scratch matric to store decomposed polynomials + 1
79 /// rlwe ciphertext temporarily.
80 ///
81 /// For example, if Auto decomposer has decompostion count `d` then the
82 /// scratch matrix must have dimension (d + 2, N) where N is the ring size.
83 fn scratch_for_rlwe_auto_and_zero_rlwe_space(
84 &mut self,
85 decompostion_count: usize,
86 ) -> (&mut [Self::R], &mut [Self::R]);
87
88 /// Returns scratch matrix for RLWE automorphism (trivial case)
89 ///
90 /// We refer to cases where RLWE(m) = [0, b] s.t. m = b as trivial cases. In
91 /// such a case a single row of length N, N being the ring dimension, is
92 /// required as scratch buffer to store automorphism of polynomial `b`
93 /// temporarily.
94 fn scratch_for_rlwe_auto_trivial_case(&mut self) -> &mut Self::R;
95
96 /// Returns scratch matrix + zeroed RLWE ciphertext space for
97 /// RLWE x RGSW
98 ///
99 /// RLWE x RGSW product requires scratch space to store decomposed
100 /// polynomials for both cases: (1) SignedDecompose(RLWE_A) x RLWE'(-sm) and
101 /// (2) SignedDecompose(RLWE_B) x RLWE'(m). Hence, scratch space returned to
102 /// store decomposed polynomials must have MAX(d_a, d_b) rows.
103 ///
104 /// Additional scratch space is required to store 1 RLWE ciphertext
105 /// temporarily. The space must be zeroed.
106 fn scratch_for_rlwe_x_rgsw_and_zero_rlwe_space<D: RlweDecomposer>(
107 &mut self,
108 decomposer: &D,
109 ) -> (&mut [Self::R], &mut [Self::R]);
110
111 /// Returns scracth matrix + zeroed RGSW ciphertext space for RGSW0 x RGSW1
112 ///
113 /// RGSW0 x RGSW1 requires `d_{0,a} + d_{0,b}` RLWE x RGSW1 products where
114 /// d_{0, a/b} are decomposition counts corresponding to decmposer used for
115 /// RGSW0. Hence, scratch space required to store decomposed polynomial for
116 /// RLWE x RGSW1 product should have MAX(d_{1, a}, d_{1, b}) rows.
117 ///
118 /// Additional scravth space is required to store RGSW0 ciphertext
119 /// temporarily. The space must be zeroed.
120 fn scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space<D: RlweDecomposer>(
121 &mut self,
122 d0: &D,
123 d1: &D,
124 ) -> (&mut [Self::R], &mut [Self::R]);
125}
126
127pub(crate) struct RlweCiphertextMutRef<'a, R> {
128 data: &'a mut [R],
129}
130
131impl<'a, R> RlweCiphertextMutRef<'a, R> {
132 pub(crate) fn new(data: &'a mut [R]) -> Self {
133 Self { data }
134 }
135}
136
137impl<'a, R: RowMut> RlweCiphertext for RlweCiphertextMutRef<'a, R> {
138 type R = R;
139 fn part_a(&self) -> &[<Self::R as Row>::Element] {
140 self.data[0].as_ref()
141 }
142 fn part_a_mut(&mut self) -> &mut [<Self::R as Row>::Element] {
143 self.data[0].as_mut()
144 }
145 fn part_b(&self) -> &[<Self::R as Row>::Element] {
146 self.data[1].as_ref()
147 }
148 fn part_b_mut(&mut self) -> &mut [<Self::R as Row>::Element] {
149 self.data[1].as_mut()
150 }
151 fn ring_size(&self) -> usize {
152 self.data[0].as_ref().len()
153 }
154}
155
156pub(crate) struct RgswCiphertextRef<'a, R> {
157 data: &'a [R],
158 d_a: usize,
159 d_b: usize,
160}
161
162impl<'a, R> RgswCiphertextRef<'a, R> {
163 pub(crate) fn new(data: &'a [R], d_a: usize, d_b: usize) -> Self {
164 RgswCiphertextRef { data, d_a, d_b }
165 }
166}
167
168impl<'a, R> RgswCiphertext for RgswCiphertextRef<'a, R>
169where
170 R: Row,
171{
172 type R = R;
173
174 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R])) {
175 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at(self.d_a * 2);
176 (
177 rlwe_dash_nsm.split_at(self.d_a),
178 rlwe_dash_m.split_at(self.d_b),
179 )
180 }
181}
182
183pub(crate) struct RgswCiphertextMutRef<'a, R> {
184 data: &'a mut [R],
185 d_a: usize,
186 d_b: usize,
187}
188
189impl<'a, R> RgswCiphertextMutRef<'a, R> {
190 pub(crate) fn new(data: &'a mut [R], d_a: usize, d_b: usize) -> Self {
191 RgswCiphertextMutRef { data, d_a, d_b }
192 }
193}
194
195impl<'a, R: RowMut> AsMut<[R]> for RgswCiphertextMutRef<'a, R> {
196 fn as_mut(&mut self) -> &mut [R] {
197 &mut self.data
198 }
199}
200
201impl<'a, R> RgswCiphertext for RgswCiphertextMutRef<'a, R>
202where
203 R: Row,
204{
205 type R = R;
206
207 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R])) {
208 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at(self.d_a * 2);
209 (
210 rlwe_dash_nsm.split_at(self.d_a),
211 rlwe_dash_m.split_at(self.d_b),
212 )
213 }
214}
215
216impl<'a, R> RgswCiphertextMut for RgswCiphertextMutRef<'a, R>
217where
218 R: RowMut,
219{
220 fn split_mut(
221 &mut self,
222 ) -> (
223 (&mut [Self::R], &mut [Self::R]),
224 (&mut [Self::R], &mut [Self::R]),
225 ) {
226 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at_mut(self.d_a * 2);
227 (
228 rlwe_dash_nsm.split_at_mut(self.d_a),
229 rlwe_dash_m.split_at_mut(self.d_b),
230 )
231 }
232}
233
234pub(crate) struct RlweKskRef<'a, R> {
235 data: &'a [R],
236 decomposition_count: usize,
237}
238impl<'a, R: Row> RlweKskRef<'a, R> {
239 pub(crate) fn new(ksk: &'a [R], decomposition_count: usize) -> Self {
240 Self {
241 data: ksk,
242 decomposition_count,
243 }
244 }
245}
246
247impl<'a, R: Row> RlweKsk for RlweKskRef<'a, R> {
248 type R = R;
249
250 fn ksk_part_a(&self) -> &[Self::R] {
251 &self.data[..self.decomposition_count]
252 }
253
254 fn ksk_part_b(&self) -> &[Self::R] {
255 &self.data[self.decomposition_count..]
256 }
257}
258
259pub(crate) struct RuntimeScratchMutRef<'a, R> {
260 data: &'a mut [R],
261}
262
263impl<'a, R> RuntimeScratchMutRef<'a, R> {
264 pub(crate) fn new(data: &'a mut [R]) -> Self {
265 Self { data }
266 }
267}
268
269impl<'a, R: RowMut> RuntimeScratchMatrix for RuntimeScratchMutRef<'a, R>
270where
271 R::Element: Zero + Clone,
272{
273 type R = R;
274 type Rgsw = RgswCiphertextRef<'a, R>;
275
276 fn scratch_for_rlwe_auto_and_zero_rlwe_space(
277 &mut self,
278 decompostion_count: usize,
279 ) -> (&mut [Self::R], &mut [Self::R]) {
280 let (decomp_poly, other) = self.data.split_at_mut(decompostion_count);
281 let (rlwe, _) = other.split_at_mut(2);
282
283 // zero fill rlwe
284 rlwe.iter_mut()
285 .for_each(|r| r.as_mut().fill(R::Element::zero()));
286
287 (decomp_poly, rlwe)
288 }
289
290 fn scratch_for_rlwe_auto_trivial_case(&mut self) -> &mut Self::R {
291 &mut self.data[0]
292 }
293
294 fn scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space<D: RlweDecomposer>(
295 &mut self,
296 rgsw0_decoposer: &D,
297 rgsw1_decoposer: &D,
298 ) -> (&mut [Self::R], &mut [Self::R]) {
299 let (decomp_poly, other) = self.data.split_at_mut(std::cmp::max(
300 rgsw1_decoposer.decomposition_count_a().0,
301 rgsw1_decoposer.decomposition_count_b().0,
302 ));
303 let (rgsw, _) = other.split_at_mut(
304 rgsw0_decoposer.decomposition_count_a().0 * 2
305 + rgsw0_decoposer.decomposition_count_b().0 * 2,
306 );
307
308 // zero fill rgsw0
309 rgsw.iter_mut()
310 .for_each(|r| r.as_mut().fill(R::Element::zero()));
311
312 (decomp_poly, rgsw)
313 }
314
315 fn scratch_for_rlwe_x_rgsw_and_zero_rlwe_space<D: RlweDecomposer>(
316 &mut self,
317 decomposer: &D,
318 ) -> (&mut [Self::R], &mut [Self::R]) {
319 let (decomp_poly, other) = self.data.split_at_mut(std::cmp::max(
320 decomposer.decomposition_count_a().0,
321 decomposer.decomposition_count_b().0,
322 ));
323
324 let (rlwe, _) = other.split_at_mut(2);
325
326 // zero fill rlwe
327 rlwe.iter_mut()
328 .for_each(|r| r.as_mut().fill(R::Element::zero()));
329
330 (decomp_poly, rlwe)
331 }
332}
333
334/// Returns no. of rows in scratch space for RGSW0 x RGSW1 product
335pub(crate) fn rgsw_x_rgsw_scratch_rows<D: DoubleDecomposerParams<Count = DecompositionCount>>(
336 rgsw0_decomposer_param: &D,
337 rgsw1_decomposer_param: &D,
338) -> usize {
339 std::cmp::max(
340 rgsw1_decomposer_param.decomposition_count_a().0,
341 rgsw1_decomposer_param.decomposition_count_b().0,
342 ) + rgsw0_decomposer_param.decomposition_count_a().0 * 2
343 + rgsw0_decomposer_param.decomposition_count_b().0 * 2
344}
345
346/// Returns no. of rows in scratch space for RLWE x RGSW product
347pub(crate) fn rlwe_x_rgsw_scratch_rows<D: DoubleDecomposerParams<Count = DecompositionCount>>(
348 rgsw_decomposer_param: &D,
349) -> usize {
350 std::cmp::max(
351 rgsw_decomposer_param.decomposition_count_a().0,
352 rgsw_decomposer_param.decomposition_count_b().0,
353 ) + 2
354}
355
356/// Returns no. of rows in scratch space for RLWE auto
357pub(crate) fn rlwe_auto_scratch_rows<D: SingleDecomposerParams<Count = DecompositionCount>>(
358 param: &D,
359) -> usize {
360 param.decomposition_count().0 + 2
361}
362
363pub(crate) fn poly_fma_routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
364 write_to_row: &mut [R::Element],
365 matrix_a: &[R],
366 matrix_b: &[R],
367 mod_op: &ModOp,
368) {
369 izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
370 mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
371 });
372}
373
374/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
375/// output matrix decomp_r
376///
377/// Note that decomposition of r(X) requires decomposition of each of
378/// coefficients.
379///
380/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
381/// will be stored at i^th row.
382pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
383 r: &[R::Element],
384 decomp_r: &mut [R],
385 decomposer: &D,
386) where
387 R::Element: Copy,
388{
389 let ring_size = r.len();
390
391 for ri in 0..ring_size {
392 decomposer
393 .decompose_iter(&r[ri])
394 .enumerate()
395 .for_each(|(index, el)| {
396 decomp_r[index].as_mut()[ri] = el;
397 });
398 }
399}
400
401/// Sends RLWE_{s(X)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
402/// element
403///
404/// - rlwe_in: Input ciphertext RLWE_{s(X)}(m(X)).
405/// - ksk: Auto key switching key with polynomials in evaluation domain
406/// - auto_map_index: If automorphism sends i^th coefficient of m(X) to j^th
407/// coefficient of m(X^k) then auto_map_index[i] = j
408/// - auto_sign_index: With a = m(X)[i], if m(X^k)[auto_map_index[i]] = -a, then
409/// auto_sign_index[i] = false, else auto_sign_index[i] = true
410/// - scratch_matrix: must have dimension at-least d+2 x ring_size. `d` rows to
411/// store decomposed polynomials nad 2 rows to store out RLWE temporarily.
412pub(crate) fn rlwe_auto<
413 Rlwe: RlweCiphertext,
414 Ksk: RlweKsk<R = Rlwe::R>,
415 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
416 ModOp: ArithmeticOps<Element = <Rlwe::R as Row>::Element>
417 + VectorOps<Element = <Rlwe::R as Row>::Element>,
418 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
419 D: Decomposer<Element = <Rlwe::R as Row>::Element>,
420>(
421 rlwe_in: &mut Rlwe,
422 ksk: &Ksk,
423 scratch_matrix: &mut Sc,
424 auto_map_index: &[usize],
425 auto_map_sign: &[bool],
426 mod_op: &ModOp,
427 ntt_op: &NttOp,
428 decomposer: &D,
429 is_trivial: bool,
430) where
431 <Rlwe::R as Row>::Element: Copy + Zero,
432{
433 // let ring_size = rlwe_in.dimension().1;
434 // assert!(rlwe_in.dimension().0 == 2);
435 // assert!(scratch_matrix.fits(d + 2, ring_size));
436
437 if !is_trivial {
438 let (decomp_poly_scratch, tmp_rlwe) = scratch_matrix
439 .scratch_for_rlwe_auto_and_zero_rlwe_space(decomposer.decomposition_count().0);
440 let mut tmp_rlwe = RlweCiphertextMutRef::new(tmp_rlwe);
441
442 // send a(X) -> a(X^k) and decompose a(X^k)
443 izip!(
444 rlwe_in.part_a(),
445 auto_map_index.iter(),
446 auto_map_sign.iter()
447 )
448 .for_each(|(el_in, to_index, sign)| {
449 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
450
451 decomposer
452 .decompose_iter(&el_out)
453 .enumerate()
454 .for_each(|(index, el)| {
455 decomp_poly_scratch[index].as_mut()[*to_index] = el;
456 });
457 });
458
459 // transform decomposed a(X^k) to evaluation domain
460 decomp_poly_scratch.iter_mut().for_each(|r| {
461 ntt_op.forward(r.as_mut());
462 });
463
464 // RLWE(m^k) = a', b'; RLWE(m) = a, b
465 // key switch: (a * RLWE'(s(X^k)))
466 // a' = decomp<a> * RLWE'_A(s(X^k))
467 poly_fma_routine(
468 tmp_rlwe.part_a_mut(),
469 decomp_poly_scratch,
470 ksk.ksk_part_a(),
471 mod_op,
472 );
473
474 // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
475 poly_fma_routine(
476 tmp_rlwe.part_b_mut(),
477 decomp_poly_scratch,
478 ksk.ksk_part_b(),
479 mod_op,
480 );
481
482 // transform RLWE(m^k) to coefficient domain
483 ntt_op.backward(tmp_rlwe.part_a_mut());
484 ntt_op.backward(tmp_rlwe.part_b_mut());
485
486 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
487 izip!(
488 rlwe_in.part_b(),
489 auto_map_index.iter(),
490 auto_map_sign.iter()
491 )
492 .for_each(|(el_in, to_index, sign)| {
493 let row = tmp_rlwe.part_b_mut();
494 if !*sign {
495 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
496 } else {
497 row[*to_index] = mod_op.add(&row[*to_index], el_in);
498 }
499 });
500
501 // copy over A; Leave B for later
502 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe.part_a());
503 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe.part_b());
504 } else {
505 // RLWE is trivial, a(X) is 0.
506 // send b(X) -> b(X^k)
507 let tmp_row = scratch_matrix.scratch_for_rlwe_auto_trivial_case();
508 izip!(
509 rlwe_in.part_b(),
510 auto_map_index.iter(),
511 auto_map_sign.iter()
512 )
513 .for_each(|(el_in, to_index, sign)| {
514 if !*sign {
515 tmp_row.as_mut()[*to_index] = mod_op.neg(el_in);
516 } else {
517 tmp_row.as_mut()[*to_index] = *el_in;
518 }
519 });
520 rlwe_in.part_b_mut().copy_from_slice(tmp_row.as_ref());
521 }
522}
523
524/// Sends RLWE_{s(X)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
525/// element
526///
527/// This is same as `galois_auto` with the difference that alongside `ksk` with
528/// key switching polynomials in evaluation domain, shoup representation,
529/// `ksk_shoup`, of the polynomials in evaluation domain is also supplied.
530pub(crate) fn rlwe_auto_shoup<
531 Rlwe: RlweCiphertext,
532 Ksk: RlweKsk<R = Rlwe::R>,
533 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
534 ModOp: ArithmeticOps<Element = <Rlwe::R as Row>::Element>
535 // + VectorOps<Element = MT::MatElement>
536 + ShoupMatrixFMA<Rlwe::R>,
537 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
538 D: Decomposer<Element = <Rlwe::R as Row>::Element>,
539>(
540 rlwe_in: &mut Rlwe,
541 ksk: &Ksk,
542 ksk_shoup: &Ksk,
543 scratch_matrix: &mut Sc,
544 auto_map_index: &[usize],
545 auto_map_sign: &[bool],
546 mod_op: &ModOp,
547 ntt_op: &NttOp,
548 decomposer: &D,
549 is_trivial: bool,
550) where
551 <Rlwe::R as Row>::Element: Copy + Zero,
552{
553 // let d = decomposer.decomposition_count();
554 // let ring_size = rlwe_in.dimension().1;
555 // assert!(rlwe_in.dimension().0 == 2);
556 // assert!(scratch_matrix.fits(d + 2, ring_size));
557
558 if !is_trivial {
559 let (decomp_poly_scratch, tmp_rlwe) = scratch_matrix
560 .scratch_for_rlwe_auto_and_zero_rlwe_space(decomposer.decomposition_count().0);
561 let mut tmp_rlwe = RlweCiphertextMutRef::new(tmp_rlwe);
562
563 // send a(X) -> a(X^k) and decompose a(X^k)
564 izip!(
565 rlwe_in.part_a(),
566 auto_map_index.iter(),
567 auto_map_sign.iter()
568 )
569 .for_each(|(el_in, to_index, sign)| {
570 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
571
572 decomposer
573 .decompose_iter(&el_out)
574 .enumerate()
575 .for_each(|(index, el)| {
576 decomp_poly_scratch[index].as_mut()[*to_index] = el;
577 });
578 });
579
580 // transform decomposed a(X^k) to evaluation domain
581 decomp_poly_scratch.iter_mut().for_each(|r| {
582 ntt_op.forward_lazy(r.as_mut());
583 });
584
585 // RLWE(m^k) = a', b'; RLWE(m) = a, b
586 // key switch: (a * RLWE'(s(X^k)))
587 // a' = decomp<a> * RLWE'_A(s(X^k))
588 mod_op.shoup_matrix_fma(
589 tmp_rlwe.part_a_mut(),
590 ksk.ksk_part_a(),
591 ksk_shoup.ksk_part_a(),
592 decomp_poly_scratch,
593 );
594
595 // b'= decomp<a(X^k)> * RLWE'_B(s(X^k))
596 mod_op.shoup_matrix_fma(
597 tmp_rlwe.part_b_mut(),
598 ksk.ksk_part_b(),
599 ksk_shoup.ksk_part_b(),
600 decomp_poly_scratch,
601 );
602
603 // transform RLWE(m^k) to coefficient domain
604 ntt_op.backward(tmp_rlwe.part_a_mut());
605 ntt_op.backward(tmp_rlwe.part_b_mut());
606
607 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
608 let row = tmp_rlwe.part_b_mut();
609 izip!(
610 rlwe_in.part_b(),
611 auto_map_index.iter(),
612 auto_map_sign.iter()
613 )
614 .for_each(|(el_in, to_index, sign)| {
615 if !*sign {
616 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
617 } else {
618 row[*to_index] = mod_op.add(&row[*to_index], el_in);
619 }
620 });
621
622 // copy over A, B
623 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe.part_a());
624 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe.part_b());
625 } else {
626 // RLWE is trivial, a(X) is 0.
627 // send b(X) -> b(X^k)
628 let row = scratch_matrix.scratch_for_rlwe_auto_trivial_case();
629 izip!(
630 rlwe_in.part_b(),
631 auto_map_index.iter(),
632 auto_map_sign.iter()
633 )
634 .for_each(|(el_in, to_index, sign)| {
635 if !*sign {
636 row.as_mut()[*to_index] = mod_op.neg(el_in);
637 } else {
638 row.as_mut()[*to_index] = *el_in;
639 }
640 });
641 rlwe_in.part_b_mut().copy_from_slice(row.as_ref());
642 }
643}
644
645/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
646///
647/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
648/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
649/// - scratch_matrix: with dimension (max(d_a, d_b) + 2) x ring_size columns.
650/// It's used to store decomposed polynomials and out RLWE temporarily
651pub(crate) fn rlwe_by_rgsw<
652 Rlwe: RlweCiphertext,
653 Rgsw: RgswCiphertext<R = Rlwe::R>,
654 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
655 D: RlweDecomposer<Element = <Rlwe::R as Row>::Element>,
656 ModOp: VectorOps<Element = <Rlwe::R as Row>::Element>,
657 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
658>(
659 rlwe_in: &mut Rlwe,
660 rgsw_in: &Rgsw,
661 scratch_matrix: &mut Sc,
662 decomposer: &D,
663 ntt_op: &NttOp,
664 mod_op: &ModOp,
665 is_trivial: bool,
666) where
667 <Rlwe::R as Row>::Element: Copy + Zero,
668{
669 let decomposer_a = decomposer.a();
670 let decomposer_b = decomposer.b();
671 let d_a = decomposer.decomposition_count_a().0;
672 let d_b = decomposer.decomposition_count_b().0;
673
674 let ((rlwe_dash_nsm_parta, rlwe_dash_nsm_partb), (rlwe_dash_m_parta, rlwe_dash_m_partb)) =
675 rgsw_in.split();
676
677 let (decomposed_poly_scratch, tmp_rlwe) =
678 scratch_matrix.scratch_for_rlwe_x_rgsw_and_zero_rlwe_space(decomposer);
679
680 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
681 if !is_trivial {
682 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
683 // decomp<a_in>
684 let mut decomposed_polys_of_rlwea = &mut decomposed_poly_scratch[..d_a];
685 decompose_r(
686 rlwe_in.part_a(),
687 &mut decomposed_polys_of_rlwea,
688 decomposer_a,
689 );
690
691 decomposed_polys_of_rlwea
692 .iter_mut()
693 .for_each(|r| ntt_op.forward(r.as_mut()));
694
695 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
696 poly_fma_routine(
697 tmp_rlwe[0].as_mut(),
698 &decomposed_polys_of_rlwea,
699 rlwe_dash_nsm_parta,
700 mod_op,
701 );
702 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
703 poly_fma_routine(
704 tmp_rlwe[1].as_mut(),
705 &decomposed_polys_of_rlwea,
706 &rlwe_dash_nsm_partb,
707 mod_op,
708 );
709 }
710
711 {
712 // decomp<b_in>
713 let mut decomposed_polys_of_rlweb = &mut decomposed_poly_scratch[..d_b];
714 decompose_r(
715 rlwe_in.part_b(),
716 &mut decomposed_polys_of_rlweb,
717 decomposer_b,
718 );
719
720 decomposed_polys_of_rlweb
721 .iter_mut()
722 .for_each(|r| ntt_op.forward(r.as_mut()));
723
724 // a_out += decomp<b_in> \cdot RLWE_A'(m)
725 poly_fma_routine(
726 tmp_rlwe[0].as_mut(),
727 &decomposed_polys_of_rlweb,
728 &rlwe_dash_m_parta,
729 mod_op,
730 );
731 // b_out += decomp<b_in> \cdot RLWE_B'(m)
732 poly_fma_routine(
733 tmp_rlwe[1].as_mut(),
734 &decomposed_polys_of_rlweb,
735 &rlwe_dash_m_partb,
736 mod_op,
737 );
738 }
739
740 // transform rlwe_out to coefficient domain
741 tmp_rlwe
742 .iter_mut()
743 .for_each(|r| ntt_op.backward(r.as_mut()));
744
745 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe[0].as_mut());
746 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe[1].as_mut());
747}
748
749/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
750///
751/// Same as `rlwe_by_rgsw` with the difference that alongside `rgsw_in` with
752/// polynomials in evaluation domain, shoup representation of polynomials in
753/// evaluation domain, `rgsw_in_shoup`, is also supplied.
754pub(crate) fn rlwe_by_rgsw_shoup<
755 Rlwe: RlweCiphertext,
756 Rgsw: RgswCiphertext<R = Rlwe::R>,
757 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
758 D: RlweDecomposer<Element = <Rlwe::R as Row>::Element>,
759 ModOp: ShoupMatrixFMA<Rlwe::R>,
760 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
761>(
762 rlwe_in: &mut Rlwe,
763 rgsw_in: &Rgsw,
764 rgsw_in_shoup: &Rgsw,
765 scratch_matrix: &mut Sc,
766 decomposer: &D,
767 ntt_op: &NttOp,
768 mod_op: &ModOp,
769 is_trivial: bool,
770) where
771 <Rlwe::R as Row>::Element: Copy + Zero,
772{
773 let decomposer_a = decomposer.a();
774 let decomposer_b = decomposer.b();
775 let d_a = decomposer.decomposition_count_a().0;
776 let d_b = decomposer.decomposition_count_b().0;
777
778 let ((rlwe_dash_nsm_parta, rlwe_dash_nsm_partb), (rlwe_dash_m_parta, rlwe_dash_m_partb)) =
779 rgsw_in.split();
780
781 let (
782 (rlwe_dash_nsm_parta_shoup, rlwe_dash_nsm_partb_shoup),
783 (rlwe_dash_m_parta_shoup, rlwe_dash_m_partb_shoup),
784 ) = rgsw_in_shoup.split();
785
786 let (decomposed_poly_scratch, tmp_rlwe) =
787 scratch_matrix.scratch_for_rlwe_x_rgsw_and_zero_rlwe_space(decomposer);
788
789 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
790 if !is_trivial {
791 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
792 // decomp<a_in>
793 let mut decomposed_polys_of_rlwea = &mut decomposed_poly_scratch[..d_a];
794 decompose_r(
795 rlwe_in.part_a(),
796 &mut decomposed_polys_of_rlwea,
797 decomposer_a,
798 );
799 decomposed_polys_of_rlwea
800 .iter_mut()
801 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
802
803 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
804 mod_op.shoup_matrix_fma(
805 tmp_rlwe[0].as_mut(),
806 &rlwe_dash_nsm_parta,
807 &rlwe_dash_nsm_parta_shoup,
808 &decomposed_polys_of_rlwea,
809 );
810
811 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
812 mod_op.shoup_matrix_fma(
813 tmp_rlwe[1].as_mut(),
814 &rlwe_dash_nsm_partb,
815 &rlwe_dash_nsm_partb_shoup,
816 &decomposed_polys_of_rlwea,
817 );
818 }
819 {
820 // decomp<b_in>
821 let mut decomposed_polys_of_rlweb = &mut decomposed_poly_scratch[..d_b];
822 decompose_r(
823 rlwe_in.part_b(),
824 &mut decomposed_polys_of_rlweb,
825 decomposer_b,
826 );
827 decomposed_polys_of_rlweb
828 .iter_mut()
829 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
830
831 // a_out += decomp<b_in> \cdot RLWE_A'(m)
832 mod_op.shoup_matrix_fma(
833 tmp_rlwe[0].as_mut(),
834 &rlwe_dash_m_parta,
835 &rlwe_dash_m_parta_shoup,
836 &decomposed_polys_of_rlweb,
837 );
838
839 // b_out += decomp<b_in> \cdot RLWE_B'(m)
840 mod_op.shoup_matrix_fma(
841 tmp_rlwe[1].as_mut(),
842 &rlwe_dash_m_partb,
843 &rlwe_dash_m_partb_shoup,
844 &decomposed_polys_of_rlweb,
845 );
846 }
847
848 // transform rlwe_out to coefficient domain
849 tmp_rlwe
850 .iter_mut()
851 .for_each(|r| ntt_op.backward(r.as_mut()));
852
853 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe[0].as_mut());
854 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe[1].as_mut());
855}
856
857/// Inplace mutates RGSW(m0) to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
858///
859/// RGSW x RGSW product requires multiple RLWE x RGSW products. For example,
860/// Define
861///
862/// RGSW(m0) = [RLWE(-sm), RLWE(\beta -sm), ..., RLWE(\beta^{d-1} -sm)
863/// RLWE(m), RLWE(\beta m), ..., RLWE(\beta^{d-1} m)]
864/// And RGSW(m1)
865///
866/// Then RGSW(m0) x RGSW(m1) equals:
867/// RGSW(m0m1) = [
868/// rlwe_x_rgsw(RLWE(-sm), RGSW(m1)),
869/// ...,
870/// rlwe_x_rgsw(RLWE(\beta^{d-1} -sm), RGSW(m1)),
871/// rlwe_x_rgsw(RLWE(m), RGSW(m1)),
872/// ...,
873/// rlwe_x_rgsw(RLWE(\beta^{d-1} m), RGSW(m1)),
874/// ]
875///
876/// Since noise growth in RLWE x RGSW depends on noise in RGSW ciphertext, it is
877/// clear to observe from above that noise in resulting RGSW(m0m1) equals noise
878/// accumulated in a single RLWE x RGSW and depends on noise in RGSW(m1) (i.e.
879/// rgsw_1_eval)
880///
881/// - rgsw_0: RGSW(m0) in coefficient domain
882/// - rgsw_1_eval: RGSW(m1) in evaluation domain
883pub(crate) fn rgsw_by_rgsw_inplace<
884 Rgsw: RgswCiphertext,
885 RgswMut: RgswCiphertextMut<R = Rgsw::R>,
886 Sc: RuntimeScratchMatrix<R = Rgsw::R, Rgsw = Rgsw>,
887 D: RlweDecomposer<Element = <Rgsw::R as Row>::Element>,
888 ModOp: VectorOps<Element = <Rgsw::R as Row>::Element>,
889 NttOp: Ntt<Element = <Rgsw::R as Row>::Element>,
890>(
891 rgsw0: &mut RgswMut,
892 rgsw1_eval: &Rgsw,
893 rgsw0_decomposer: &D,
894 rgsw1_decomposer: &D,
895 scratch_matrix: &mut Sc,
896 ntt_op: &NttOp,
897 mod_op: &ModOp,
898) where
899 <Rgsw::R as Row>::Element: Copy + Zero,
900 RgswMut: AsMut<[Rgsw::R]>,
901 RgswMut::R: RowMut,
902 // Rgsw: AsRef<[Rgsw::R]>,
903{
904 let (decomp_r_space, rgsw_space) = scratch_matrix
905 .scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space(rgsw0_decomposer, rgsw1_decomposer);
906
907 let mut rgsw_space = RgswCiphertextMutRef::new(
908 rgsw_space,
909 rgsw0_decomposer.decomposition_count_a().0,
910 rgsw0_decomposer.decomposition_count_b().0,
911 );
912 let (
913 (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb),
914 (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb),
915 ) = rgsw_space.split_mut();
916
917 let ((rgsw0_nsm_parta, rgsw0_nsm_partb), (rgsw0_m_parta, rgsw0_m_partb)) = rgsw0.split();
918 let ((rgsw1_nsm_parta, rgsw1_nsm_partb), (rgsw1_m_parta, rgsw1_m_partb)) = rgsw1_eval.split();
919
920 // RGSW x RGSW
921 izip!(
922 rgsw0_nsm_parta.iter().chain(rgsw0_m_parta),
923 rgsw0_nsm_partb.iter().chain(rgsw0_m_partb),
924 rlwe_dash_space_nsm_parta
925 .iter_mut()
926 .chain(rlwe_dash_space_m_parta.iter_mut()),
927 rlwe_dash_space_nsm_partb
928 .iter_mut()
929 .chain(rlwe_dash_space_m_partb.iter_mut()),
930 )
931 .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
932 // RLWE(m0) x RGSW(m1)
933
934 // Part A: Decomp<RLWE(m0)[A]> \cdot RLWE'(-sm1)
935 {
936 let decomp_r_parta = &mut decomp_r_space[..rgsw1_decomposer.decomposition_count_a().0];
937 decompose_r(
938 rlwe_a.as_ref(),
939 decomp_r_parta.as_mut(),
940 rgsw1_decomposer.a(),
941 );
942 decomp_r_parta
943 .iter_mut()
944 .for_each(|ri| ntt_op.forward(ri.as_mut()));
945 poly_fma_routine(
946 rlwe_out_a.as_mut(),
947 &decomp_r_parta,
948 &rgsw1_nsm_parta,
949 mod_op,
950 );
951 poly_fma_routine(
952 rlwe_out_b.as_mut(),
953 &decomp_r_parta,
954 &rgsw1_nsm_partb,
955 mod_op,
956 );
957 }
958
959 // Part B: Decompose<RLWE(m0)[B]> \cdot RLWE'(m1)
960 {
961 let decomp_r_partb = &mut decomp_r_space[..rgsw1_decomposer.decomposition_count_b().0];
962 decompose_r(
963 rlwe_b.as_ref(),
964 decomp_r_partb.as_mut(),
965 rgsw1_decomposer.b(),
966 );
967 decomp_r_partb
968 .iter_mut()
969 .for_each(|ri| ntt_op.forward(ri.as_mut()));
970 poly_fma_routine(rlwe_out_a.as_mut(), &decomp_r_partb, &rgsw1_m_parta, mod_op);
971 poly_fma_routine(rlwe_out_b.as_mut(), &decomp_r_partb, &rgsw1_m_partb, mod_op);
972 }
973 });
974
975 // copy over RGSW(m0m1) to RGSW(m0)
976 // let d = rgsw0.as_mut();
977 izip!(rgsw0.as_mut().iter_mut(), rgsw_space.data.iter())
978 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
979
980 // send back to coefficient domain
981 rgsw0
982 .as_mut()
983 .iter_mut()
984 .for_each(|ri| ntt_op.backward(ri.as_mut()));
985}
986
987/// Key switches input RLWE_{s'}(m) -> RLWE_{s}(m)
988///
989/// Let RLWE_{s'}(m) = [a, b] s.t. m+e = b - as'
990///
991/// Given key switchin key Ksk(s' -> s) = RLWE'_{s}(s') = [RLWE_{s}(beta^i s')]
992/// = [a, a*s + e + beta^i s'] for i \in [0,d), key switching computes:
993/// 1. RLWE_{s}(-s'a) = \sum signed_decompose(-a)[i] RLWE_{s}(beta^i s')
994/// 2. RLWE_{s}(m) = (b, 0) + RLWE_{s}(-s'a)
995///
996/// - rlwe_in: Input rlwe ciphertext
997/// - ksk: Key switching key Ksk(s' -> s) with polynomials in evaluation domain
998/// - ksk_shoup: Key switching key Ksk(s' -> s) with polynomials in evaluation
999/// domain in shoup representation
1000/// - decomposer: Decomposer used for key switching
1001pub(crate) fn rlwe_key_switch<
1002 M: MatrixMut + MatrixEntity,
1003 ModOp: GetModulus<Element = M::MatElement> + ShoupMatrixFMA<M::R> + VectorOps<Element = M::MatElement>,
1004 NttOp: Ntt<Element = M::MatElement>,
1005 D: Decomposer<Element = M::MatElement>,
1006>(
1007 rlwe_in: &M,
1008 ksk: &M,
1009 ksk_shoup: &M,
1010 decomposer: &D,
1011 ntt_op: &NttOp,
1012 mod_op: &ModOp,
1013) -> M
1014where
1015 <M as Matrix>::R: RowMut + RowEntity,
1016 M::MatElement: Copy,
1017{
1018 let ring_size = rlwe_in.dimension().1;
1019 assert!(rlwe_in.dimension().0 == 2);
1020 assert!(ksk.dimension() == (decomposer.decomposition_count().0 * 2, ring_size));
1021
1022 let mut rlwe_out = M::zeros(2, ring_size);
1023
1024 let mut tmp = M::zeros(decomposer.decomposition_count().0, ring_size);
1025 let mut tmp_row = M::R::zeros(ring_size);
1026
1027 // key switch RLWE part -A
1028 // negative A
1029 tmp_row.as_mut().copy_from_slice(rlwe_in.get_row_slice(0));
1030 mod_op.elwise_neg_mut(tmp_row.as_mut());
1031 // decompose -A and send to evaluation domain
1032 decompose_r(tmp_row.as_ref(), tmp.as_mut(), decomposer);
1033 tmp.iter_rows_mut()
1034 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
1035
1036 // RLWE_s(-A u) = B' + B, A' = (decomp(-A) * Ksk(u -> s)) + (B, 0)
1037 let (ksk_part_a, ksk_part_b) = ksk.split_at_row(decomposer.decomposition_count().0);
1038 let (ksk_part_a_shoup, ksk_part_b_shoup) =
1039 ksk_shoup.split_at_row(decomposer.decomposition_count().0);
1040 // Part A'
1041 mod_op.shoup_matrix_fma(
1042 rlwe_out.get_row_mut(0),
1043 &ksk_part_a,
1044 &ksk_part_a_shoup,
1045 tmp.as_ref(),
1046 );
1047 // Part B'
1048 mod_op.shoup_matrix_fma(
1049 rlwe_out.get_row_mut(1),
1050 &ksk_part_b,
1051 &ksk_part_b_shoup,
1052 tmp.as_ref(),
1053 );
1054 // back to coefficient domain
1055 rlwe_out
1056 .iter_rows_mut()
1057 .for_each(|r| ntt_op.backward(r.as_mut()));
1058
1059 // B' + B
1060 mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), rlwe_in.get_row_slice(1));
1061
1062 rlwe_out
1063}