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