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