poulpy/poulpy-hal/src/oep/vmp_pmat.rs
poulpy-hal/src/oep/vmp_pmat.rsRaw
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_bytes_of_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}
148
149/// # THIS TRAIT IS AN OPEN EXTENSION POINT (unsafe)
150/// * See TODO.
151/// * See [crate::api::VmpZero] for corresponding public API.
152/// # Safety [crate::doc::backend_safety] for safety contract.
153pub unsafe trait VmpZeroImpl<B: Backend> {
154 fn vmp_zero_impl<R>(module: &Module<B>, res: &mut R)
155 where
156 R: VmpPMatToMut<B>;
157}