1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{ModuleNew, VecZnxNegate, VecZnxNegateInplace},
  7    layouts::{Backend, FillUniform, Module, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
  8    reference::znx::{ZnxNegate, ZnxNegateInplace, ZnxZero},
  9    source::Source,
 10};
 11
 12pub fn vec_znx_negate<R, A, ZNXARI>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 13where
 14    R: VecZnxToMut,
 15    A: VecZnxToRef,
 16    ZNXARI: ZnxNegate + ZnxZero,
 17{
 18    let a: VecZnx<&[u8]> = a.to_ref();
 19    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 20
 21    #[cfg(debug_assertions)]
 22    {
 23        assert_eq!(a.n(), res.n());
 24    }
 25
 26    let min_size: usize = res.size().min(a.size());
 27
 28    for j in 0..min_size {
 29        ZNXARI::znx_negate(res.at_mut(res_col, j), a.at(a_col, j));
 30    }
 31
 32    for j in min_size..res.size() {
 33        ZNXARI::znx_zero(res.at_mut(res_col, j));
 34    }
 35}
 36
 37pub fn vec_znx_negate_inplace<R, ZNXARI>(res: &mut R, res_col: usize)
 38where
 39    R: VecZnxToMut,
 40    ZNXARI: ZnxNegateInplace,
 41{
 42    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 43    for j in 0..res.size() {
 44        ZNXARI::znx_negate_inplace(res.at_mut(res_col, j));
 45    }
 46}
 47
 48pub fn bench_vec_znx_negate<B: Backend>(c: &mut Criterion, label: &str)
 49where
 50    Module<B>: VecZnxNegate + ModuleNew<B>,
 51{
 52    let group_name: String = format!("vec_znx_negate::{label}");
 53
 54    let mut group = c.benchmark_group(group_name);
 55
 56    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
 57    where
 58        Module<B>: VecZnxNegate + ModuleNew<B>,
 59    {
 60        let n: usize = 1 << params[0];
 61        let cols: usize = params[1];
 62        let size: usize = params[2];
 63
 64        let module: Module<B> = Module::<B>::new(n as u64);
 65
 66        let mut source: Source = Source::new([0u8; 32]);
 67
 68        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
 69        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
 70
 71        // Fill a with random i64
 72        a.fill_uniform(50, &mut source);
 73        b.fill_uniform(50, &mut source);
 74
 75        move || {
 76            for i in 0..cols {
 77                module.vec_znx_negate(&mut b, i, &a, i);
 78            }
 79            black_box(());
 80        }
 81    }
 82
 83    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
 84        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
 85        let mut runner = runner::<B>(params);
 86        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 87    }
 88
 89    group.finish();
 90}
 91
 92pub fn bench_vec_znx_negate_inplace<B: Backend>(c: &mut Criterion, label: &str)
 93where
 94    Module<B>: VecZnxNegateInplace + ModuleNew<B>,
 95{
 96    let group_name: String = format!("vec_znx_negate_inplace::{label}");
 97
 98    let mut group = c.benchmark_group(group_name);
 99
100    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
101    where
102        Module<B>: VecZnxNegateInplace + ModuleNew<B>,
103    {
104        let n: usize = 1 << params[0];
105        let cols: usize = params[1];
106        let size: usize = params[2];
107
108        let module: Module<B> = Module::<B>::new(n as u64);
109
110        let mut source: Source = Source::new([0u8; 32]);
111
112        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
113
114        // Fill a with random i64
115        a.fill_uniform(50, &mut source);
116        move || {
117            for i in 0..cols {
118                module.vec_znx_negate_inplace(&mut a, i);
119            }
120            black_box(());
121        }
122    }
123
124    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
125        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2]));
126        let mut runner = runner::<B>(params);
127        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
128    }
129
130    group.finish();
131}