1use crate::hal::layouts::{
 2    Backend, MatZnxToRef, Module, Scratch, VecZnxDftToMut, VecZnxDftToRef, VmpPMatOwned, VmpPMatToMut, VmpPMatToRef,
 3};
 4
 5pub unsafe trait VmpPMatAllocImpl<B: Backend> {
 6    fn vmp_pmat_alloc_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<B>;
 7}
 8
 9pub unsafe trait VmpPMatAllocBytesImpl<B: Backend> {
10    fn vmp_pmat_alloc_bytes_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
11}
12
13pub unsafe trait VmpPMatFromBytesImpl<B: Backend> {
14    fn vmp_pmat_from_bytes_impl(
15        n: usize,
16        rows: usize,
17        cols_in: usize,
18        cols_out: usize,
19        size: usize,
20        bytes: Vec<u8>,
21    ) -> VmpPMatOwned<B>;
22}
23
24pub unsafe trait VmpPrepareTmpBytesImpl<B: Backend> {
25    fn vmp_prepare_tmp_bytes_impl(
26        module: &Module<B>,
27        n: usize,
28        rows: usize,
29        cols_in: usize,
30        cols_out: usize,
31        size: usize,
32    ) -> usize;
33}
34
35pub unsafe trait VmpPMatPrepareImpl<B: Backend> {
36    fn vmp_prepare_impl<R, A>(module: &Module<B>, res: &mut R, a: &A, scratch: &mut Scratch<B>)
37    where
38        R: VmpPMatToMut<B>,
39        A: MatZnxToRef;
40}
41
42pub unsafe trait VmpApplyTmpBytesImpl<B: Backend> {
43    fn vmp_apply_tmp_bytes_impl(
44        module: &Module<B>,
45        n: usize,
46        res_size: usize,
47        a_size: usize,
48        b_rows: usize,
49        b_cols_in: usize,
50        b_cols_out: usize,
51        b_size: usize,
52    ) -> usize;
53}
54
55pub unsafe trait VmpApplyImpl<B: Backend> {
56    fn vmp_apply_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
57    where
58        R: VecZnxDftToMut<B>,
59        A: VecZnxDftToRef<B>,
60        C: VmpPMatToRef<B>;
61}
62
63pub unsafe trait VmpApplyAddTmpBytesImpl<B: Backend> {
64    fn vmp_apply_add_tmp_bytes_impl(
65        module: &Module<B>,
66        n: usize,
67        res_size: usize,
68        a_size: usize,
69        b_rows: usize,
70        b_cols_in: usize,
71        b_cols_out: usize,
72        b_size: usize,
73    ) -> usize;
74}
75
76pub unsafe trait VmpApplyAddImpl<B: Backend> {
77    // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
78    fn vmp_apply_add_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<B>)
79    where
80        R: VecZnxDftToMut<B>,
81        A: VecZnxDftToRef<B>,
82        C: VmpPMatToRef<B>;
83}