poulpy/poulpy-hal/src/delegates/convolution.rs
poulpy-hal/src/delegates/convolution.rsRaw
1use crate::{
2 api::{CnvPVecAlloc, CnvPVecBytesOf, Convolution},
3 layouts::{
4 Backend, CnvPVecL, CnvPVecLToMut, CnvPVecLToRef, CnvPVecR, CnvPVecRToMut, CnvPVecRToRef, Module, Scratch, VecZnxBigToMut,
5 VecZnxDftToMut, VecZnxToRef, ZnxInfos, ZnxViewMut,
6 },
7 oep::{CnvPVecBytesOfImpl, CnvPVecLAllocImpl, ConvolutionImpl},
8};
9
10impl<BE: Backend> CnvPVecAlloc<BE> for Module<BE>
11where
12 BE: CnvPVecLAllocImpl<BE>,
13{
14 fn cnv_pvec_left_alloc(&self, cols: usize, size: usize) -> CnvPVecL<Vec<u8>, BE> {
15 BE::cnv_pvec_left_alloc_impl(self.n(), cols, size)
16 }
17
18 fn cnv_pvec_right_alloc(&self, cols: usize, size: usize) -> CnvPVecR<Vec<u8>, BE> {
19 BE::cnv_pvec_right_alloc_impl(self.n(), cols, size)
20 }
21}
22
23impl<BE: Backend> CnvPVecBytesOf for Module<BE>
24where
25 BE: CnvPVecBytesOfImpl,
26{
27 fn bytes_of_cnv_pvec_left(&self, cols: usize, size: usize) -> usize {
28 BE::bytes_of_cnv_pvec_left_impl(self.n(), cols, size)
29 }
30
31 fn bytes_of_cnv_pvec_right(&self, cols: usize, size: usize) -> usize {
32 BE::bytes_of_cnv_pvec_right_impl(self.n(), cols, size)
33 }
34}
35
36impl<BE: Backend> Convolution<BE> for Module<BE>
37where
38 BE: ConvolutionImpl<BE>,
39{
40 fn cnv_prepare_left_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize {
41 BE::cnv_prepare_left_tmp_bytes_impl(self, res_size, a_size)
42 }
43 fn cnv_prepare_left<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
44 where
45 R: CnvPVecLToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
46 A: VecZnxToRef + ZnxInfos,
47 {
48 BE::cnv_prepare_left_impl(self, res, a, scratch);
49 }
50
51 fn cnv_prepare_right_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize {
52 BE::cnv_prepare_right_tmp_bytes_impl(self, res_size, a_size)
53 }
54 fn cnv_prepare_right<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
55 where
56 R: CnvPVecRToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
57 A: VecZnxToRef + ZnxInfos,
58 {
59 BE::cnv_prepare_right_impl(self, res, a, scratch);
60 }
61
62 fn cnv_apply_dft_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize {
63 BE::cnv_apply_dft_tmp_bytes_impl(self, res_size, res_offset, a_size, b_size)
64 }
65
66 fn cnv_by_const_apply_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize {
67 BE::cnv_by_const_apply_tmp_bytes_impl(self, res_size, res_offset, a_size, b_size)
68 }
69
70 fn cnv_by_const_apply<R, A>(
71 &self,
72 res: &mut R,
73 res_offset: usize,
74 res_col: usize,
75 a: &A,
76 a_col: usize,
77 b: &[i64],
78 scratch: &mut Scratch<BE>,
79 ) where
80 R: VecZnxBigToMut<BE>,
81 A: VecZnxToRef,
82 {
83 BE::cnv_by_const_apply_impl(self, res, res_offset, res_col, a, a_col, b, scratch);
84 }
85
86 fn cnv_apply_dft<R, A, B>(
87 &self,
88 res: &mut R,
89 res_offset: usize,
90 res_col: usize,
91 a: &A,
92 a_col: usize,
93 b: &B,
94 b_col: usize,
95 scratch: &mut Scratch<BE>,
96 ) where
97 R: VecZnxDftToMut<BE>,
98 A: CnvPVecLToRef<BE>,
99 B: CnvPVecRToRef<BE>,
100 {
101 BE::cnv_apply_dft_impl(self, res, res_offset, res_col, a, a_col, b, b_col, scratch);
102 }
103
104 fn cnv_pairwise_apply_dft_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize {
105 BE::cnv_pairwise_apply_dft_tmp_bytes(self, res_size, res_offset, a_size, b_size)
106 }
107
108 fn cnv_pairwise_apply_dft<R, A, B>(
109 &self,
110 res: &mut R,
111 res_offset: usize,
112 res_col: usize,
113 a: &A,
114 b: &B,
115 i: usize,
116 j: usize,
117 scratch: &mut Scratch<BE>,
118 ) where
119 R: VecZnxDftToMut<BE>,
120 A: CnvPVecLToRef<BE>,
121 B: CnvPVecRToRef<BE>,
122 {
123 BE::cnv_pairwise_apply_dft_impl(self, res, res_offset, res_col, a, b, i, j, scratch);
124 }
125}