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::{
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::tfhe::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]
198 .data_mut()
199 .encode_coeff_i64(res.base2k().into(), 0, 2, 0, 1);
200
201 let mut level_ref: Vec<&mut GLWE<&mut [u8]>> = level.iter_mut().collect_vec();
202 let (mut prev_level, mut next_level) = level_ref.split_at_mut(state_size);
203
204 let (all_but_last, last) = nodes.split_at(nodes.len() - state_size);
205
206 for nodes_lvl in all_but_last.chunks_exact(state_size) {
207 for (j, node) in nodes_lvl.iter().enumerate() {
208 match node {
209 Node::Cmux(in_idx, hi_idx, lo_idx) => {
210 module.cmux(
211 next_level[j],
212 prev_level[*hi_idx],
213 prev_level[*lo_idx],
214 &inputs.get_bit(*in_idx),
215 scratch_1,
216 );
217 }
218 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 */
219 Node::None => {}
220 }
221 }
222
223 (prev_level, next_level) = (next_level, prev_level);
224 }
225
226 // Last chunck of max_inter_state Nodes is always structured as
227 // [CMUX, NONE, NONE, ..., NONE]
228 match &last[0] {
229 Node::Cmux(in_idx, hi_idx, lo_idx) => {
230 module.cmux(
231 res,
232 prev_level[*hi_idx],
233 prev_level[*lo_idx],
234 &inputs.get_bit(*in_idx),
235 scratch_1,
236 );
237 }
238 _ => {
239 panic!("invalid last node, should be CMUX")
240 }
241 }
242}
243
244impl<const N: usize> BitCircuit<N> {
245 pub const fn new(nodes: [Node; N], max_inter_state: usize) -> Self {
246 Self {
247 nodes,
248 max_inter_state,
249 }
250 }
251}
252impl<const N: usize> BitCircuitInfo for BitCircuit<N> {
253 fn info(&self) -> (&[Node], usize) {
254 (self.nodes.as_ref(), self.max_inter_state)
255 }
256}
257
258#[derive(Debug)]
259pub enum Node {
260 Cmux(usize, usize, usize),
261 Copy,
262 None,
263}
264
265impl<BE: Backend> Cswap<BE> for Module<BE> where
266 Self: Sized
267 + GLWEExternalProductInternal<BE>
268 + GLWESub
269 + VecZnxBigAddSmallInplace<BE>
270 + GLWENormalize<BE>
271 + VecZnxDftBytesOf
272 + VecZnxBigNormalize<BE>
273 + VecZnxBigNormalizeTmpBytes
274 + GLWENormalize<BE>
275 + VecZnxBigAddSmall<BE>
276 + VecZnxBigSubSmallA<BE>
277 + VecZnxBigBytesOf
278{
279}
280
281pub trait Cswap<BE: Backend>
282where
283 Self: Sized
284 + GLWEExternalProductInternal<BE>
285 + GLWESub
286 + GLWECopy
287 + VecZnxBigAddSmallInplace<BE>
288 + GLWENormalize<BE>
289 + VecZnxDftBytesOf
290 + VecZnxBigNormalize<BE>
291 + VecZnxBigNormalizeTmpBytes
292 + GLWENormalize<BE>
293 + VecZnxBigAddSmall<BE>
294 + VecZnxBigSubSmallA<BE>
295 + VecZnxBigBytesOf,
296{
297 fn cswap_tmp_bytes<R, A, S>(&self, res_a_infos: &R, res_b_infos: &A, s_infos: &S) -> usize
298 where
299 R: GLWEInfos,
300 A: GLWEInfos,
301 S: GGSWInfos,
302 {
303 let res_dft: usize = self.bytes_of_vec_znx_dft((s_infos.rank() + 1).into(), s_infos.size());
304 let mut tot = res_dft
305 + (self.glwe_external_product_internal_tmp_bytes(res_a_infos, res_b_infos, s_infos)
306 + GLWE::bytes_of_from_infos(&GLWELayout {
307 n: s_infos.n(),
308 base2k: s_infos.base2k(),
309 k: res_a_infos.k().max(res_b_infos.k()),
310 rank: s_infos.rank(),
311 }))
312 .max(self.vec_znx_big_normalize_tmp_bytes());
313
314 if res_a_infos.base2k() != s_infos.base2k() {
315 tot += GLWE::bytes_of_from_infos(&GLWELayout {
316 n: res_a_infos.n(),
317 base2k: s_infos.base2k(),
318 k: res_a_infos.k(),
319 rank: res_a_infos.rank(),
320 });
321
322 tot += GLWE::bytes_of_from_infos(&GLWELayout {
323 n: res_b_infos.n(),
324 base2k: s_infos.base2k(),
325 k: res_b_infos.k(),
326 rank: res_b_infos.rank(),
327 });
328 }
329
330 tot += self.bytes_of_vec_znx_big(1, s_infos.size());
331
332 tot
333 }
334
335 fn cswap<A, B, S>(&self, res_a: &mut A, res_b: &mut B, s: &S, scratch: &mut Scratch<BE>)
336 where
337 A: GLWEToMut,
338 B: GLWEToMut,
339 S: GGSWPreparedToRef<BE> + GGSWInfos,
340 Scratch<BE>: ScratchTakeCore<BE>,
341 {
342 let res_a: &mut GLWE<&mut [u8]> = &mut res_a.to_mut();
343 let res_b: &mut GLWE<&mut [u8]> = &mut res_b.to_mut();
344 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
345 assert_eq!(res_a.base2k(), res_b.base2k());
346
347 let res_base2k: usize = res_a.base2k().as_usize();
348 let s_base2k: usize = s.base2k().as_usize();
349
350 if res_base2k == s_base2k {
351 let res_big: VecZnxBig<&mut [u8], BE>;
352 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (s.rank() + 1).into(), s.size()); // Todo optimise
353 {
354 // Temporary value storing a - b
355 let tmp_c_infos: GLWELayout = GLWELayout {
356 n: s.n(),
357 base2k: s.base2k(),
358 k: res_a.k().max(res_b.k()),
359 rank: s.rank(),
360 };
361 let (mut tmp_c, scratch_2) = scratch_1.take_glwe(&tmp_c_infos);
362 self.glwe_sub(&mut tmp_c, res_b, res_a);
363 res_big = self.glwe_external_product_internal(res_dft, &tmp_c, s, scratch_2);
364 }
365
366 // Single column res_big to store temporary value before normalization
367 let (mut res_big_tmp, scratch_2) = scratch_1.take_vec_znx_big::<_, BE>(self, 1, res_big.size());
368
369 // res_a = (b-a) * bit + a
370 for j in 0..(res_a.rank() + 1).into() {
371 self.vec_znx_big_add_small(&mut res_big_tmp, 0, &res_big, j, res_a.data(), j);
372 self.vec_znx_big_normalize(
373 res_base2k,
374 res_a.data_mut(),
375 j,
376 s_base2k,
377 &res_big_tmp,
378 0,
379 scratch_2,
380 );
381 }
382
383 // res_b = a - (a - b) * bit = (b - a) * bit + a
384 for j in 0..(res_b.rank() + 1).into() {
385 self.vec_znx_big_sub_small_a(&mut res_big_tmp, 0, res_b.data(), j, &res_big, j);
386 self.vec_znx_big_normalize(
387 res_base2k,
388 res_b.data_mut(),
389 j,
390 s_base2k,
391 &res_big_tmp,
392 0,
393 scratch_2,
394 );
395 }
396 } else {
397 let (mut tmp_a, scratch_1) = scratch.take_glwe(&GLWELayout {
398 n: res_a.n(),
399 base2k: s.base2k(),
400 k: res_a.k(),
401 rank: res_a.rank(),
402 });
403
404 let (mut tmp_b, scratch_2) = scratch_1.take_glwe(&GLWELayout {
405 n: res_b.n(),
406 base2k: s.base2k(),
407 k: res_b.k(),
408 rank: res_b.rank(),
409 });
410
411 self.glwe_normalize(&mut tmp_a, res_a, scratch_2);
412 self.glwe_normalize(&mut tmp_b, res_b, scratch_2);
413
414 let res_big: VecZnxBig<&mut [u8], BE>;
415 let (res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, (s.rank() + 1).into(), s.size()); // Todo optimise
416 {
417 // Temporary value storing a - b
418 let tmp_c_infos: GLWELayout = GLWELayout {
419 n: s.n(),
420 base2k: s.base2k(),
421 k: res_a.k().max(res_b.k()),
422 rank: s.rank(),
423 };
424 let (mut tmp_c, scratch_4) = scratch_3.take_glwe(&tmp_c_infos);
425 self.glwe_sub(&mut tmp_c, res_b, res_a);
426 res_big = self.glwe_external_product_internal(res_dft, &tmp_c, s, scratch_4);
427 }
428
429 // Single column res_big to store temporary value before normalization
430 let (mut res_big_tmp, scratch_4) = scratch_3.take_vec_znx_big::<_, BE>(self, 1, res_big.size());
431
432 // res_a = (b-a) * bit + a
433 for j in 0..(res_a.rank() + 1).into() {
434 self.vec_znx_big_add_small(&mut res_big_tmp, 0, &res_big, j, tmp_a.data(), j);
435 self.vec_znx_big_normalize(
436 res_base2k,
437 res_a.data_mut(),
438 j,
439 s_base2k,
440 &res_big_tmp,
441 0,
442 scratch_4,
443 );
444 }
445
446 // res_b = a - (a - b) * bit = (b - a) * bit + a
447 for j in 0..(res_b.rank() + 1).into() {
448 self.vec_znx_big_sub_small_a(&mut res_big_tmp, 0, tmp_b.data(), j, &res_big, j);
449 self.vec_znx_big_normalize(
450 res_base2k,
451 res_b.data_mut(),
452 j,
453 s_base2k,
454 &res_big_tmp,
455 0,
456 scratch_4,
457 );
458 }
459 }
460 }
461}
462
463pub trait Cmux<BE: Backend>
464where
465 Self: Sized
466 + GLWEExternalProductInternal<BE>
467 + GLWESub
468 + VecZnxBigAddSmallInplace<BE>
469 + GLWENormalize<BE>
470 + VecZnxDftBytesOf
471 + VecZnxBigNormalize<BE>
472 + VecZnxBigNormalizeTmpBytes,
473{
474 fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
475 where
476 R: GLWEInfos,
477 A: GLWEInfos,
478 B: GGSWInfos,
479 {
480 let res_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), b_infos.size());
481 res_dft
482 + self
483 .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, b_infos)
484 .max(self.vec_znx_big_normalize_tmp_bytes())
485 }
486
487 // res = (t - f) * s + f
488 fn cmux<R, T, F, S>(&self, res: &mut R, t: &T, f: &F, s: &S, scratch: &mut Scratch<BE>)
489 where
490 R: GLWEToMut,
491 T: GLWEToRef,
492 F: GLWEToRef,
493 S: GGSWPreparedToRef<BE> + GGSWInfos,
494 Scratch<BE>: ScratchTakeCore<BE>,
495 {
496 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
497 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
498 let f: GLWE<&[u8]> = f.to_ref();
499
500 let res_base2k: usize = res.base2k().into();
501 let ggsw_base2k: usize = s.base2k().into();
502
503 self.glwe_sub(res, t, &f);
504 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
505 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
506 for j in 0..(res.rank() + 1).into() {
507 self.vec_znx_big_add_small_inplace(&mut res_big, j, f.data(), j);
508 self.vec_znx_big_normalize(
509 res_base2k,
510 res.data_mut(),
511 j,
512 ggsw_base2k,
513 &res_big,
514 j,
515 scratch_1,
516 );
517 }
518 }
519
520 // res = (a - res) * s + res
521 fn cmux_inplace_neg<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
522 where
523 R: GLWEToMut,
524 A: GLWEToRef,
525 S: GGSWPreparedToRef<BE> + GGSWInfos,
526 Scratch<BE>: ScratchTakeCore<BE>,
527 {
528 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
529 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
530 let a: &GLWE<&[u8]> = &a.to_ref();
531
532 assert_eq!(res.base2k(), a.base2k());
533
534 let res_base2k: usize = res.base2k().into();
535 let ggsw_base2k: usize = s.base2k().into();
536 let (mut tmp, scratch_1) = scratch.take_glwe(&GLWELayout {
537 n: s.n(),
538 base2k: res.base2k(),
539 k: res.k().max(a.k()),
540 rank: res.rank(),
541 });
542 self.glwe_sub(&mut tmp, a, res);
543 let (res_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
544 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, &tmp, s, scratch_2);
545 for j in 0..(res.rank() + 1).into() {
546 self.vec_znx_big_add_small_inplace(&mut res_big, j, res.data(), j);
547 self.vec_znx_big_normalize(
548 res_base2k,
549 res.data_mut(),
550 j,
551 ggsw_base2k,
552 &res_big,
553 j,
554 scratch_2,
555 );
556 }
557 }
558
559 // res = (res - a) * s + a
560 fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
561 where
562 R: GLWEToMut,
563 A: GLWEToRef,
564 S: GGSWPreparedToRef<BE> + GGSWInfos,
565 Scratch<BE>: ScratchTakeCore<BE>,
566 {
567 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
568 let s: &GGSWPrepared<&[u8], BE> = &s.to_ref();
569 let a: GLWE<&[u8]> = a.to_ref();
570 let res_base2k: usize = res.base2k().into();
571 let ggsw_base2k: usize = s.base2k().into();
572 self.glwe_sub_inplace(res, &a);
573 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), s.size()); // Todo optimise
574 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_external_product_internal(res_dft, res, s, scratch_1);
575 for j in 0..(res.rank() + 1).into() {
576 self.vec_znx_big_add_small_inplace(&mut res_big, j, a.data(), j);
577 self.vec_znx_big_normalize(
578 res_base2k,
579 res.data_mut(),
580 j,
581 ggsw_base2k,
582 &res_big,
583 j,
584 scratch_1,
585 );
586 }
587 }
588}
589
590impl<BE: Backend> Cmux<BE> for Module<BE>
591where
592 Self: Sized
593 + GLWEExternalProductInternal<BE>
594 + GLWESub
595 + VecZnxBigAddSmallInplace<BE>
596 + GLWENormalize<BE>
597 + VecZnxDftBytesOf
598 + VecZnxBigNormalize<BE>
599 + VecZnxBigNormalizeTmpBytes,
600 Scratch<BE>: ScratchTakeCore<BE>,
601{
602}