1use crate::layouts::{
  2    Backend, MatZnxToRef, Module, Scratch, VecZnxDftToMut, VecZnxDftToRef, VmpPMatOwned, VmpPMatToMut, VmpPMatToRef,
  3};
  4
  5/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
  6/// * See TODO for reference code.
  7/// * See TODO for corresponding public API.
  8/// # Safety [crate::doc::backend_safety] for safety contract.
  9pub unsafe trait VmpPMatAllocImpl<B: Backend> {
 10    fn vmp_pmat_alloc_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<B>;
 11}
 12
 13/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 14/// * See TODO for reference code.
 15/// * See TODO for corresponding public API.
 16/// # Safety [crate::doc::backend_safety] for safety contract.
 17pub unsafe trait VmpPMatAllocBytesImpl<B: Backend> {
 18    fn vmp_pmat_alloc_bytes_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
 19}
 20
 21/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 22/// * See TODO for reference code.
 23/// * See TODO for corresponding public API.
 24/// # Safety [crate::doc::backend_safety] for safety contract.
 25pub unsafe trait VmpPMatFromBytesImpl<B: Backend> {
 26    fn vmp_pmat_from_bytes_impl(
 27        n: usize,
 28        rows: usize,
 29        cols_in: usize,
 30        cols_out: usize,
 31        size: usize,
 32        bytes: Vec<u8>,
 33    ) -> VmpPMatOwned<B>;
 34}
 35
 36/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 37/// * See TODO for reference code.
 38/// * See TODO for corresponding public API.
 39/// # Safety [crate::doc::backend_safety] for safety contract.
 40pub unsafe trait VmpPrepareTmpBytesImpl<B: Backend> {
 41    fn vmp_prepare_tmp_bytes_impl(
 42        module: &Module<B>,
 43        n: usize,
 44        rows: usize,
 45        cols_in: usize,
 46        cols_out: usize,
 47        size: usize,
 48    ) -> usize;
 49}
 50
 51/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 52/// * See TODO for reference code.
 53/// * See TODO for corresponding public API.
 54/// # Safety [crate::doc::backend_safety] for safety contract.
 55pub unsafe trait VmpPMatPrepareImpl<B: Backend> {
 56    fn vmp_prepare_impl<R, A>(module: &Module<B>, res: &mut R, a: &A, scratch: &mut Scratch<B>)
 57    where
 58        R: VmpPMatToMut<B>,
 59        A: MatZnxToRef;
 60}
 61
 62#[allow(clippy::too_many_arguments)]
 63/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 64/// * See TODO for reference code.
 65/// * See TODO for corresponding public API.
 66/// # Safety [crate::doc::backend_safety] for safety contract.
 67pub unsafe trait VmpApplyTmpBytesImpl<B: Backend> {
 68    fn vmp_apply_tmp_bytes_impl(
 69        module: &Module<B>,
 70        n: usize,
 71        res_size: usize,
 72        a_size: usize,
 73        b_rows: usize,
 74        b_cols_in: usize,
 75        b_cols_out: usize,
 76        b_size: usize,
 77    ) -> usize;
 78}
 79
 80/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 81/// * See TODO for reference code.
 82/// * See TODO for corresponding public API.
 83/// # Safety [crate::doc::backend_safety] for safety contract.
 84pub unsafe trait VmpApplyImpl<B: Backend> {
 85    fn vmp_apply_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
 86    where
 87        R: VecZnxDftToMut<B>,
 88        A: VecZnxDftToRef<B>,
 89        C: VmpPMatToRef<B>;
 90}
 91
 92#[allow(clippy::too_many_arguments)]
 93/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 94/// * See TODO for reference code.
 95/// * See TODO for corresponding public API.
 96/// # Safety [crate::doc::backend_safety] for safety contract.
 97pub unsafe trait VmpApplyAddTmpBytesImpl<B: Backend> {
 98    fn vmp_apply_add_tmp_bytes_impl(
 99        module: &Module<B>,
100        n: usize,
101        res_size: usize,
102        a_size: usize,
103        b_rows: usize,
104        b_cols_in: usize,
105        b_cols_out: usize,
106        b_size: usize,
107    ) -> usize;
108}
109
110/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
111/// * See TODO for reference code.
112/// * See TODO for corresponding public API.
113/// # Safety [crate::doc::backend_safety] for safety contract.
114pub unsafe trait VmpApplyAddImpl<B: Backend> {
115    // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
116    fn vmp_apply_add_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<B>)
117    where
118        R: VecZnxDftToMut<B>,
119        A: VecZnxDftToRef<B>,
120        C: VmpPMatToRef<B>;
121}