arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/rgsw/runtime.rs
src/rgsw/runtime.rsRaw
  1use itertools::izip;
  2use num_traits::Zero;
  3
  4use crate::{
  5    backend::{ArithmeticOps, GetModulus, Modulus, ShoupMatrixFMA, VectorOps},
  6    decomposer::{Decomposer, RlweDecomposer},
  7    ntt::Ntt,
  8    Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
  9};
 10
 11use super::IsTrivial;
 12
 13pub(crate) fn routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
 14    write_to_row: &mut [R::Element],
 15    matrix_a: &[R],
 16    matrix_b: &[R],
 17    mod_op: &ModOp,
 18) {
 19    izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
 20        mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
 21    });
 22}
 23
 24/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
 25/// output matrix decomp_r
 26///
 27/// Note that decomposition of r(X) requires decomposition of each of
 28/// coefficients.
 29///
 30/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
 31///   will be stored at i^th row.
 32pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
 33    r: &[R::Element],
 34    decomp_r: &mut [R],
 35    decomposer: &D,
 36) where
 37    R::Element: Copy,
 38{
 39    let ring_size = r.len();
 40
 41    for ri in 0..ring_size {
 42        decomposer
 43            .decompose_iter(&r[ri])
 44            .enumerate()
 45            .for_each(|(index, el)| {
 46                decomp_r[index].as_mut()[ri] = el;
 47            });
 48    }
 49}
 50
 51/// Sends RLWE_{s}(X) -> RLWE_{s}(X^k) where k is some galois element
 52///
 53/// - scratch_matrix: must have dimension at-least d+2 x ring_size. d rows to
 54///   store decomposed polynomials and 2 for rlwe
 55pub(crate) fn galois_auto<
 56    MT: Matrix + IsTrivial + MatrixMut,
 57    Mmut: MatrixMut<MatElement = MT::MatElement>,
 58    ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
 59    NttOp: Ntt<Element = MT::MatElement>,
 60    D: Decomposer<Element = MT::MatElement>,
 61>(
 62    rlwe_in: &mut MT,
 63    ksk: &Mmut,
 64    scratch_matrix: &mut Mmut,
 65    auto_map_index: &[usize],
 66    auto_map_sign: &[bool],
 67    mod_op: &ModOp,
 68    ntt_op: &NttOp,
 69    decomposer: &D,
 70) where
 71    <Mmut as Matrix>::R: RowMut,
 72    <MT as Matrix>::R: RowMut,
 73    MT::MatElement: Copy + Zero,
 74{
 75    let d = decomposer.decomposition_count();
 76    let ring_size = rlwe_in.dimension().1;
 77    assert!(rlwe_in.dimension().0 == 2);
 78    assert!(scratch_matrix.fits(d + 2, ring_size));
 79
 80    // scratch matrix is guaranteed to have at-least d+2 rows but can have more than
 81    // d+2 rows. We require to split them into sub-matrices of exact sizes one with
 82    // d rows for storing decomposed polynomial and second with 2 rows to act
 83    // tomperary space for RLWE ciphertext. Exact sizes is necessary to avoid any
 84    // irrelevant extra FMA or NTT ops.
 85    let (scratch_matrix_d_ring, other_half) = scratch_matrix.split_at_row_mut(d);
 86    let (tmp_rlwe_out, _) = other_half.split_at_mut(2);
 87
 88    debug_assert!(tmp_rlwe_out.len() == 2);
 89    debug_assert!(scratch_matrix_d_ring.len() == d);
 90
 91    if !rlwe_in.is_trivial() {
 92        tmp_rlwe_out.iter_mut().for_each(|r| {
 93            r.as_mut().fill(Mmut::MatElement::zero());
 94        });
 95
 96        // send a(X) -> a(X^k) and decompose a(X^k)
 97        izip!(
 98            rlwe_in.get_row(0),
 99            auto_map_index.iter(),
100            auto_map_sign.iter()
101        )
102        .for_each(|(el_in, to_index, sign)| {
103            let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
104
105            decomposer
106                .decompose_iter(&el_out)
107                .enumerate()
108                .for_each(|(index, el)| {
109                    scratch_matrix_d_ring[index].as_mut()[*to_index] = el;
110                });
111        });
112
113        // transform decomposed a(X^k) to evaluation domain
114        scratch_matrix_d_ring.iter_mut().for_each(|r| {
115            ntt_op.forward(r.as_mut());
116        });
117
118        // RLWE(m^k) = a', b'; RLWE(m) = a, b
119        // key switch: (a * RLWE'(s(X^k)))
120        let (ksk_a, ksk_b) = ksk.split_at_row(d);
121        // a' = decomp<a> * RLWE'_A(s(X^k))
122        routine(
123            tmp_rlwe_out[0].as_mut(),
124            scratch_matrix_d_ring,
125            ksk_a,
126            mod_op,
127        );
128
129        // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
130        routine(
131            tmp_rlwe_out[1].as_mut(),
132            scratch_matrix_d_ring,
133            ksk_b,
134            mod_op,
135        );
136
137        // transform RLWE(m^k) to coefficient domain
138        tmp_rlwe_out
139            .iter_mut()
140            .for_each(|r| ntt_op.backward(r.as_mut()));
141
142        // send b(X) -> b(X^k) and then b'(X) += b(X^k)
143        izip!(
144            rlwe_in.get_row(1),
145            auto_map_index.iter(),
146            auto_map_sign.iter()
147        )
148        .for_each(|(el_in, to_index, sign)| {
149            let row = tmp_rlwe_out[1].as_mut();
150            if !*sign {
151                row[*to_index] = mod_op.sub(&row[*to_index], el_in);
152            } else {
153                row[*to_index] = mod_op.add(&row[*to_index], el_in);
154            }
155        });
156
157        // copy over A; Leave B for later
158        rlwe_in
159            .get_row_mut(0)
160            .copy_from_slice(tmp_rlwe_out[0].as_ref());
161    } else {
162        // RLWE is trivial, a(X) is 0.
163        // send b(X) -> b(X^k)
164        izip!(
165            rlwe_in.get_row(1),
166            auto_map_index.iter(),
167            auto_map_sign.iter()
168        )
169        .for_each(|(el_in, to_index, sign)| {
170            if !*sign {
171                tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
172            } else {
173                tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
174            }
175        });
176    }
177
178    // Copy over B
179    rlwe_in
180        .get_row_mut(1)
181        .copy_from_slice(tmp_rlwe_out[1].as_ref());
182}
183
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
549pub(crate) fn rgsw_by_rgsw_inplace<
550    Mmut: MatrixMut,
551    D: RlweDecomposer<Element = Mmut::MatElement>,
552    ModOp: VectorOps<Element = Mmut::MatElement>,
553    NttOp: Ntt<Element = Mmut::MatElement>,
554>(
555    rgsw_0: &mut Mmut,
556    rgsw_1_eval: &Mmut,
557    decomposer: &D,
558    scratch_matrix: &mut Mmut,
559    ntt_op: &NttOp,
560    mod_op: &ModOp,
561) where
562    <Mmut as Matrix>::R: RowMut,
563    Mmut::MatElement: Copy + Zero,
564{
565    let decomposer_a = decomposer.a();
566    let decomposer_b = decomposer.b();
567    let d_a = decomposer_a.decomposition_count();
568    let d_b = decomposer_b.decomposition_count();
569    let max_d = std::cmp::max(d_a, d_b);
570    let rgsw_rows = d_a * 2 + d_b * 2;
571    assert!(rgsw_0.dimension().0 == rgsw_rows);
572    let ring_size = rgsw_0.dimension().1;
573    assert!(rgsw_1_eval.dimension() == (rgsw_rows, ring_size));
574    assert!(scratch_matrix.fits(max_d + rgsw_rows, ring_size));
575
576    let (decomp_r_space, rgsw_space) = scratch_matrix.split_at_row_mut(max_d);
577
578    // zero rgsw_space
579    rgsw_space
580        .iter_mut()
581        .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
582    let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(d_a * 2);
583    let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
584        rlwe_dash_space_nsm.split_at_mut(d_a);
585    let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) = rlwe_dash_space_m.split_at_mut(d_b);
586
587    let (rgsw0_nsm, rgsw0_m) = rgsw_0.split_at_row(d_a * 2);
588    let (rgsw1_nsm, rgsw1_m) = rgsw_1_eval.split_at_row(d_a * 2);
589
590    // RGSW x RGSW
591    izip!(
592        rgsw0_nsm.iter().take(d_a).chain(rgsw0_m.iter().take(d_b)),
593        rgsw0_nsm.iter().skip(d_a).chain(rgsw0_m.iter().skip(d_b)),
594        rlwe_dash_space_nsm_parta
595            .iter_mut()
596            .chain(rlwe_dash_space_m_parta.iter_mut()),
597        rlwe_dash_space_nsm_partb
598            .iter_mut()
599            .chain(rlwe_dash_space_m_partb.iter_mut()),
600    )
601    .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
602        // Part A
603        decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer_a);
604        decomp_r_space
605            .iter_mut()
606            .take(d_a)
607            .for_each(|ri| ntt_op.forward(ri.as_mut()));
608        routine(
609            rlwe_out_a.as_mut(),
610            &decomp_r_space[..d_a],
611            &rgsw1_nsm[..d_a],
612            mod_op,
613        );
614        routine(
615            rlwe_out_b.as_mut(),
616            &decomp_r_space[..d_a],
617            &rgsw1_nsm[d_a..],
618            mod_op,
619        );
620
621        // Part B
622        decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer_b);
623        decomp_r_space
624            .iter_mut()
625            .take(d_b)
626            .for_each(|ri| ntt_op.forward(ri.as_mut()));
627        routine(
628            rlwe_out_a.as_mut(),
629            &decomp_r_space[..d_b],
630            &rgsw1_m[..d_b],
631            mod_op,
632        );
633        routine(
634            rlwe_out_b.as_mut(),
635            &decomp_r_space[..d_b],
636            &rgsw1_m[d_b..],
637            mod_op,
638        );
639    });
640
641    // copy over RGSW(m0m1) into RGSW(m0)
642    izip!(rgsw_0.iter_rows_mut(), rgsw_space.iter())
643        .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
644
645    // send back to coefficient domain
646    rgsw_0
647        .iter_rows_mut()
648        .for_each(|ri| ntt_op.backward(ri.as_mut()));
649}