poulpy/rns/src/ring/impl_u64/rescaling_rns.rs
rns/src/ring/impl_u64/rescaling_rns.rsRaw
1use crate::modulus::barrett::Barrett;
2use crate::modulus::{BARRETT, NONE, ONCE};
3use crate::poly::{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 /// buf is unused if <ROUND=false,NTT=false>
12 pub fn div_by_last_modulus<const ROUND: bool, const NTT: bool>(
13 &self,
14 a: &PolyRNS<u64>,
15 buf: &mut [Poly<u64>; 2],
16 b: &mut PolyRNS<u64>,
17 ) {
18 debug_assert!(self.level() != 0, "invalid call: self.level()=0");
19 debug_assert!(
20 a.level() >= self.level(),
21 "invalid input a: a.level()={} < self.level()={}",
22 a.level(),
23 self.level()
24 );
25 debug_assert!(
26 b.level() >= self.level() - 1,
27 "invalid input b: b.level()={} < self.level()-1={}",
28 b.level(),
29 self.level() - 1
30 );
31
32 let level = self.level();
33 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
34 let r_last: &Ring<u64> = &self.0[level];
35
36 if ROUND {
37 let q_level_half: u64 = r_last.modulus.q >> 1;
38
39 let (buf_q_scaling, buf_qi_scaling) = buf.split_at_mut(1);
40
41 if NTT {
42 r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
43 r_last.a_add_b_scalar_into_a::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
44 for (i, r) in self.0[0..level].iter().enumerate() {
45 r_last.a_add_b_scalar_into_c::<NONE>(
46 &buf_q_scaling[0],
47 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
48 &mut buf_qi_scaling[0],
49 );
50 r.ntt_inplace::<true>(&mut buf_qi_scaling[0]);
51 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
52 &buf_qi_scaling[0],
53 a.at(i),
54 &rescaling_constants.0[i],
55 b.at_mut(i),
56 );
57 }
58 } else {
59 r_last.a_add_b_scalar_into_c::<ONCE>(
60 a.at(self.level()),
61 &q_level_half,
62 &mut buf_q_scaling[0],
63 );
64 for (i, r) in self.0[0..level].iter().enumerate() {
65 r_last.a_add_b_scalar_into_c::<NONE>(
66 &buf_q_scaling[0],
67 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
68 &mut buf_qi_scaling[0],
69 );
70 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
71 &buf_qi_scaling[0],
72 a.at(i),
73 &rescaling_constants.0[i],
74 b.at_mut(i),
75 );
76 }
77 }
78 } else {
79 if NTT {
80 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.split_at_mut(1);
81 self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
82 for (i, r) in self.0[0..level].iter().enumerate() {
83 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
84 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
85 &buf_ntt_qi_scaling[0],
86 a.at(i),
87 &rescaling_constants.0[i],
88 b.at_mut(i),
89 );
90 }
91 } else {
92 for (i, r) in self.0[0..level].iter().enumerate() {
93 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
94 a.at(level),
95 a.at(i),
96 &rescaling_constants.0[i],
97 b.at_mut(i),
98 );
99 }
100 }
101 }
102 }
103
104 /// Updates a to floor(a / q[b.level()]).
105 /// Expects a to be in the NTT domain.
106 pub fn div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(
107 &self,
108 buf: &mut [Poly<u64>; 2],
109 a: &mut PolyRNS<u64>,
110 ) {
111 debug_assert!(
112 self.level() <= a.level(),
113 "invalid input a: a.level()={} < self.level()={}",
114 self.level(),
115 a.level()
116 );
117
118 let level = self.level();
119 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
120 let r_last: &Ring<u64> = &self.0[level];
121
122 if ROUND {
123 let q_level_half: u64 = r_last.modulus.q >> 1;
124 let (buf_q_scaling, buf_qi_scaling) = buf.split_at_mut(1);
125
126 if NTT {
127 r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
128 r_last.a_add_b_scalar_into_a::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
129 for (i, r) in self.0[0..level].iter().enumerate() {
130 r_last.a_add_b_scalar_into_c::<NONE>(
131 &buf_q_scaling[0],
132 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
133 &mut buf_qi_scaling[0],
134 );
135 r.ntt_inplace::<false>(&mut buf_qi_scaling[0]);
136 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
137 &buf_qi_scaling[0],
138 &rescaling_constants.0[i],
139 a.at_mut(i),
140 );
141 }
142 } else {
143 let (a_qi, a_q_last) = a.0.split_at_mut(self.level());
144 r_last.a_add_b_scalar_into_a::<ONCE>(&q_level_half, &mut a_q_last[0]);
145 for (i, r) in self.0[0..level].iter().enumerate() {
146 r.b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a::<1, ONCE>(
147 &a_q_last[0],
148 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
149 &rescaling_constants.0[i],
150 &mut a_qi[i],
151 );
152 }
153 }
154 } else {
155 if NTT {
156 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.split_at_mut(1);
157 r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
158 for (i, r) in self.0[0..level].iter().enumerate() {
159 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
160 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
161 &buf_ntt_qi_scaling[0],
162 &rescaling_constants.0[i],
163 a.at_mut(i),
164 );
165 }
166 } else {
167 let (a_i, a_level) = a.0.split_at_mut(level);
168 for (i, r) in self.0[0..level].iter().enumerate() {
169 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
170 &a_level[0],
171 &rescaling_constants.0[i],
172 &mut a_i[i],
173 );
174 }
175 }
176 }
177 }
178
179 /// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i])
180 pub fn div_by_last_moduli<const ROUND: bool, const NTT: bool>(
181 &self,
182 nb_moduli_dropped: usize,
183 a: &PolyRNS<u64>,
184 buf0: &mut [Poly<u64>; 2],
185 buf1: &mut PolyRNS<u64>,
186 c: &mut PolyRNS<u64>,
187 ) {
188 debug_assert!(
189 nb_moduli_dropped <= self.level(),
190 "invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
191 nb_moduli_dropped,
192 self.level()
193 );
194 debug_assert!(
195 a.level() >= self.level(),
196 "invalid input a: a.level()={} < self.level()={}",
197 a.level(),
198 self.level()
199 );
200 debug_assert!(
201 buf1.level() >= self.level(),
202 "invalid input buf: buf.level()={} < self.level()={}",
203 buf1.level(),
204 self.level()
205 );
206 debug_assert!(
207 c.level() >= self.level() - nb_moduli_dropped,
208 "invalid input c: c.level()={} < self.level()-nb_moduli_dropped={}",
209 c.level(),
210 self.level() - nb_moduli_dropped
211 );
212
213 if nb_moduli_dropped == 0 {
214 if a != c {
215 c.copy(a);
216 }
217 } else {
218 if NTT {
219 self.intt::<false>(a, buf1);
220 (0..nb_moduli_dropped).for_each(|i| {
221 self.at_level(self.level() - i)
222 .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
223 });
224 self.at_level(self.level() - nb_moduli_dropped)
225 .ntt::<false>(buf1, c);
226 } else {
227 self.div_by_last_modulus::<ROUND, false>(a, buf0, buf1);
228
229 (1..nb_moduli_dropped - 1).for_each(|i| {
230 self.at_level(self.level() - i)
231 .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1);
232 });
233
234 self.at_level(self.level() - nb_moduli_dropped + 1)
235 .div_by_last_modulus::<ROUND, false>(buf1, buf0, c);
236 }
237 }
238 }
239
240 /// Updates a to floor(a / prod_{level - nb_moduli_dropped}^{level} q[i])
241 pub fn div_by_last_moduli_inplace<const ROUND: bool, const NTT: bool>(
242 &self,
243 nb_moduli_dropped: usize,
244 buf0: &mut [Poly<u64>; 2],
245 buf1: &mut PolyRNS<u64>,
246 a: &mut PolyRNS<u64>,
247 ) {
248 debug_assert!(
249 nb_moduli_dropped <= self.level(),
250 "invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
251 nb_moduli_dropped,
252 self.level()
253 );
254 debug_assert!(
255 a.level() >= self.level(),
256 "invalid input a: a.level()={} < self.level()={}",
257 a.level(),
258 self.level()
259 );
260 debug_assert!(
261 buf1.level() >= self.level(),
262 "invalid input buf: buf.level()={} < self.level()={}",
263 buf1.level(),
264 self.level()
265 );
266 if nb_moduli_dropped == 0 {
267 return;
268 }
269
270 if NTT {
271 self.intt::<false>(a, buf1);
272 (0..nb_moduli_dropped).for_each(|i| {
273 self.at_level(self.level() - i)
274 .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
275 });
276 self.at_level(self.level() - nb_moduli_dropped)
277 .ntt::<false>(buf1, a);
278 } else {
279 (0..nb_moduli_dropped).for_each(|i| {
280 self.at_level(self.level() - i)
281 .div_by_last_modulus_inplace::<ROUND, false>(buf0, a)
282 });
283 }
284 }
285}