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(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
 42}
 43
 44/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 45/// * See TODO for reference code.
 46/// * See TODO for corresponding public API.
 47/// # Safety [crate::doc::backend_safety] for safety contract.
 48pub unsafe trait VmpPMatPrepareImpl<B: Backend> {
 49    fn vmp_prepare_impl<R, A>(module: &Module<B>, res: &mut R, a: &A, scratch: &mut Scratch<B>)
 50    where
 51        R: VmpPMatToMut<B>,
 52        A: MatZnxToRef;
 53}
 54
 55#[allow(clippy::too_many_arguments)]
 56/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 57/// * See TODO for reference code.
 58/// * See TODO for corresponding public API.
 59/// # Safety [crate::doc::backend_safety] for safety contract.
 60pub unsafe trait VmpApplyDftToDftTmpBytesImpl<B: Backend> {
 61    fn vmp_apply_dft_to_dft_tmp_bytes_impl(
 62        module: &Module<B>,
 63        res_size: usize,
 64        a_size: usize,
 65        b_rows: usize,
 66        b_cols_in: usize,
 67        b_cols_out: usize,
 68        b_size: usize,
 69    ) -> usize;
 70}
 71
 72/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 73/// * See TODO for reference code.
 74/// * See TODO for corresponding public API.
 75/// # Safety [crate::doc::backend_safety] for safety contract.
 76pub unsafe trait VmpApplyDftToDftImpl<B: Backend> {
 77    fn vmp_apply_dft_to_dft_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
 78    where
 79        R: VecZnxDftToMut<B>,
 80        A: VecZnxDftToRef<B>,
 81        C: VmpPMatToRef<B>;
 82}
 83
 84#[allow(clippy::too_many_arguments)]
 85/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 86/// * See TODO for reference code.
 87/// * See TODO for corresponding public API.
 88/// # Safety [crate::doc::backend_safety] for safety contract.
 89pub unsafe trait VmpApplyDftToDftAddTmpBytesImpl<B: Backend> {
 90    fn vmp_apply_dft_to_dft_add_tmp_bytes_impl(
 91        module: &Module<B>,
 92        res_size: usize,
 93        a_size: usize,
 94        b_rows: usize,
 95        b_cols_in: usize,
 96        b_cols_out: usize,
 97        b_size: usize,
 98    ) -> usize;
 99}
100
101/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
102/// * See TODO for reference code.
103/// * See TODO for corresponding public API.
104/// # Safety [crate::doc::backend_safety] for safety contract.
105pub unsafe trait VmpApplyDftToDftAddImpl<B: Backend> {
106    // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
107    fn vmp_apply_dft_to_dft_add_impl<R, A, C>(
108        module: &Module<B>,
109        res: &mut R,
110        a: &A,
111        b: &C,
112        scale: usize,
113        scratch: &mut Scratch<B>,
114    ) where
115        R: VecZnxDftToMut<B>,
116        A: VecZnxDftToRef<B>,
117        C: VmpPMatToRef<B>;
118}