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