1use crate::{
  2    api::{
  3        VmpApplyDft, VmpApplyDftTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftAddTmpBytes,
  4        VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPMatBytesOf, VmpPMatFromBytes, VmpPrepare, VmpPrepareTmpBytes, VmpZero,
  5    },
  6    layouts::{
  7        Backend, MatZnxToRef, Module, Scratch, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef, VmpPMatOwned, VmpPMatToMut,
  8        VmpPMatToRef,
  9    },
 10    oep::{
 11        VmpApplyDftImpl, VmpApplyDftTmpBytesImpl, VmpApplyDftToDftAddImpl, VmpApplyDftToDftAddTmpBytesImpl, VmpApplyDftToDftImpl,
 12        VmpApplyDftToDftTmpBytesImpl, VmpPMatAllocBytesImpl, VmpPMatAllocImpl, VmpPMatFromBytesImpl, VmpPrepareImpl,
 13        VmpPrepareTmpBytesImpl, VmpZeroImpl,
 14    },
 15};
 16
 17impl<B> VmpPMatAlloc<B> for Module<B>
 18where
 19    B: Backend + VmpPMatAllocImpl<B>,
 20{
 21    fn vmp_pmat_alloc(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<B> {
 22        B::vmp_pmat_alloc_impl(self.n(), rows, cols_in, cols_out, size)
 23    }
 24}
 25
 26impl<B> VmpPMatBytesOf for Module<B>
 27where
 28    B: Backend + VmpPMatAllocBytesImpl<B>,
 29{
 30    fn bytes_of_vmp_pmat(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 31        B::vmp_pmat_bytes_of_impl(self.n(), rows, cols_in, cols_out, size)
 32    }
 33}
 34
 35impl<B> VmpPMatFromBytes<B> for Module<B>
 36where
 37    B: Backend + VmpPMatFromBytesImpl<B>,
 38{
 39    fn vmp_pmat_from_bytes(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize, bytes: Vec<u8>) -> VmpPMatOwned<B> {
 40        B::vmp_pmat_from_bytes_impl(self.n(), rows, cols_in, cols_out, size, bytes)
 41    }
 42}
 43
 44impl<B> VmpPrepareTmpBytes for Module<B>
 45where
 46    B: Backend + VmpPrepareTmpBytesImpl<B>,
 47{
 48    fn vmp_prepare_tmp_bytes(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 49        B::vmp_prepare_tmp_bytes_impl(self, rows, cols_in, cols_out, size)
 50    }
 51}
 52
 53impl<B> VmpPrepare<B> for Module<B>
 54where
 55    B: Backend + VmpPrepareImpl<B>,
 56{
 57    fn vmp_prepare<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<B>)
 58    where
 59        R: VmpPMatToMut<B>,
 60        A: MatZnxToRef,
 61    {
 62        B::vmp_prepare_impl(self, res, a, scratch)
 63    }
 64}
 65
 66impl<B> VmpApplyDftTmpBytes for Module<B>
 67where
 68    B: Backend + VmpApplyDftTmpBytesImpl<B>,
 69{
 70    fn vmp_apply_dft_tmp_bytes(
 71        &self,
 72        res_size: usize,
 73        a_size: usize,
 74        b_rows: usize,
 75        b_cols_in: usize,
 76        b_cols_out: usize,
 77        b_size: usize,
 78    ) -> usize {
 79        B::vmp_apply_dft_tmp_bytes_impl(
 80            self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size,
 81        )
 82    }
 83}
 84
 85impl<B> VmpApplyDft<B> for Module<B>
 86where
 87    B: Backend + VmpApplyDftImpl<B>,
 88{
 89    fn vmp_apply_dft<R, A, C>(&self, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
 90    where
 91        R: VecZnxDftToMut<B>,
 92        A: VecZnxToRef,
 93        C: VmpPMatToRef<B>,
 94    {
 95        B::vmp_apply_dft_impl(self, res, a, b, scratch);
 96    }
 97}
 98
 99impl<B> VmpApplyDftToDftTmpBytes for Module<B>
100where
101    B: Backend + VmpApplyDftToDftTmpBytesImpl<B>,
102{
103    fn vmp_apply_dft_to_dft_tmp_bytes(
104        &self,
105        res_size: usize,
106        a_size: usize,
107        b_rows: usize,
108        b_cols_in: usize,
109        b_cols_out: usize,
110        b_size: usize,
111    ) -> usize {
112        B::vmp_apply_dft_to_dft_tmp_bytes_impl(
113            self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size,
114        )
115    }
116}
117
118impl<B> VmpApplyDftToDft<B> for Module<B>
119where
120    B: Backend + VmpApplyDftToDftImpl<B>,
121{
122    fn vmp_apply_dft_to_dft<R, A, C>(&self, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
123    where
124        R: VecZnxDftToMut<B>,
125        A: VecZnxDftToRef<B>,
126        C: VmpPMatToRef<B>,
127    {
128        B::vmp_apply_dft_to_dft_impl(self, res, a, b, scratch);
129    }
130}
131
132impl<B> VmpApplyDftToDftAddTmpBytes for Module<B>
133where
134    B: Backend + VmpApplyDftToDftAddTmpBytesImpl<B>,
135{
136    fn vmp_apply_dft_to_dft_add_tmp_bytes(
137        &self,
138        res_size: usize,
139        a_size: usize,
140        b_rows: usize,
141        b_cols_in: usize,
142        b_cols_out: usize,
143        b_size: usize,
144    ) -> usize {
145        B::vmp_apply_dft_to_dft_add_tmp_bytes_impl(
146            self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size,
147        )
148    }
149}
150
151impl<B> VmpApplyDftToDftAdd<B> for Module<B>
152where
153    B: Backend + VmpApplyDftToDftAddImpl<B>,
154{
155    fn vmp_apply_dft_to_dft_add<R, A, C>(&self, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<B>)
156    where
157        R: VecZnxDftToMut<B>,
158        A: VecZnxDftToRef<B>,
159        C: VmpPMatToRef<B>,
160    {
161        B::vmp_apply_dft_to_dft_add_impl(self, res, a, b, scale, scratch);
162    }
163}
164
165impl<B> VmpZero<B> for Module<B>
166where
167    B: Backend + VmpZeroImpl<B>,
168{
169    fn vmp_zero<R>(&self, res: &mut R)
170    where
171        R: VmpPMatToMut<B>,
172    {
173        B::vmp_zero_impl(self, res);
174    }
175}