poulpy/poulpy-schemes/src/bin_fhe/bdd_arithmetic/eval.rs
poulpy-schemes/src/bin_fhe/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::{
8 GGSWInfos, GGSWPrepared, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, LWEInfos, prepared::GGSWPreparedToRef,
9 },
10};
11use poulpy_hal::{
12 api::{
13 ScratchAvailable, ScratchTakeBasic, VecZnxBigAddSmall, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize,
14 VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallA, VecZnxDftBytesOf,
15 },
16 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig, ZnxInfos, ZnxZero},
17};
18
19use crate::bin_fhe::bdd_arithmetic::GetGGSWBit;
20
21pub trait BitCircuitInfo: Sync {
22 fn info(&self) -> (&[Node], usize);
23}
24
25pub trait GetBitCircuitInfo: Sync {
26 fn input_size(&self) -> usize;
27 fn output_size(&self) -> usize;
28 fn get_circuit(&self, bit: usize) -> (&[Node], usize);
29}
30
31pub struct BitCircuit<const N: usize> {
32 pub nodes: [Node; N],
33 pub max_inter_state: usize,
34}
35
36pub trait BitCircuitFamily {
37 const INPUT_BITS: usize;
38 const OUTPUT_BITS: usize;
39}
40
41pub struct Circuit<C: BitCircuitInfo, const N: usize>(pub [C; N]);
42
43impl<C, const N: usize> GetBitCircuitInfo for Circuit<C, N>
44where
45 C: BitCircuitInfo + BitCircuitFamily,
46{
47 fn input_size(&self) -> usize {
48 C::INPUT_BITS
49 }
50 fn output_size(&self) -> usize {
51 C::OUTPUT_BITS
52 }
53 fn get_circuit(&self, bit: usize) -> (&[Node], usize) {
54 self.0[bit].info()
55 }
56}
57
58pub trait ExecuteBDDCircuit<BE: Backend> {
59 fn execute_bdd_circuit_tmp_bytes<R, G>(&self, res_infos: &R, state_size: usize, ggsw_infos: &G) -> usize
60 where
61 R: GLWEInfos,
62 G: GGSWInfos;
63
64 fn execute_bdd_circuit<C, G, O>(&self, out: &mut [GLWE<O>], inputs: &G, circuit: &C, scratch: &mut Scratch<BE>)
65 where
66 G: GetGGSWBit<BE> + BitSize,
67 C: GetBitCircuitInfo,
68 O: DataMut,
69 {
70 self.execute_bdd_circuit_multi_thread(1, out, inputs, circuit, scratch);
71 }
72
73 fn execute_bdd_circuit_multi_thread<C, G, O>(
74 &self,
75 threads: usize,
76 out: &mut [GLWE<O>],
77 inputs: &G,
78 circuit: &C,
79 scratch: &mut Scratch<BE>,
80 ) where
81 G: GetGGSWBit<BE> + BitSize,
82 C: GetBitCircuitInfo,
83 O: DataMut;
84}
85
86pub trait BitSize {
87 fn bit_size(&self) -> usize;
88}
89
90impl<BE: Backend> ExecuteBDDCircuit<BE> for Module<BE>
91where
92 Self: Cmux<BE> + GLWECopy,
93 Scratch<BE>: ScratchTakeCore<BE>,
94{
95 fn execute_bdd_circuit_tmp_bytes<R, G>(&self, res_infos: &R, state_size: usize, ggsw_infos: &G) -> usize
96 where
97 R: GLWEInfos,
98 G: GGSWInfos,
99 {
100 2 * state_size * GLWE::bytes_of_from_infos(res_infos) + self.cmux_tmp_bytes(res_infos, res_infos, ggsw_infos)
101 }
102
103 fn execute_bdd_circuit_multi_thread<C, G, O>(
104 &self,
105 threads: usize,
106 out: &mut [GLWE<O>],
107 inputs: &G,
108 circuit: &C,
109 scratch: &mut Scratch<BE>,
110 ) where
111 G: GetGGSWBit<BE> + BitSize,
112 C: GetBitCircuitInfo,
113 O: DataMut,
114 {
115 #[cfg(debug_assertions)]
116 {
117 assert!(
118 inputs.bit_size() >= circuit.input_size(),
119 "inputs.bit_size(): {} < circuit.input_size():{}",
120 inputs.bit_size(),
121 circuit.input_size()
122 );
123 assert!(
124 out.len() >= circuit.output_size(),
125 "out.len(): {} < circuit.output_size(): {}",
126 out.len(),
127 circuit.output_size()
128 );
129 }
130
131 let mut max_state_size = 0;
132 for i in 0..circuit.output_size() {
133 let (_, state_size) = circuit.get_circuit(i);
134 max_state_size = max_state_size.max(state_size)
135 }
136
137 let scratch_thread_size: usize = self.execute_bdd_circuit_tmp_bytes(&out[0], max_state_size, &inputs.get_bit(0));
138
139 assert!(
140 scratch.available() >= threads * scratch_thread_size,
141 "scratch.available(): {} < threads:{threads} * scratch_thread_size: {scratch_thread_size}",
142 scratch.available()
143 );
144
145 let (mut scratches, _) = scratch.split_mut(threads, scratch_thread_size);
146
147 let chunk_size: usize = circuit.output_size().div_ceil(threads);
148
149 thread::scope(|scope| {
150 for (thread_idx, (scratch_thread, out_chunk)) in scratches
151 .iter_mut()
152 .zip(out[..circuit.output_size()].chunks_mut(chunk_size))
153 .enumerate()
154 {
155 // Capture chunk + thread scratch by move
156 scope.spawn(move || {
157 for (idx, out_i) in out_chunk.iter_mut().enumerate() {
158 let (nodes, state_size) = circuit.get_circuit(thread_idx * chunk_size + idx);
159
160 if state_size == 0 {
161 out_i.data_mut().zero();
162 } else {
163 eval_level(self, out_i, inputs, nodes, state_size, *scratch_thread);
164 }
165 }
166 });
167 }
168 });
169
170 for out_i in out.iter_mut().skip(circuit.output_size()) {
171 out_i.data_mut().zero();
172 }
173 }
174}
175
176fn eval_level<M, R, G, BE: Backend>(
177 module: &M,
178 res: &mut R,
179 inputs: &G,
180 nodes: &[Node],
181 state_size: usize,
182 scratch: &mut Scratch<BE>,
183) where
184 M: Cmux<BE> + GLWECopy,
185 R: GLWEToMut,
186 G: GetGGSWBit<BE> + BitSize,
187 Scratch<BE>: ScratchTakeCore<BE>,
188{
189 assert!(nodes.len().is_multiple_of(state_size));
190 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
191
192 let (mut level, scratch_1) = scratch.take_glwe_slice(state_size * 2, res);
193
194 level.iter_mut().for_each(|ct| ct.data_mut().zero());
195
196 // TODO: implement API on GLWE
197 level[1].data_mut().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 { nodes, max_inter_state }
245 }
246}
247impl<const N: usize> BitCircuitInfo for BitCircuit<N> {
248 fn info(&self) -> (&[Node], usize) {
249 (self.nodes.as_ref(), self.max_inter_state)
250 }
251}
252
253#[derive(Debug)]
254pub enum Node {
255 Cmux(usize, usize, usize),
256 Copy,
257 None,
258}
259
260impl<BE: Backend> Cswap<BE> for Module<BE> where
261 Self: Sized
262 + GLWEExternalProductInternal<BE>
263 + GLWESub
264 + VecZnxBigAddSmallInplace<BE>
265 + GLWENormalize<BE>
266 + VecZnxDftBytesOf
267 + VecZnxBigNormalize<BE>
268 + VecZnxBigNormalizeTmpBytes
269 + GLWENormalize<BE>
270 + VecZnxBigAddSmall<BE>
271 + VecZnxBigSubSmallA<BE>
272 + VecZnxBigBytesOf
273{
274}
275
276pub trait Cswap<BE: Backend>
277where
278 Self: Sized
279 + GLWEExternalProductInternal<BE>
280 + GLWESub
281 + GLWECopy
282 + VecZnxBigAddSmallInplace<BE>
283 + GLWENormalize<BE>
284 + VecZnxDftBytesOf
285 + VecZnxBigNormalize<BE>
286 + VecZnxBigNormalizeTmpBytes
287 + GLWENormalize<BE>
288 + VecZnxBigAddSmall<BE>
289 + VecZnxBigSubSmallA<BE>
290 + VecZnxBigBytesOf,
291{
292 fn cswap_tmp_bytes<R, A, S>(&self, res_a_infos: &R, res_b_infos: &A, s_infos: &S) -> usize
293 where
294 R: GLWEInfos,
295 A: GLWEInfos,
296 S: GGSWInfos,
297 {
298 let res_dft: usize = self.bytes_of_vec_znx_dft((s_infos.rank() + 1).into(), s_infos.size());
299 let mut tot = res_dft
300 + (self.glwe_external_product_internal_tmp_bytes(res_a_infos, res_b_infos, s_infos)
301 + GLWE::bytes_of_from_infos(&GLWELayout {
302 n: s_infos.n(),
303 base2k: s_infos.base2k(),
304 k: res_a_infos.k().max(res_b_infos.k()),
305 rank: s_infos.rank(),
306 }))
307 .max(self.vec_znx_big_normalize_tmp_bytes());
308
309 if res_a_infos.base2k() != s_infos.base2k() {
310 tot += GLWE::bytes_of_from_infos(&GLWELayout {
311 n: res_a_infos.n(),
312 base2k: s_infos.base2k(),
313 k: res_a_infos.k(),
314 rank: res_a_infos.rank(),
315 });
316
317 tot += GLWE::bytes_of_from_infos(&GLWELayout {
318 n: res_b_infos.n(),
319 base2k: s_infos.base2k(),
320 k: res_b_infos.k(),
321 rank: res_b_infos.rank(),
322 });
323 }
324
325 tot += self.bytes_of_vec_znx_big(1, s_infos.size());
326
327 tot
328 }
329
330 fn cswap<A, B, S>(&self, res_a: &mut A, res_b: &mut B, s: &S, scratch: &mut Scratch<BE>)
331 where
332 A: GLWEToMut,
333 B: GLWEToMut,
334 S: GGSWPreparedToRef<BE> + GGSWInfos,
335 Scratch<BE>: ScratchTakeCore<BE>,
336 {
337 let res_a: &mut GLWE<&mut [u8]> = &mut res_a.to_mut();
338 let res_b: &mut GLWE<&mut [u8]> = &mut res_b.to_mut();
339 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
340 assert_eq!(res_a.base2k(), res_b.base2k());
341
342 let res_base2k: usize = res_a.base2k().as_usize();
343 let s_base2k: usize = s.base2k().as_usize();
344
345 if res_base2k == s_base2k {
346 let res_big: VecZnxBig<&mut [u8], BE>;
347 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (s.rank() + 1).into(), s.size()); // Todo optimise
348 {
349 // Temporary value storing a - b
350 let tmp_c_infos: GLWELayout = GLWELayout {
351 n: s.n(),
352 base2k: s.base2k(),
353 k: res_a.k().max(res_b.k()),
354 rank: s.rank(),
355 };
356 let (mut tmp_c, scratch_2) = scratch_1.take_glwe(&tmp_c_infos);
357 self.glwe_sub(&mut tmp_c, res_b, res_a);
358 res_big = self.glwe_external_product_internal(res_dft, &tmp_c, s, scratch_2);
359 }
360
361 // Single column res_big to store temporary value before normalization
362 let (mut res_big_tmp, scratch_2) = scratch_1.take_vec_znx_big::<_, BE>(self, 1, res_big.size());
363
364 // res_a = (b-a) * bit + a
365 for j in 0..(res_a.rank() + 1).into() {
366 self.vec_znx_big_add_small(&mut res_big_tmp, 0, &res_big, j, res_a.data(), j);
367 self.vec_znx_big_normalize(res_a.data_mut(), res_base2k, 0, j, &res_big_tmp, s_base2k, 0, scratch_2);
368 }
369
370 // res_b = a - (a - b) * bit = (b - a) * bit + a
371 for j in 0..(res_b.rank() + 1).into() {
372 self.vec_znx_big_sub_small_a(&mut res_big_tmp, 0, res_b.data(), j, &res_big, j);
373 self.vec_znx_big_normalize(res_b.data_mut(), res_base2k, 0, j, &res_big_tmp, s_base2k, 0, scratch_2);
374 }
375 } else {
376 let (mut tmp_a, scratch_1) = scratch.take_glwe(&GLWELayout {
377 n: res_a.n(),
378 base2k: s.base2k(),
379 k: res_a.k(),
380 rank: res_a.rank(),
381 });
382
383 let (mut tmp_b, scratch_2) = scratch_1.take_glwe(&GLWELayout {
384 n: res_b.n(),
385 base2k: s.base2k(),
386 k: res_b.k(),
387 rank: res_b.rank(),
388 });
389
390 self.glwe_normalize(&mut tmp_a, res_a, scratch_2);
391 self.glwe_normalize(&mut tmp_b, res_b, scratch_2);
392
393 let res_big: VecZnxBig<&mut [u8], BE>;
394 let (res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, (s.rank() + 1).into(), s.size()); // Todo optimise
395 {
396 // Temporary value storing a - b
397 let tmp_c_infos: GLWELayout = GLWELayout {
398 n: s.n(),
399 base2k: s.base2k(),
400 k: res_a.k().max(res_b.k()),
401 rank: s.rank(),
402 };
403 let (mut tmp_c, scratch_4) = scratch_3.take_glwe(&tmp_c_infos);
404 self.glwe_sub(&mut tmp_c, res_b, res_a);
405 res_big = self.glwe_external_product_internal(res_dft, &tmp_c, s, scratch_4);
406 }
407
408 // Single column res_big to store temporary value before normalization
409 let (mut res_big_tmp, scratch_4) = scratch_3.take_vec_znx_big::<_, BE>(self, 1, res_big.size());
410
411 // res_a = (b-a) * bit + a
412 for j in 0..(res_a.rank() + 1).into() {
413 self.vec_znx_big_add_small(&mut res_big_tmp, 0, &res_big, j, tmp_a.data(), j);
414 self.vec_znx_big_normalize(res_a.data_mut(), res_base2k, 0, j, &res_big_tmp, s_base2k, 0, scratch_4);
415 }
416
417 // res_b = a - (a - b) * bit = (b - a) * bit + a
418 for j in 0..(res_b.rank() + 1).into() {
419 self.vec_znx_big_sub_small_a(&mut res_big_tmp, 0, tmp_b.data(), j, &res_big, j);
420 self.vec_znx_big_normalize(res_b.data_mut(), res_base2k, 0, j, &res_big_tmp, s_base2k, 0, scratch_4);
421 }
422 }
423 }
424}
425
426pub trait Cmux<BE: Backend>
427where
428 Self: Sized
429 + GLWEExternalProductInternal<BE>
430 + GLWESub
431 + VecZnxBigAddSmallInplace<BE>
432 + GLWENormalize<BE>
433 + VecZnxDftBytesOf
434 + VecZnxBigNormalize<BE>
435 + VecZnxBigNormalizeTmpBytes,
436{
437 fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, selector_infos: &B) -> usize
438 where
439 R: GLWEInfos,
440 A: GLWEInfos,
441 B: GGSWInfos,
442 {
443 let res_dft: usize = self.bytes_of_vec_znx_dft((selector_infos.rank() + 1).into(), selector_infos.size());
444 res_dft
445 + self
446 .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, selector_infos)
447 .max(self.vec_znx_big_normalize_tmp_bytes())
448 }
449
450 // res = (t - f) * s + f
451 fn cmux<R, T, F, S>(&self, res: &mut R, t: &T, f: &F, s: &S, scratch: &mut Scratch<BE>)
452 where
453 R: GLWEToMut,
454 T: GLWEToRef,
455 F: GLWEToRef,
456 S: GGSWPreparedToRef<BE> + GGSWInfos,
457 Scratch<BE>: ScratchTakeCore<BE>,
458 {
459 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
460 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
461 let f: GLWE<&[u8]> = f.to_ref();
462
463 let res_base2k: usize = res.base2k().into();
464 let ggsw_base2k: usize = s.base2k().into();
465
466 self.glwe_sub(res, t, &f);
467 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
468 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
469 for j in 0..(res.rank() + 1).into() {
470 self.vec_znx_big_add_small_inplace(&mut res_big, j, f.data(), j);
471 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, j, &res_big, ggsw_base2k, j, scratch_1);
472 }
473 }
474
475 // res = (a - res) * s + res
476 fn cmux_inplace_neg<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
477 where
478 R: GLWEToMut,
479 A: GLWEToRef,
480 S: GGSWPreparedToRef<BE> + GGSWInfos,
481 Scratch<BE>: ScratchTakeCore<BE>,
482 {
483 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
484 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
485 let a: &GLWE<&[u8]> = &a.to_ref();
486
487 assert_eq!(res.base2k(), a.base2k());
488
489 let res_base2k: usize = res.base2k().into();
490 let ggsw_base2k: usize = s.base2k().into();
491 let (mut tmp, scratch_1) = scratch.take_glwe(&GLWELayout {
492 n: s.n(),
493 base2k: res.base2k(),
494 k: res.k().max(a.k()),
495 rank: res.rank(),
496 });
497 self.glwe_sub(&mut tmp, a, res);
498 let (res_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
499 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, &tmp, s, scratch_2);
500 for j in 0..(res.rank() + 1).into() {
501 self.vec_znx_big_add_small_inplace(&mut res_big, j, res.data(), j);
502 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, j, &res_big, ggsw_base2k, j, scratch_2);
503 }
504 }
505
506 // res = (res - a) * s + a
507 fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
508 where
509 R: GLWEToMut,
510 A: GLWEToRef,
511 S: GGSWPreparedToRef<BE> + GGSWInfos,
512 Scratch<BE>: ScratchTakeCore<BE>,
513 {
514 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
515 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
516 let a: GLWE<&[u8]> = a.to_ref();
517 let res_base2k: usize = res.base2k().into();
518 let ggsw_base2k: usize = s.base2k().into();
519 self.glwe_sub_inplace(res, &a);
520 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
521 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
522 for j in 0..(res.rank() + 1).into() {
523 self.vec_znx_big_add_small_inplace(&mut res_big, j, a.data(), j);
524 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, j, &res_big, ggsw_base2k, j, scratch_1);
525 }
526 }
527}
528
529impl<BE: Backend> Cmux<BE> for Module<BE>
530where
531 Self: Sized
532 + GLWEExternalProductInternal<BE>
533 + GLWESub
534 + VecZnxBigAddSmallInplace<BE>
535 + GLWENormalize<BE>
536 + VecZnxDftBytesOf
537 + VecZnxBigNormalize<BE>
538 + VecZnxBigNormalizeTmpBytes,
539 Scratch<BE>: ScratchTakeCore<BE>,
540{
541}