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