poulpy/math/src/ring/impl_u64/rescaling_rns.rs
math/src/ring/impl_u64/rescaling_rns.rsRaw
1use crate::modulus::barrett::Barrett;
2use crate::modulus::{NONE, ONCE, BARRETT};
3use crate::poly::PolyRNS;
4use crate::ring::Ring;
5use crate::ring::RingRNS;
6use crate::scalar::ScalarRNS;
7extern crate test;
8
9impl RingRNS<u64> {
10 /// Updates b to floor(a / q[b.level()]).
11 pub fn div_floor_by_last_modulus<const NTT: bool>(
12 &self,
13 a: &PolyRNS<u64>,
14 buf: &mut PolyRNS<u64>,
15 b: &mut PolyRNS<u64>,
16 ) {
17 debug_assert!(self.level() != 0, "invalid call: self.level()=0");
18 debug_assert!(
19 self.level() <= a.level(),
20 "invalid input a: self.level()={} > a.level()={}",
21 self.level(),
22 a.level()
23 );
24 debug_assert!(
25 b.level() >= self.level() - 1,
26 "invalid input b: b.level()={} < self.level()-1={}",
27 b.level(),
28 self.level() - 1
29 );
30
31 let level = self.level();
32 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
33
34 if NTT {
35 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
36 self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
37 for (i, r) in self.0[0..level].iter().enumerate() {
38 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
39 r.a_sub_b_mul_c_scalar_barrett::<2, ONCE>(
40 &buf_ntt_qi_scaling[0],
41 a.at(i),
42 &rescaling_constants.0[i],
43 b.at_mut(i),
44 );
45 }
46 } else {
47 for (i, r) in self.0[0..level].iter().enumerate() {
48 r.a_sub_b_mul_c_scalar_barrett::<2, ONCE>(
49 a.at(level),
50 a.at(i),
51 &rescaling_constants.0[i],
52 b.at_mut(i),
53 );
54 }
55 }
56 }
57
58 /// Updates a to floor(a / q[b.level()]).
59 /// Expects a to be in the NTT domain.
60 pub fn div_floor_by_last_modulus_inplace<const NTT: bool>(
61 &self,
62 buf: &mut PolyRNS<u64>,
63 a: &mut PolyRNS<u64>,
64 ) {
65 debug_assert!(
66 self.level() <= a.level(),
67 "invalid input a: self.level()={} > a.level()={}",
68 self.level(),
69 a.level()
70 );
71
72 let level = self.level();
73 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
74
75 if NTT {
76 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
77 self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
78 for (i, r) in self.0[0..level].iter().enumerate() {
79 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
80 r.a_sub_b_mul_c_scalar_barrett_inplace::<2, ONCE>(
81 &buf_ntt_qi_scaling[0],
82 &rescaling_constants.0[i],
83 a.at_mut(i),
84 );
85 }
86 } else {
87 let (a_i, a_level) = a.0.split_at_mut(level);
88 for (i, r) in self.0[0..level].iter().enumerate() {
89 r.a_sub_b_mul_c_scalar_barrett_inplace::<2, ONCE>(
90 &a_level[0],
91 &rescaling_constants.0[i],
92 &mut a_i[i],
93 );
94 }
95 }
96 }
97
98 /// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i])
99 pub fn div_floor_by_last_moduli<const NTT: bool>(
100 &self,
101 nb_moduli: usize,
102 a: &PolyRNS<u64>,
103 buf: &mut PolyRNS<u64>,
104 c: &mut PolyRNS<u64>,
105 ) {
106 debug_assert!(
107 nb_moduli <= self.level(),
108 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
109 nb_moduli,
110 a.level()
111 );
112 debug_assert!(
113 a.level() <= self.level(),
114 "invalid input a: self.level()={} > a.level()={}",
115 self.level(),
116 a.level()
117 );
118 debug_assert!(
119 buf.level() >= self.level() - 1,
120 "invalid input buf: buf.level()={} < a.level()-1={}",
121 buf.level(),
122 a.level() - 1
123 );
124 debug_assert!(
125 c.level() >= self.level() - nb_moduli,
126 "invalid input c: c.level()={} < c.level()-nb_moduli={}",
127 c.level(),
128 a.level() - nb_moduli
129 );
130
131 if nb_moduli == 0 {
132 if a != c {
133 c.copy(a);
134 }
135 } else {
136 if NTT {
137 self.intt::<false>(a, buf);
138 (0..nb_moduli).for_each(|i| {
139 self.at_level(self.level() - i)
140 .div_floor_by_last_modulus_inplace::<false>(
141 &mut PolyRNS::<u64>::default(),
142 buf,
143 )
144 });
145 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c);
146 } else {
147
148 let empty_buf: &mut PolyRNS<u64> = &mut PolyRNS::<u64>::default();
149
150 if nb_moduli == 1{
151 self.div_floor_by_last_modulus::<false>(a, empty_buf, c);
152 }else{
153 self.div_floor_by_last_modulus::<false>(a, empty_buf, buf);
154 }
155
156 (1..nb_moduli-1).for_each(|i| {
157 self.at_level(self.level() - i)
158 .div_floor_by_last_modulus_inplace::<false>(empty_buf, buf);
159 });
160
161 self.at_level(self.level()-nb_moduli+1).div_floor_by_last_modulus::<false>(buf, empty_buf, c);
162 }
163 }
164 }
165
166 /// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i])
167 pub fn div_floor_by_last_moduli_inplace<const NTT: bool>(
168 &self,
169 nb_moduli: usize,
170 buf: &mut PolyRNS<u64>,
171 a: &mut PolyRNS<u64>,
172 ) {
173 debug_assert!(
174 self.level() <= a.level(),
175 "invalid input a: self.level()={} > a.level()={}",
176 self.level(),
177 a.level()
178 );
179 debug_assert!(
180 nb_moduli <= a.level(),
181 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
182 nb_moduli,
183 a.level()
184 );
185 if nb_moduli == 0{
186 return
187 }
188 if NTT {
189 self.intt::<false>(a, buf);
190 (0..nb_moduli).for_each(|i| {
191 self.at_level(self.level() - i)
192 .div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)
193 });
194 self.at_level(self.level() - nb_moduli+1).ntt::<false>(buf, a);
195 } else {
196 (0..nb_moduli).for_each(|i| {
197 self.at_level(self.level() - i)
198 .div_floor_by_last_modulus_inplace::<false>(buf, a);
199 });
200 }
201 }
202
203 /// Updates b to round(a / q[b.level()]).
204 /// Expects b to be in the NTT domain.
205 pub fn div_round_by_last_modulus<const NTT: bool>(
206 &self,
207 a: &PolyRNS<u64>,
208 buf: &mut PolyRNS<u64>,
209 b: &mut PolyRNS<u64>,
210 ) {
211 debug_assert!(
212 self.level() <= a.level(),
213 "invalid input a: self.level()={} > a.level()={}",
214 self.level(),
215 a.level()
216 );
217 debug_assert!(
218 b.level() >= a.level() - 1,
219 "invalid input b: b.level()={} < a.level()-1={}",
220 b.level(),
221 a.level() - 1
222 );
223
224 let level: usize = self.level();
225 let r_last: &Ring<u64> = &self.0[level];
226 let q_level_half: u64 = r_last.modulus.q >> 1;
227 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
228 let (buf_q_scaling, buf_qi_scaling) = buf.0.split_at_mut(1);
229
230 if NTT {
231 r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
232 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
233 for (i, r) in self.0[0..level].iter().enumerate() {
234 r_last.add_scalar::<NONE>(
235 &buf_q_scaling[0],
236 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
237 &mut buf_qi_scaling[0],
238 );
239 r.ntt_inplace::<true>(&mut buf_qi_scaling[0]);
240 r.a_sub_b_mul_c_scalar_barrett::<2, ONCE>(
241 &buf_qi_scaling[0],
242 a.at(i),
243 &rescaling_constants.0[i],
244 b.at_mut(i),
245 );
246 }
247 } else {
248 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
249 for (i, r) in self.0[0..level].iter().enumerate() {
250 r_last.add_scalar::<NONE>(
251 &buf_q_scaling[0],
252 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
253 &mut buf_qi_scaling[0],
254 );
255 r.a_sub_b_mul_c_scalar_barrett::<2, ONCE>(
256 &buf_qi_scaling[0],
257 a.at(i),
258 &rescaling_constants.0[i],
259 b.at_mut(i),
260 );
261 }
262 }
263 }
264
265 /// Updates a to round(a / q[b.level()]).
266 /// Expects a to be in the NTT domain.
267 pub fn div_round_by_last_modulus_inplace<const NTT: bool>(
268 &self,
269 buf: &mut PolyRNS<u64>,
270 a: &mut PolyRNS<u64>,
271 ) {
272 debug_assert!(
273 self.level() <= a.level(),
274 "invalid input a: self.level()={} > a.level()={}",
275 self.level(),
276 a.level()
277 );
278
279 let level = self.level();
280 let r_last: &Ring<u64> = &self.0[level];
281 let q_level_half: u64 = r_last.modulus.q >> 1;
282 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
283 let (buf_q_scaling, buf_qi_scaling) = buf.0.split_at_mut(1);
284
285 if NTT {
286 r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
287 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
288 for (i, r) in self.0[0..level].iter().enumerate() {
289 r_last.add_scalar::<NONE>(
290 &buf_q_scaling[0],
291 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
292 &mut buf_qi_scaling[0],
293 );
294 r.ntt_inplace::<false>(&mut buf_qi_scaling[0]);
295 r.a_sub_b_mul_c_scalar_barrett_inplace::<2, ONCE>(
296 &buf_qi_scaling[0],
297 &rescaling_constants.0[i],
298 a.at_mut(i),
299 );
300 }
301 } else {
302 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
303 for (i, r) in self.0[0..level].iter().enumerate() {
304 r_last.add_scalar::<NONE>(
305 &buf_q_scaling[0],
306 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
307 &mut buf_qi_scaling[0],
308 );
309 r.a_sub_b_mul_c_scalar_barrett_inplace::<2, ONCE>(
310 &buf_qi_scaling[0],
311 &rescaling_constants.0[i],
312 a.at_mut(i),
313 );
314 }
315 }
316 }
317
318 /// Updates b to round(a / prod_{level - nb_moduli}^{level} q[i])
319 pub fn div_round_by_last_moduli<const NTT: bool>(
320 &self,
321 nb_moduli: usize,
322 a: &PolyRNS<u64>,
323 buf: &mut PolyRNS<u64>,
324 c: &mut PolyRNS<u64>,
325 ) {
326 debug_assert!(
327 self.level() <= a.level(),
328 "invalid input a: self.level()={} > a.level()={}",
329 self.level(),
330 a.level()
331 );
332 debug_assert!(
333 c.level() >= a.level() - 1,
334 "invalid input b: b.level()={} < a.level()-1={}",
335 c.level(),
336 a.level() - 1
337 );
338 debug_assert!(
339 nb_moduli <= a.level(),
340 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
341 nb_moduli,
342 a.level()
343 );
344
345 if nb_moduli == 0 {
346 if a != c {
347 c.copy(a);
348 }
349 } else {
350 if NTT {
351 self.intt::<false>(a, buf);
352 (0..nb_moduli).for_each(|i| {
353 self.at_level(self.level() - i)
354 .div_round_by_last_modulus_inplace::<false>(
355 &mut PolyRNS::<u64>::default(),
356 buf,
357 )
358 });
359 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c);
360 } else {
361 self.div_round_by_last_modulus::<false>(a, buf, c);
362 (1..nb_moduli).for_each(|i| {
363 self.at_level(self.level() - i)
364 .div_round_by_last_modulus_inplace::<false>(buf, c)
365 });
366 }
367 }
368 }
369
370 /// Updates a to round(a / prod_{level - nb_moduli}^{level} q[i])
371 pub fn div_round_by_last_moduli_inplace<const NTT: bool>(
372 &self,
373 nb_moduli: usize,
374 buf: &mut PolyRNS<u64>,
375 a: &mut PolyRNS<u64>,
376 ) {
377 debug_assert!(
378 self.level() <= a.level(),
379 "invalid input a: self.level()={} > a.level()={}",
380 self.level(),
381 a.level()
382 );
383 debug_assert!(
384 nb_moduli <= a.level(),
385 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
386 nb_moduli,
387 a.level()
388 );
389 if NTT {
390 self.intt::<false>(a, buf);
391 (0..nb_moduli).for_each(|i| {
392 self.at_level(self.level() - i)
393 .div_round_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)
394 });
395 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, a);
396 } else {
397 (0..nb_moduli).for_each(|i| {
398 self.at_level(self.level() - i)
399 .div_round_by_last_modulus_inplace::<false>(buf, a)
400 });
401 }
402 }
403}