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::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_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::<false>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
38 r.sum_aqqmb_prod_c_scalar_barrett::<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.sum_aqqmb_prod_c_scalar_barrett::<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::<true>(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.sum_aqqmb_prod_c_scalar_barrett_inplace::<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) = buf.0.split_at_mut(level);
87 for (i, r) in self.0[0..level].iter().enumerate() {
88 r.sum_aqqmb_prod_c_scalar_barrett_inplace::<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 debug_assert!(
106 self.level() <= a.level(),
107 "invalid input a: self.level()={} > a.level()={}",
108 self.level(),
109 a.level()
110 );
111 debug_assert!(
112 c.level() >= a.level() - 1,
113 "invalid input b: b.level()={} < a.level()-1={}",
114 c.level(),
115 a.level() - 1
116 );
117 debug_assert!(
118 nb_moduli <= a.level(),
119 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
120 nb_moduli,
121 a.level()
122 );
123
124 if nb_moduli == 0 {
125 if a != c {
126 c.copy(a);
127 }
128 } else {
129 if NTT {
130 self.intt::<false>(a, buf);
131 (0..nb_moduli).for_each(|i| {
132 self.at_level(self.level() - i)
133 .div_floor_by_last_modulus_inplace::<false>(
134 &mut PolyRNS::<u64>::default(),
135 buf,
136 )
137 });
138 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c);
139 } else {
140 self.div_floor_by_last_modulus::<false>(a, buf, c);
141 (1..nb_moduli).for_each(|i| {
142 self.at_level(self.level() - i)
143 .div_floor_by_last_modulus_inplace::<false>(buf, c)
144 });
145 }
146 }
147 }
148
149 /// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i])
150 pub fn div_floor_by_last_moduli_inplace<const NTT: bool>(
151 &self,
152 nb_moduli: usize,
153 buf: &mut PolyRNS<u64>,
154 a: &mut PolyRNS<u64>,
155 ) {
156 debug_assert!(
157 self.level() <= a.level(),
158 "invalid input a: self.level()={} > a.level()={}",
159 self.level(),
160 a.level()
161 );
162 debug_assert!(
163 nb_moduli <= a.level(),
164 "invalid input nb_moduli: nb_moduli={} > a.level()={}",
165 nb_moduli,
166 a.level()
167 );
168 if NTT {
169 self.intt::<false>(a, buf);
170 (0..nb_moduli).for_each(|i| {
171 self.at_level(self.level() - i)
172 .div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)
173 });
174 self.at_level(self.level() - nb_moduli).ntt::<false>(buf, a);
175 } else {
176 (0..nb_moduli).for_each(|i| {
177 self.at_level(self.level() - i)
178 .div_floor_by_last_modulus_inplace::<false>(buf, a)
179 });
180 }
181 }
182
183 /// Updates b to round(a / q[b.level()]).
184 /// Expects b to be in the NTT domain.
185 pub fn div_round_by_last_modulus<const NTT: bool>(
186 &self,
187 a: &PolyRNS<u64>,
188 buf: &mut PolyRNS<u64>,
189 b: &mut PolyRNS<u64>,
190 ) {
191 debug_assert!(
192 self.level() <= a.level(),
193 "invalid input a: self.level()={} > a.level()={}",
194 self.level(),
195 a.level()
196 );
197 debug_assert!(
198 b.level() >= a.level() - 1,
199 "invalid input b: b.level()={} < a.level()-1={}",
200 b.level(),
201 a.level() - 1
202 );
203
204 let level: usize = self.level();
205 let r_last: &Ring<u64> = &self.0[level];
206 let q_level_half: u64 = r_last.modulus.q >> 1;
207 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
208 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
209
210 if NTT {
211 r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
212 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
213 for (i, r) in self.0[0..level].iter().enumerate() {
214 r_last.add_scalar::<ONCE>(
215 &buf_ntt_q_scaling[0],
216 &q_level_half,
217 &mut buf_ntt_qi_scaling[0],
218 );
219 r.ntt_inplace::<false>(&mut buf_ntt_qi_scaling[0]);
220 r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(
221 &buf_ntt_qi_scaling[0],
222 a.at(i),
223 &rescaling_constants.0[i],
224 b.at_mut(i),
225 );
226 }
227 } else {
228 }
229 }
230
231 /// Updates a to round(a / q[b.level()]).
232 /// Expects a to be in the NTT domain.
233 pub fn div_round_by_last_modulus_inplace<const NTT: bool>(
234 &self,
235 buf: &mut PolyRNS<u64>,
236 a: &mut PolyRNS<u64>,
237 ) {
238 debug_assert!(
239 self.level() <= a.level(),
240 "invalid input a: self.level()={} > a.level()={}",
241 self.level(),
242 a.level()
243 );
244
245 let level = self.level();
246 let r_last: &Ring<u64> = &self.0[level];
247 let q_level_half: u64 = r_last.modulus.q >> 1;
248 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
249 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
250
251 if NTT {
252 r_last.intt::<true>(a.at(level), &mut buf_ntt_q_scaling[0]);
253 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
254 for (i, r) in self.0[0..level].iter().enumerate() {
255 r_last.add_scalar::<ONCE>(
256 &buf_ntt_q_scaling[0],
257 &q_level_half,
258 &mut buf_ntt_qi_scaling[0],
259 );
260 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
261 r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(
262 &buf_ntt_qi_scaling[0],
263 &rescaling_constants.0[i],
264 a.at_mut(i),
265 );
266 }
267 }
268 }
269}