poulpy/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs
poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rsRaw
1use core::panic;
2
3use itertools::Itertools;
4use poulpy_core::{
5 GLWECopy, GLWEExternalProductInternal, GLWENormalize, GLWESub, ScratchTakeCore,
6 layouts::{GGSWInfos, GGSWPrepared, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::GGSWPreparedToRef},
7};
8use poulpy_hal::{
9 api::{ScratchTakeBasic, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftBytesOf},
10 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig, ZnxZero},
11};
12
13use crate::tfhe::bdd_arithmetic::{GetGGSWBit, UnsignedInteger};
14
15pub trait BitCircuitInfo {
16 fn info(&self) -> (&[Node], usize);
17}
18
19pub trait GetBitCircuitInfo<T: UnsignedInteger> {
20 fn input_size(&self) -> usize;
21 fn output_size(&self) -> usize;
22 fn get_circuit(&self, bit: usize) -> (&[Node], usize);
23}
24
25pub struct BitCircuit<const N: usize> {
26 pub nodes: [Node; N],
27 pub max_inter_state: usize,
28}
29
30pub trait BitCircuitFamily {
31 const INPUT_BITS: usize;
32 const OUTPUT_BITS: usize;
33}
34
35pub struct Circuit<C: BitCircuitInfo, const N: usize>(pub [C; N]);
36
37impl<C, T: UnsignedInteger, const N: usize> GetBitCircuitInfo<T> for Circuit<C, N>
38where
39 C: BitCircuitInfo + BitCircuitFamily,
40{
41 fn input_size(&self) -> usize {
42 C::INPUT_BITS
43 }
44 fn output_size(&self) -> usize {
45 C::OUTPUT_BITS
46 }
47 fn get_circuit(&self, bit: usize) -> (&[Node], usize) {
48 self.0[bit].info()
49 }
50}
51
52pub trait ExecuteBDDCircuit<T: UnsignedInteger, BE: Backend> {
53 fn execute_bdd_circuit<C, G, O>(&self, out: &mut [GLWE<O>], inputs: &G, circuit: &C, scratch: &mut Scratch<BE>)
54 where
55 G: GetGGSWBit<BE> + BitSize,
56 C: GetBitCircuitInfo<T>,
57 O: DataMut;
58}
59
60pub trait BitSize {
61 fn bit_size(&self) -> usize;
62}
63
64impl<T: UnsignedInteger, BE: Backend> ExecuteBDDCircuit<T, BE> for Module<BE>
65where
66 Self: Cmux<BE> + GLWECopy,
67 Scratch<BE>: ScratchTakeCore<BE>,
68{
69 fn execute_bdd_circuit<C, G, O>(&self, out: &mut [GLWE<O>], inputs: &G, circuit: &C, scratch: &mut Scratch<BE>)
70 where
71 G: GetGGSWBit<BE> + BitSize,
72 C: GetBitCircuitInfo<T>,
73 O: DataMut,
74 {
75 #[cfg(debug_assertions)]
76 {
77 assert!(
78 inputs.bit_size() >= circuit.input_size(),
79 "inputs.bit_size(): {} < circuit.input_size():{}",
80 inputs.bit_size(),
81 circuit.input_size()
82 );
83 assert!(
84 out.len() >= circuit.output_size(),
85 "out.len(): {} < circuit.output_size(): {}",
86 out.len(),
87 circuit.output_size()
88 );
89 }
90
91 for (i, out_i) in out.iter_mut().enumerate().take(circuit.output_size()) {
92 let (nodes, max_inter_state) = circuit.get_circuit(i);
93
94 if max_inter_state == 0 {
95 out_i.data_mut().zero();
96 } else {
97 assert!(nodes.len().is_multiple_of(max_inter_state));
98
99 let (mut level, scratch_1) = scratch.take_glwe_slice(max_inter_state * 2, out_i);
100
101 level.iter_mut().for_each(|ct| ct.data_mut().zero());
102
103 // TODO: implement API on GLWE
104 level[1]
105 .data_mut()
106 .encode_coeff_i64(out_i.base2k().into(), 0, 2, 0, 1);
107
108 let mut level_ref = level.iter_mut().collect_vec();
109 let (mut prev_level, mut next_level) = level_ref.split_at_mut(max_inter_state);
110
111 let (all_but_last, last) = nodes.split_at(nodes.len() - max_inter_state);
112
113 for nodes_lvl in all_but_last.chunks_exact(max_inter_state) {
114 for (j, node) in nodes_lvl.iter().enumerate() {
115 match node {
116 Node::Cmux(in_idx, hi_idx, lo_idx) => {
117 self.cmux(
118 next_level[j],
119 prev_level[*hi_idx],
120 prev_level[*lo_idx],
121 &inputs.get_bit(*in_idx),
122 scratch_1,
123 );
124 }
125 Node::Copy => self.glwe_copy(next_level[j], prev_level[j]), /* Update BDD circuits to order Cmux -> Copy -> None so that mem swap can be used */
126 Node::None => {}
127 }
128 }
129
130 (prev_level, next_level) = (next_level, prev_level);
131 }
132
133 // Last chunck of max_inter_state Nodes is always structured as
134 // [CMUX, NONE, NONE, ..., NONE]
135 match &last[0] {
136 Node::Cmux(in_idx, hi_idx, lo_idx) => {
137 self.cmux(
138 out_i,
139 prev_level[*hi_idx],
140 prev_level[*lo_idx],
141 &inputs.get_bit(*in_idx),
142 scratch_1,
143 );
144 }
145 _ => {
146 panic!("invalid last node, should be CMUX")
147 }
148 }
149 }
150 }
151
152 for out_i in out.iter_mut().skip(circuit.output_size()) {
153 out_i.data_mut().zero();
154 }
155 }
156}
157
158impl<const N: usize> BitCircuit<N> {
159 pub const fn new(nodes: [Node; N], max_inter_state: usize) -> Self {
160 Self {
161 nodes,
162 max_inter_state,
163 }
164 }
165}
166impl<const N: usize> BitCircuitInfo for BitCircuit<N> {
167 fn info(&self) -> (&[Node], usize) {
168 (self.nodes.as_ref(), self.max_inter_state)
169 }
170}
171
172#[derive(Debug)]
173pub enum Node {
174 Cmux(usize, usize, usize),
175 Copy,
176 None,
177}
178
179pub trait Cmux<BE: Backend>
180where
181 Self: Sized
182 + GLWEExternalProductInternal<BE>
183 + GLWESub
184 + VecZnxBigAddSmallInplace<BE>
185 + GLWENormalize<BE>
186 + VecZnxDftBytesOf
187 + VecZnxBigNormalize<BE>
188 + VecZnxBigNormalizeTmpBytes,
189{
190 fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
191 where
192 R: GLWEInfos,
193 A: GLWEInfos,
194 B: GGSWInfos,
195 {
196 let res_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), b_infos.size());
197 res_dft
198 + self
199 .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, b_infos)
200 .max(self.vec_znx_big_normalize_tmp_bytes())
201 }
202
203 fn cmux<R, T, F, S>(&self, res: &mut R, t: &T, f: &F, s: &S, scratch: &mut Scratch<BE>)
204 where
205 R: GLWEToMut,
206 T: GLWEToRef,
207 F: GLWEToRef,
208 S: GGSWPreparedToRef<BE> + GGSWInfos,
209 Scratch<BE>: ScratchTakeCore<BE>,
210 {
211 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
212 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
213 let f: GLWE<&[u8]> = f.to_ref();
214
215 let res_base2k: usize = res.base2k().into();
216 let ggsw_base2k: usize = s.base2k().into();
217
218 self.glwe_sub(res, t, &f);
219 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
220 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
221 for j in 0..(res.rank() + 1).into() {
222 self.vec_znx_big_add_small_inplace(&mut res_big, j, f.data(), j);
223 self.vec_znx_big_normalize(
224 res_base2k,
225 res.data_mut(),
226 j,
227 ggsw_base2k,
228 &res_big,
229 j,
230 scratch_1,
231 );
232 }
233 }
234
235 fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
236 where
237 R: GLWEToMut,
238 A: GLWEToRef,
239 S: GGSWPreparedToRef<BE> + GGSWInfos,
240 Scratch<BE>: ScratchTakeCore<BE>,
241 {
242 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
243 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
244 let a: GLWE<&[u8]> = a.to_ref();
245 let res_base2k: usize = res.base2k().into();
246 let ggsw_base2k: usize = s.base2k().into();
247 self.glwe_sub_inplace(res, &a);
248 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
249 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
250 for j in 0..(res.rank() + 1).into() {
251 self.vec_znx_big_add_small_inplace(&mut res_big, j, a.data(), j);
252 self.vec_znx_big_normalize(
253 res_base2k,
254 res.data_mut(),
255 j,
256 ggsw_base2k,
257 &res_big,
258 j,
259 scratch_1,
260 );
261 }
262 }
263}
264
265impl<BE: Backend> Cmux<BE> for Module<BE>
266where
267 Self: Sized
268 + GLWEExternalProductInternal<BE>
269 + GLWESub
270 + VecZnxBigAddSmallInplace<BE>
271 + GLWENormalize<BE>
272 + VecZnxDftBytesOf
273 + VecZnxBigNormalize<BE>
274 + VecZnxBigNormalizeTmpBytes,
275 Scratch<BE>: ScratchTakeCore<BE>,
276{
277}