1use crate::layouts::{
  2    Backend, MatZnxToRef, Module, Scratch, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef, VmpPMatOwned, VmpPMatToMut, VmpPMatToRef,
  3};
  4
  5/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
  6/// * See [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
  7/// * See [crate::api::VmpPMatAlloc] 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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 15/// * See [crate::api::VmpPMatAllocBytes] 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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 23/// * See [crate::api::VmpPMatFromBytes] 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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 38/// * See [crate::api::VmpPrepareTmpBytes] 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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 46/// * See [crate::api::VmpPrepare] for corresponding public API.
 47/// # Safety [crate::doc::backend_safety] for safety contract.
 48pub unsafe trait VmpPrepareImpl<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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 58/// * See [crate::api::VmpApplyDftTmpBytes] for corresponding public API.
 59/// # Safety [crate::doc::backend_safety] for safety contract.
 60pub unsafe trait VmpApplyDftTmpBytesImpl<B: Backend> {
 61    fn vmp_apply_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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 74/// * See [crate::api::VmpApplyDft] for corresponding public API.
 75/// # Safety [crate::doc::backend_safety] for safety contract.
 76pub unsafe trait VmpApplyDftImpl<B: Backend> {
 77    fn vmp_apply_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: VecZnxToRef,
 81        C: VmpPMatToRef<B>;
 82}
 83
 84#[allow(clippy::too_many_arguments)]
 85/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
 86/// * See [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
 87/// * See [crate::api::VmpApplyDftToDftTmpBytes] for corresponding public API.
 88/// # Safety [crate::doc::backend_safety] for safety contract.
 89pub unsafe trait VmpApplyDftToDftTmpBytesImpl<B: Backend> {
 90    fn vmp_apply_dft_to_dft_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 [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
103/// * See [crate::api::VmpApplyDftToDft] for corresponding public API.
104/// # Safety [crate::doc::backend_safety] for safety contract.
105pub unsafe trait VmpApplyDftToDftImpl<B: Backend> {
106    fn vmp_apply_dft_to_dft_impl<R, A, C>(module: &Module<B>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
107    where
108        R: VecZnxDftToMut<B>,
109        A: VecZnxDftToRef<B>,
110        C: VmpPMatToRef<B>;
111}
112
113#[allow(clippy::too_many_arguments)]
114/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
115/// * See [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
116/// * See [crate::api::VmpApplyDftToDftAddTmpBytes] for corresponding public API.
117/// # Safety [crate::doc::backend_safety] for safety contract.
118pub unsafe trait VmpApplyDftToDftAddTmpBytesImpl<B: Backend> {
119    fn vmp_apply_dft_to_dft_add_tmp_bytes_impl(
120        module: &Module<B>,
121        res_size: usize,
122        a_size: usize,
123        b_rows: usize,
124        b_cols_in: usize,
125        b_cols_out: usize,
126        b_size: usize,
127    ) -> usize;
128}
129
130/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
131/// * See [poulpy-backend/src/cpu_fft64_ref/vmp.rs](https://github.com/phantomzone-org/poulpy/blob/main/poulpy-backend/src/cpu_fft64_ref/vmp.rs) for reference implementation.
132/// * See [crate::api::VmpApplyDftToDftAdd] for corresponding public API.
133/// # Safety [crate::doc::backend_safety] for safety contract.
134pub unsafe trait VmpApplyDftToDftAddImpl<B: Backend> {
135    // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
136    fn vmp_apply_dft_to_dft_add_impl<R, A, C>(
137        module: &Module<B>,
138        res: &mut R,
139        a: &A,
140        b: &C,
141        scale: usize,
142        scratch: &mut Scratch<B>,
143    ) where
144        R: VecZnxDftToMut<B>,
145        A: VecZnxDftToRef<B>,
146        C: VmpPMatToRef<B>;
147}