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(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
 80    }
 81}
 82
 83impl<B> VmpApplyDft<B> for Module<B>
 84where
 85    B: Backend + VmpApplyDftImpl<B>,
 86{
 87    fn vmp_apply_dft<R, A, C>(&self, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
 88    where
 89        R: VecZnxDftToMut<B>,
 90        A: VecZnxToRef,
 91        C: VmpPMatToRef<B>,
 92    {
 93        B::vmp_apply_dft_impl(self, res, a, b, scratch);
 94    }
 95}
 96
 97impl<B> VmpApplyDftToDftTmpBytes for Module<B>
 98where
 99    B: Backend + VmpApplyDftToDftTmpBytesImpl<B>,
100{
101    fn vmp_apply_dft_to_dft_tmp_bytes(
102        &self,
103        res_size: usize,
104        a_size: usize,
105        b_rows: usize,
106        b_cols_in: usize,
107        b_cols_out: usize,
108        b_size: usize,
109    ) -> usize {
110        B::vmp_apply_dft_to_dft_tmp_bytes_impl(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
111    }
112}
113
114impl<B> VmpApplyDftToDft<B> for Module<B>
115where
116    B: Backend + VmpApplyDftToDftImpl<B>,
117{
118    fn vmp_apply_dft_to_dft<R, A, C>(&self, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<B>)
119    where
120        R: VecZnxDftToMut<B>,
121        A: VecZnxDftToRef<B>,
122        C: VmpPMatToRef<B>,
123    {
124        B::vmp_apply_dft_to_dft_impl(self, res, a, b, scratch);
125    }
126}
127
128impl<B> VmpApplyDftToDftAddTmpBytes for Module<B>
129where
130    B: Backend + VmpApplyDftToDftAddTmpBytesImpl<B>,
131{
132    fn vmp_apply_dft_to_dft_add_tmp_bytes(
133        &self,
134        res_size: usize,
135        a_size: usize,
136        b_rows: usize,
137        b_cols_in: usize,
138        b_cols_out: usize,
139        b_size: usize,
140    ) -> usize {
141        B::vmp_apply_dft_to_dft_add_tmp_bytes_impl(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
142    }
143}
144
145impl<B> VmpApplyDftToDftAdd<B> for Module<B>
146where
147    B: Backend + VmpApplyDftToDftAddImpl<B>,
148{
149    fn vmp_apply_dft_to_dft_add<R, A, C>(&self, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<B>)
150    where
151        R: VecZnxDftToMut<B>,
152        A: VecZnxDftToRef<B>,
153        C: VmpPMatToRef<B>,
154    {
155        B::vmp_apply_dft_to_dft_add_impl(self, res, a, b, scale, scratch);
156    }
157}
158
159impl<B> VmpZero<B> for Module<B>
160where
161    B: Backend + VmpZeroImpl<B>,
162{
163    fn vmp_zero<R>(&self, res: &mut R)
164    where
165        R: VmpPMatToMut<B>,
166    {
167        B::vmp_zero_impl(self, res);
168    }
169}