1use crate::hal::layouts::{
  2    Backend, MatZnxToRef, Scratch, VecZnxDftToMut, VecZnxDftToRef, VmpPMatOwned, VmpPMatToMut, VmpPMatToRef,
  3};
  4
  5pub trait VmpPMatAlloc<B: Backend> {
  6    fn vmp_pmat_alloc(&self, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<B>;
  7}
  8
  9pub trait VmpPMatAllocBytes {
 10    fn vmp_pmat_alloc_bytes(&self, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
 11}
 12
 13pub trait VmpPMatFromBytes<B: Backend> {
 14    fn vmp_pmat_from_bytes(
 15        &self,
 16        n: usize,
 17        rows: usize,
 18        cols_in: usize,
 19        cols_out: usize,
 20        size: usize,
 21        bytes: Vec<u8>,
 22    ) -> VmpPMatOwned<B>;
 23}
 24
 25pub trait VmpPrepareTmpBytes {
 26    fn vmp_prepare_tmp_bytes(&self, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
 27}
 28
 29pub trait VmpPrepare<B: Backend> {
 30    fn vmp_prepare<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<B>)
 31    where
 32        R: VmpPMatToMut<B>,
 33        A: MatZnxToRef;
 34}
 35
 36#[allow(clippy::too_many_arguments)]
 37pub trait VmpApplyTmpBytes {
 38    fn vmp_apply_tmp_bytes(
 39        &self,
 40        n: usize,
 41        res_size: usize,
 42        a_size: usize,
 43        b_rows: usize,
 44        b_cols_in: usize,
 45        b_cols_out: usize,
 46        b_size: usize,
 47    ) -> usize;
 48}
 49
 50pub trait VmpApply<B: Backend> {
 51    /// Applies the vector matrix product [crate::hal::layouts::VecZnxDft] x [crate::hal::layouts::VmpPMat].
 52    ///
 53    /// A vector matrix product numerically equivalent to a sum of [crate::hal::api::SvpApply],
 54    /// where each [crate::hal::layouts::SvpPPol] is a limb of the input [crate::hal::layouts::VecZnx] in DFT,
 55    /// and each vector a [crate::hal::layouts::VecZnxDft] (row) of the [crate::hal::layouts::VmpPMat].
 56    ///
 57    /// As such, given an input [crate::hal::layouts::VecZnx] of `i` size and a [crate::hal::layouts::VmpPMat] of `i` rows and
 58    /// `j` size, the output is a [crate::hal::layouts::VecZnx] of `j` size.
 59    ///
 60    /// If there is a mismatch between the dimensions the largest valid ones are used.
 61    ///
 62    /// ```text
 63    /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
 64    ///             |h i j|
 65    ///             |k l m|
 66    /// ```
 67    /// where each element is a [crate::hal::layouts::VecZnxDft].
 68    ///
 69    /// # Arguments
 70    ///
 71    /// * `c`: the output of the vector matrix product, as a [crate::hal::layouts::VecZnxDft].
 72    /// * `a`: the left operand [crate::hal::layouts::VecZnxDft] of the vector matrix product.
 73    /// * `b`: the right operand [crate::hal::layouts::VmpPMat] of the vector matrix product.
 74    /// * `buf`: scratch space, the size can be obtained with [VmpApplyTmpBytes::vmp_apply_tmp_bytes].
 75    fn vmp_apply<R, A, C>(&self, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
 76    where
 77        R: VecZnxDftToMut<B>,
 78        A: VecZnxDftToRef<B>,
 79        C: VmpPMatToRef<B>;
 80}
 81
 82#[allow(clippy::too_many_arguments)]
 83pub trait VmpApplyAddTmpBytes {
 84    fn vmp_apply_add_tmp_bytes(
 85        &self,
 86        n: usize,
 87        res_size: usize,
 88        a_size: usize,
 89        b_rows: usize,
 90        b_cols_in: usize,
 91        b_cols_out: usize,
 92        b_size: usize,
 93    ) -> usize;
 94}
 95
 96pub trait VmpApplyAdd<B: Backend> {
 97    fn vmp_apply_add<R, A, C>(&self, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<B>)
 98    where
 99        R: VecZnxDftToMut<B>,
100        A: VecZnxDftToRef<B>,
101        C: VmpPMatToRef<B>;
102}