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_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
37 let q_level_half: u64 = r_last.modulus.q >> 1;
38
39 let (buf_q_scaling, buf_qi_scaling) = buf.0.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>(a.at(self.level()), &q_level_half, &mut buf_q_scaling[0]);
60 for (i, r) in self.0[0..level].iter().enumerate() {
61 r_last.a_add_b_scalar_into_c::<NONE>(
62 &buf_q_scaling[0],
63 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
64 &mut buf_qi_scaling[0],
65 );
66 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
67 &buf_qi_scaling[0],
68 a.at(i),
69 &rescaling_constants.0[i],
70 b.at_mut(i),
71 );
72 }
73 }
74 }else{
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_into_d::<2, ONCE>(
81 &buf_ntt_qi_scaling[0],
82 a.at(i),
83 &rescaling_constants.0[i],
84 b.at_mut(i),
85 );
86 }
87 } else {
88 for (i, r) in self.0[0..level].iter().enumerate() {
89 r.a_sub_b_mul_c_scalar_barrett_into_d::<2, ONCE>(
90 a.at(level),
91 a.at(i),
92 &rescaling_constants.0[i],
93 b.at_mut(i),
94 );
95 }
96 }
97 }
98 }
99
100 /// Updates a to floor(a / q[b.level()]).
101 /// Expects a to be in the NTT domain.
102 pub fn div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(
103 &self,
104 buf: &mut PolyRNS<u64>,
105 a: &mut PolyRNS<u64>,
106 ) {
107 debug_assert!(
108 self.level() <= a.level(),
109 "invalid input a: self.level()={} > a.level()={}",
110 self.level(),
111 a.level()
112 );
113
114 let level = self.level();
115 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
116 let r_last: &Ring<u64> = &self.0[level];
117
118 if ROUND{
119
120 let q_level_half: u64 = r_last.modulus.q >> 1;
121 let (buf_q_scaling, buf_qi_scaling) = buf.0.split_at_mut(1);
122
123 if NTT {
124 r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
125 r_last.a_add_b_scalar_into_a::<ONCE>(&q_level_half, &mut buf_q_scaling[0]);
126 for (i, r) in self.0[0..level].iter().enumerate() {
127 r_last.a_add_b_scalar_into_c::<NONE>(
128 &buf_q_scaling[0],
129 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
130 &mut buf_qi_scaling[0],
131 );
132 r.ntt_inplace::<false>(&mut buf_qi_scaling[0]);
133 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
134 &buf_qi_scaling[0],
135 &rescaling_constants.0[i],
136 a.at_mut(i),
137 );
138 }
139 } else {
140 let (a_qi, a_q_last) = a.0.split_at_mut(self.level());
141 r_last.a_add_b_scalar_into_a::<ONCE>(&q_level_half, &mut a_q_last[0]);
142 for (i, r) in self.0[0..level].iter().enumerate() {
143 r.b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a::<1, ONCE>(
144 &a_q_last[0],
145 &(r.modulus.q - r_last.modulus.barrett.reduce::<BARRETT>(&q_level_half)),
146 &rescaling_constants.0[i],
147 &mut a_qi[i],
148 );
149 }
150 }
151 }else{
152
153 if NTT {
154 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
155 r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
156 for (i, r) in self.0[0..level].iter().enumerate() {
157 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
158 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
159 &buf_ntt_qi_scaling[0],
160 &rescaling_constants.0[i],
161 a.at_mut(i),
162 );
163 }
164 }else{
165 let (a_i, a_level) = a.0.split_at_mut(level);
166 for (i, r) in self.0[0..level].iter().enumerate() {
167 r.b_sub_a_mul_c_scalar_barrett_into_a::<2, ONCE>(
168 &a_level[0],
169 &rescaling_constants.0[i],
170 &mut a_i[i],
171 );
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
217 if NTT{
218 self.intt::<false>(a, buf);
219 (0..nb_moduli).for_each(|i| {
220 self.at_level(self.level() - i)
221 .div_by_last_modulus_inplace::<ROUND, false>(
222 &mut PolyRNS::<u64>::default(),
223 buf,
224 )
225 });
226 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c);
227 }else{
228
229 println!("{} {:?}", self.level(), buf.level());
230 self.div_by_last_modulus::<ROUND, false>(a, buf, c);
231
232 (1..nb_moduli-1).for_each(|i| {
233 println!("{} {:?}", self.level() - i, buf.level());
234 self.at_level(self.level() - i)
235 .div_by_last_modulus_inplace::<ROUND, false>(buf, c);
236 });
237
238 self.at_level(self.level()-nb_moduli+1).div_by_last_modulus_inplace::<ROUND, false>(buf, c);
239 }
240 }
241 }
242
243 /// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i])
244 pub fn div_by_last_moduli_inplace<const ROUND:bool, const NTT: bool>(
245 &self,
246 nb_moduli: usize,
247 buf: &mut PolyRNS<u64>,
248 a: &mut PolyRNS<u64>,
249 ) {
250 debug_assert!(
251 self.level() <= a.level(),
252 "invalid input a: self.level()={} > a.level()={}",
253 self.level(),
254 a.level()
255 );
256 debug_assert!(
257 nb_moduli <= a.level(),
258 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
259 nb_moduli,
260 a.level()
261 );
262 if nb_moduli == 0{
263 return
264 }
265
266 if NTT {
267 self.intt::<false>(a, buf);
268 (0..nb_moduli).for_each(|i| {
269 self.at_level(self.level() - i)
270 .div_by_last_modulus_inplace::<ROUND, false>(&mut PolyRNS::<u64>::default(), buf)
271 });
272 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, a);
273 } else {
274 (0..nb_moduli).for_each(|i| {
275 self.at_level(self.level() - i)
276 .div_by_last_modulus_inplace::<ROUND, false>(buf, a)
277 });
278 }
279 }
280}