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