poulpy_hal/api/convolution.rs
1use crate::layouts::{
2 Backend, CnvDftAccTerm, CnvPVecLBackendMut, CnvPVecLBackendRef, CnvPVecLOwned, CnvPVecRBackendMut, CnvPVecRBackendRef,
3 CnvPVecROwned, ScratchArena, VecZnxBackendRef, VecZnxBigBackendMut, VecZnxDftBackendMut,
4};
5
6/// Allocates prepared convolution operands ([`CnvPVecL`](crate::layouts::CnvPVecL), [`CnvPVecR`](crate::layouts::CnvPVecR)).
7pub trait CnvPVecAlloc<BE: Backend> {
8 fn cnv_pvec_left_alloc(&self, cols: usize, size: usize) -> CnvPVecLOwned<BE>;
9 fn cnv_pvec_right_alloc(&self, cols: usize, size: usize) -> CnvPVecROwned<BE>;
10}
11
12/// Returns the byte sizes for prepared convolution operands.
13pub trait CnvPVecBytesOf {
14 fn bytes_of_cnv_pvec_left(&self, cols: usize, size: usize) -> usize;
15 fn bytes_of_cnv_pvec_right(&self, cols: usize, size: usize) -> usize;
16}
17
18/// Bivariate convolution over `Z[X, Y] mod (X^N + 1)` where `Y = 2^{-K}`.
19///
20/// Provides methods to prepare left/right operands and apply the convolution.
21/// See method-level documentation for the mathematical formulation.
22pub trait Convolution<BE: Backend> {
23 /// Returns scratch bytes required for [`cnv_prepare_left`](Convolution::cnv_prepare_left).
24 fn cnv_prepare_left_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
25 /// Prepares a coefficient-domain [`VecZnx`](crate::layouts::VecZnx) as the left
26 /// operand of a bivariate convolution.
27 fn cnv_prepare_left(
28 &self,
29 res: &mut CnvPVecLBackendMut<'_, BE>,
30 a: &VecZnxBackendRef<'_, BE>,
31 mask: i64,
32 scratch: &mut ScratchArena<'_, BE>,
33 );
34
35 /// Returns scratch bytes required for [`cnv_prepare_right`](Convolution::cnv_prepare_right).
36 fn cnv_prepare_right_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
37 /// Prepares a coefficient-domain [`VecZnx`](crate::layouts::VecZnx) as the right
38 /// operand of a bivariate convolution.
39 fn cnv_prepare_right(
40 &self,
41 res: &mut CnvPVecRBackendMut<'_, BE>,
42 a: &VecZnxBackendRef<'_, BE>,
43 mask: i64,
44 scratch: &mut ScratchArena<'_, BE>,
45 );
46
47 /// Returns scratch bytes required for [`cnv_apply_dft`](Convolution::cnv_apply_dft).
48 fn cnv_apply_dft_tmp_bytes(&self, cnv_offset: usize, res_size: usize, a_size: usize, b_size: usize) -> usize;
49
50 /// Returns scratch bytes required for [`cnv_by_const_apply`](Convolution::cnv_by_const_apply).
51 fn cnv_by_const_apply_tmp_bytes(&self, cnv_offset: usize, res_size: usize, a_size: usize, b_size: usize) -> usize;
52
53 /// Evaluates a bivariate convolution over Z\[X, Y\] (x) Z\[Y\] mod (X^N + 1) where Y = 2^-K over the
54 /// selected columns and stores the result on the selected column, scaled by 2^{cnv_offset * Base2K}
55 ///
56 /// Behavior is identical to [Convolution::cnv_apply_dft] with `b` treated as a constant polynomial
57 /// in the X variable, for example:
58 ///```text
59 /// 1 X X^2 X^3
60 /// a = 1 [a00, a10, a20, a30] = (a00 + a01 * 2^-K) + (a10 + a11 * 2^-K) * X ...
61 /// Y [a01, a11, a21, a31]
62 ///
63 /// b = 1 [b0] = (b00 + b01 * 2^-K)
64 /// Y [b0]
65 /// ```
66 /// This method is intended to be used for multiplications by constants that are greater than the base2k.
67 #[allow(clippy::too_many_arguments)]
68 fn cnv_by_const_apply(
69 &self,
70 cnv_offset: usize,
71 res: &mut VecZnxBigBackendMut<'_, BE>,
72 res_col: usize,
73 a: &VecZnxBackendRef<'_, BE>,
74 a_col: usize,
75 b: &VecZnxBackendRef<'_, BE>,
76 b_col: usize,
77 b_coeff: usize,
78 scratch: &mut ScratchArena<'_, BE>,
79 );
80
81 /// `res[res_col] +=` the [`Convolution::cnv_by_const_apply`] result; limbs
82 /// the convolution would zero-fill are left untouched.
83 #[allow(clippy::too_many_arguments)]
84 fn cnv_by_const_apply_add(
85 &self,
86 cnv_offset: usize,
87 res: &mut VecZnxBigBackendMut<'_, BE>,
88 res_col: usize,
89 a: &VecZnxBackendRef<'_, BE>,
90 a_col: usize,
91 b: &VecZnxBackendRef<'_, BE>,
92 b_col: usize,
93 b_coeff: usize,
94 scratch: &mut ScratchArena<'_, BE>,
95 );
96
97 #[allow(clippy::too_many_arguments)]
98 /// Evaluates a bivariate convolution over Z\[X, Y\] (x) Z\[X, Y\] mod (X^N + 1) where Y = 2^-K over the
99 /// selected columns and stores the result on the selected column, scaled by 2^{cnv_offset * Base2K}
100 ///
101 /// # Example
102 ///```text
103 /// 1 X X^2 X^3
104 /// a = 1 [a00, a10, a20, a30] = (a00 + a01 * 2^-K) + (a10 + a11 * 2^-K) * X ...
105 /// Y [a01, a11, a21, a31]
106 ///
107 /// b = 1 [b00, b10, b20, b30] = (b00 + b01 * 2^-K) + (b10 + b11 * 2^-K) * X ...
108 /// Y [b01, b11, b21, b31]
109 ///
110 /// If cnv_offset = 0:
111 ///
112 /// 1 X X^2 X^3
113 /// res = 1 [r00, r10, r20, r30] = (r00 + r01 * 2^-K + r02 * 2^-2K + r03 * 2^-3K) + ... * X + ...
114 /// Y [r01, r11, r21, r31]
115 /// Y^2[r02, r12, r22, r32]
116 /// Y^3[r03, r13, r23, r33]
117 ///
118 /// If cnv_offset = 1:
119 ///
120 /// 1 X X^2 X^3
121 /// res = 1 [r01, r11, r21, r31] = (r01 + r02 * 2^-K + r03 * 2^-2K) + ... * X + ...
122 /// Y [r02, r12, r22, r32]
123 /// Y^2[r03, r13, r23, r33]
124 /// Y^3[ 0, 0, 0 , 0]
125 /// ```
126 /// If res.size() < a.size() + b.size() + k, result is truncated accordingly in the Y dimension.
127 fn cnv_apply_dft(
128 &self,
129 cnv_offset: usize,
130 res: &mut VecZnxDftBackendMut<'_, BE>,
131 res_col: usize,
132 a: &CnvPVecLBackendRef<'_, BE>,
133 a_col: usize,
134 b: &CnvPVecRBackendRef<'_, BE>,
135 b_col: usize,
136 scratch: &mut ScratchArena<'_, BE>,
137 );
138
139 /// Lazy-canonicalization convolution used by `glwe_mul_plain`; bit-identical to
140 /// the eager `cnv_prepare_left/right` + `cnv_apply_dft`.
141 fn cnv_prepare_left_lazy_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
142 fn cnv_prepare_left_lazy(
143 &self,
144 res: &mut CnvPVecLBackendMut<'_, BE>,
145 a: &VecZnxBackendRef<'_, BE>,
146 mask: i64,
147 scratch: &mut ScratchArena<'_, BE>,
148 );
149 fn cnv_prepare_right_lazy_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
150 fn cnv_prepare_right_lazy(
151 &self,
152 res: &mut CnvPVecRBackendMut<'_, BE>,
153 a: &VecZnxBackendRef<'_, BE>,
154 mask: i64,
155 scratch: &mut ScratchArena<'_, BE>,
156 );
157 fn cnv_apply_dft_lazy_tmp_bytes(&self, cnv_offset: usize, res_size: usize, a_size: usize, b_size: usize) -> usize;
158 #[allow(clippy::too_many_arguments)]
159 fn cnv_apply_dft_lazy(
160 &self,
161 cnv_offset: usize,
162 res: &mut VecZnxDftBackendMut<'_, BE>,
163 res_col: usize,
164 a: &CnvPVecLBackendRef<'_, BE>,
165 a_col: usize,
166 b: &CnvPVecRBackendRef<'_, BE>,
167 b_col: usize,
168 scratch: &mut ScratchArena<'_, BE>,
169 );
170
171 /// Accumulating variant of [`cnv_apply_dft`](Convolution::cnv_apply_dft):
172 /// `res[res_col] += a[a_col] (x) b[b_col]`, bit-identical to `cnv_apply_dft`
173 /// followed by a DFT-domain add. Limbs `>= min(res.size(), a.size() + b.size())`
174 /// are left untouched. Scratch requirement is
175 /// [`cnv_apply_dft_tmp_bytes`](Convolution::cnv_apply_dft_tmp_bytes).
176 #[allow(clippy::too_many_arguments)]
177 fn cnv_apply_dft_accumulate(
178 &self,
179 cnv_offset: usize,
180 res: &mut VecZnxDftBackendMut<'_, BE>,
181 res_col: usize,
182 a: &CnvPVecLBackendRef<'_, BE>,
183 a_col: usize,
184 b: &CnvPVecRBackendRef<'_, BE>,
185 b_col: usize,
186 scratch: &mut ScratchArena<'_, BE>,
187 );
188
189 /// Returns scratch bytes required for [`cnv_accumulate_dft`](Convolution::cnv_accumulate_dft).
190 ///
191 /// `a_size` and `b_size` are upper bounds over the sizes of the term operands.
192 fn cnv_accumulate_dft_tmp_bytes(&self, cnv_offset: usize, res_size: usize, a_size: usize, b_size: usize) -> usize;
193
194 /// Evaluates a sum of bivariate convolutions: `res[res_col] = Σ_t a_t ⊛ b_t`,
195 /// scaled by `2^{cnv_offset * Base2K}`, overwriting `res[res_col]`.
196 ///
197 /// Each term behaves like one [`Convolution::cnv_apply_dft`] call over the
198 /// selected columns and the per-term results are summed; with an empty
199 /// `terms` slice the output column is zeroed. Backends may fuse the
200 /// accumulation (one lazy reduction per output limb, destination written
201 /// once), so the result is congruent to — but not necessarily bit-identical
202 /// with — a sequence of [`Convolution::cnv_apply_dft_accumulate`] calls.
203 fn cnv_accumulate_dft<'a>(
204 &self,
205 cnv_offset: usize,
206 res: &mut VecZnxDftBackendMut<'_, BE>,
207 res_col: usize,
208 terms: &[CnvDftAccTerm<'a, BE>],
209 scratch: &mut ScratchArena<'_, BE>,
210 ) where
211 BE: 'a;
212
213 /// Returns scratch bytes required for [`cnv_pairwise_apply_dft`](Convolution::cnv_pairwise_apply_dft).
214 fn cnv_pairwise_apply_dft_tmp_bytes(&self, cnv_offset: usize, res_size: usize, a_size: usize, b_size: usize) -> usize;
215
216 #[allow(clippy::too_many_arguments)]
217 /// Evaluates the bivariate pair-wise convolution res = (a\[i\] + a\[j\]) * (b\[i\] + b\[j\]).
218 /// If i == j then calls [Convolution::cnv_apply_dft], i.e. res = a\[i\] * b\[i\].
219 /// See [Convolution::cnv_apply_dft] for information about the bivariate convolution.
220 fn cnv_pairwise_apply_dft(
221 &self,
222 cnv_offset: usize,
223 res: &mut VecZnxDftBackendMut<'_, BE>,
224 res_col: usize,
225 a: &CnvPVecLBackendRef<'_, BE>,
226 b: &CnvPVecRBackendRef<'_, BE>,
227 i: usize,
228 j: usize,
229 scratch: &mut ScratchArena<'_, BE>,
230 );
231
232 /// Returns scratch bytes required for [`cnv_prepare_self`](Convolution::cnv_prepare_self).
233 fn cnv_prepare_self_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
234
235 /// Prepares both left and right convolution operands from the same input polynomial,
236 /// sharing the FFT/NTT computation. This is an optimization for self-convolution
237 /// (squaring) where both operands are the same polynomial.
238 fn cnv_prepare_self(
239 &self,
240 left: &mut CnvPVecLBackendMut<'_, BE>,
241 right: &mut CnvPVecRBackendMut<'_, BE>,
242 a: &VecZnxBackendRef<'_, BE>,
243 mask: i64,
244 scratch: &mut ScratchArena<'_, BE>,
245 );
246}