1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{
  7        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxMulXpMinusOne, VecZnxMulXpMinusOneInplace,
  8        VecZnxMulXpMinusOneInplaceTmpBytes,
  9    },
 10    layouts::{Backend, FillUniform, Module, ScratchOwned, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
 11    reference::{
 12        vec_znx::{vec_znx_rotate, vec_znx_sub_ab_inplace},
 13        znx::{ZnxNegate, ZnxRotate, ZnxSubABInplace, ZnxSubBAInplace, ZnxZero},
 14    },
 15    source::Source,
 16};
 17
 18pub fn vec_znx_mul_xp_minus_one_inplace_tmp_bytes(n: usize) -> usize {
 19    n * size_of::<i64>()
 20}
 21
 22pub fn vec_znx_mul_xp_minus_one<R, A, ZNXARI>(p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
 23where
 24    R: VecZnxToMut,
 25    A: VecZnxToRef,
 26    ZNXARI: ZnxRotate + ZnxZero + ZnxSubABInplace,
 27{
 28    vec_znx_rotate::<_, _, ZNXARI>(p, res, res_col, a, a_col);
 29    vec_znx_sub_ab_inplace::<_, _, ZNXARI>(res, res_col, a, a_col);
 30}
 31
 32pub fn vec_znx_mul_xp_minus_one_inplace<R, ZNXARI>(p: i64, res: &mut R, res_col: usize, tmp: &mut [i64])
 33where
 34    R: VecZnxToMut,
 35    ZNXARI: ZnxRotate + ZnxNegate + ZnxSubBAInplace,
 36{
 37    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 38    #[cfg(debug_assertions)]
 39    {
 40        assert_eq!(res.n(), tmp.len());
 41    }
 42    for j in 0..res.size() {
 43        ZNXARI::znx_rotate(p, tmp, res.at(res_col, j));
 44        ZNXARI::znx_sub_ba_inplace(res.at_mut(res_col, j), tmp);
 45    }
 46}
 47
 48pub fn bench_vec_znx_mul_xp_minus_one<B: Backend>(c: &mut Criterion, label: &str)
 49where
 50    Module<B>: VecZnxMulXpMinusOne + ModuleNew<B>,
 51{
 52    let group_name: String = format!("vec_znx_mul_xp_minus_one::{}", 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>: VecZnxMulXpMinusOne + 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 res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
 70
 71        // Fill a with random i64
 72        a.fill_uniform(50, &mut source);
 73        res.fill_uniform(50, &mut source);
 74
 75        move || {
 76            for i in 0..cols {
 77                module.vec_znx_mul_xp_minus_one(-7, &mut res, 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_mul_xp_minus_one_inplace<B: Backend>(c: &mut Criterion, label: &str)
 93where
 94    Module<B>: VecZnxMulXpMinusOneInplace<B> + VecZnxMulXpMinusOneInplaceTmpBytes + ModuleNew<B>,
 95    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 96{
 97    let group_name: String = format!("vec_znx_mul_xp_minus_one_inplace::{}", label);
 98
 99    let mut group = c.benchmark_group(group_name);
100
101    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
102    where
103        Module<B>: VecZnxMulXpMinusOneInplace<B> + ModuleNew<B> + VecZnxMulXpMinusOneInplaceTmpBytes,
104        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
105    {
106        let n: usize = 1 << params[0];
107        let cols: usize = params[1];
108        let size: usize = params[2];
109
110        let module: Module<B> = Module::<B>::new(n as u64);
111
112        let mut source: Source = Source::new([0u8; 32]);
113
114        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
115
116        let mut scratch = ScratchOwned::alloc(module.vec_znx_mul_xp_minus_one_inplace_tmp_bytes());
117
118        // Fill a with random i64
119        res.fill_uniform(50, &mut source);
120
121        move || {
122            for i in 0..cols {
123                module.vec_znx_mul_xp_minus_one_inplace(-7, &mut res, i, scratch.borrow());
124            }
125            black_box(());
126        }
127    }
128
129    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
130        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
131        let mut runner = runner::<B>(params);
132        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
133    }
134
135    group.finish();
136}