poulpy/math/src/ring/impl_u64/rescaling_rns.rs
math/src/ring/impl_u64/rescaling_rns.rsRaw
1use crate::ring::Ring;
2use crate::ring::RingRNS;
3use crate::poly::PolyRNS;
4use crate::modulus::barrett::Barrett;
5use crate::scalar::ScalarRNS;
6use crate::modulus::ONCE;
7extern crate test;
8
9impl RingRNS<'_, u64>{
10
11 /// Updates b to floor(a / q[b.level()]).
12 pub fn div_floor_by_last_modulus<const NTT:bool>(&self, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>){
13 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
14 debug_assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
15
16 let level = self.level();
17 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
18
19 if NTT{
20 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
21 self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
22 for (i, r) in self.0[0..level].iter().enumerate(){
23 r.ntt::<false>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
24 r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(&buf_ntt_qi_scaling[0], a.at(i), &rescaling_constants.0[i], b.at_mut(i));
25 }
26 }else{
27 for (i, r) in self.0[0..level].iter().enumerate(){
28 r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(a.at(level), a.at(i), &rescaling_constants.0[i], b.at_mut(i));
29 }
30 }
31 }
32
33 /// Updates a to floor(a / q[b.level()]).
34 /// Expects a to be in the NTT domain.
35 pub fn div_floor_by_last_modulus_inplace<const NTT:bool>(&self, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
36 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
37
38 let level = self.level();
39 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
40
41 if NTT{
42 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
43 self.0[level].intt::<true>(a.at(level), &mut buf_ntt_q_scaling[0]);
44 for (i, r) in self.0[0..level].iter().enumerate(){
45 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
46 r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&buf_ntt_qi_scaling[0], &rescaling_constants.0[i], a.at_mut(i));
47 }
48 }else{
49 let (a_i, a_level) = buf.0.split_at_mut(level);
50 for (i, r) in self.0[0..level].iter().enumerate(){
51 r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&a_level[0], &rescaling_constants.0[i], &mut a_i[i]);
52 }
53 }
54 }
55
56 /// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i])
57 pub fn div_floor_by_last_moduli<const NTT:bool>(&self, nb_moduli:usize, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, c: &mut PolyRNS<u64>){
58
59 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
60 debug_assert!(c.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", c.level(), a.level()-1);
61 debug_assert!(nb_moduli <= a.level(), "invalid input nb_moduli: nb_moduli={} > a.level()={}", nb_moduli, a.level());
62
63 if nb_moduli == 0{
64 if a != c{
65 c.copy(a);
66 }
67 }else{
68 if NTT{
69 self.intt::<false>(a, buf);
70 (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)});
71 self.at_level(self.level()-nb_moduli).ntt::<false>(buf, c);
72 }else{
73 self.div_floor_by_last_modulus::<false>(a, buf, c);
74 (1..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(buf, c)});
75 }
76 }
77 }
78
79 /// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i])
80 pub fn div_floor_by_last_moduli_inplace<const NTT:bool>(&self, nb_moduli:usize, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
81 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
82 debug_assert!(nb_moduli <= a.level(), "invalid input nb_moduli: nb_moduli={} > a.level()={}", nb_moduli, a.level());
83 if NTT{
84 self.intt::<false>(a, buf);
85 (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)});
86 self.at_level(self.level()-nb_moduli).ntt::<false>(buf, a);
87 }else{
88 (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(buf, a)});
89 }
90 }
91
92 /// Updates b to round(a / q[b.level()]).
93 /// Expects b to be in the NTT domain.
94 pub fn div_round_by_last_modulus<const NTT:bool>(&self, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>){
95 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
96 debug_assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
97
98 let level: usize = self.level();
99 let r_last: &Ring<u64> = &self.0[level];
100 let q_level_half: u64 = r_last.modulus.q>>1;
101 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
102 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
103
104 if NTT{
105 r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
106 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
107 for (i, r) in self.0[0..level].iter().enumerate(){
108 r_last.add_scalar::<ONCE>(&buf_ntt_q_scaling[0], &q_level_half, &mut buf_ntt_qi_scaling[0]);
109 r.ntt_inplace::<false>(&mut buf_ntt_qi_scaling[0]);
110 r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(&buf_ntt_qi_scaling[0], a.at(i), &rescaling_constants.0[i], b.at_mut(i));
111 }
112 }else{
113
114 }
115
116 }
117
118 /// Updates a to round(a / q[b.level()]).
119 /// Expects a to be in the NTT domain.
120 pub fn div_round_by_last_modulus_inplace<const NTT:bool>(&self, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
121 debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
122
123 let level = self.level();
124 let r_last: &Ring<u64> = &self.0[level];
125 let q_level_half: u64 = r_last.modulus.q>>1;
126 let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
127 let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
128
129 if NTT{
130 r_last.intt::<true>(a.at(level), &mut buf_ntt_q_scaling[0]);
131 r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
132 for (i, r) in self.0[0..level].iter().enumerate(){
133 r_last.add_scalar::<ONCE>(&buf_ntt_q_scaling[0], &q_level_half, &mut buf_ntt_qi_scaling[0]);
134 r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
135 r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&buf_ntt_qi_scaling[0], &rescaling_constants.0[i], a.at_mut(i));
136 }
137 }
138
139 }
140}
141
142