poulpy_hal/delegates/
vmp_pmat.rs

1use crate::{
2    api::{
3        VmpApplyDft, VmpApplyDftTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAccumulate, VmpApplyDftToDftAccumulateTmpBytes,
4        VmpApplyDftToDftTmpBytes, VmpExtractSelectedRows, VmpPMatAlloc, VmpPMatBytesOf, VmpPrepare, VmpPrepareTmpBytes, VmpZero,
5    },
6    layouts::{
7        Backend, MatZnxBackendRef, Module, ScratchArena, VecZnxBackendRef, VecZnxDftBackendMut, VecZnxDftBackendRef,
8        VecZnxDftToBackendMut, VmpPMatBackendMut, VmpPMatBackendRef, VmpPMatOwned,
9    },
10    oep::HalVmpImpl,
11};
12
13macro_rules! impl_vmp_delegate {
14    ($trait:ty, $($body:item)+) => {
15        impl<B> $trait for Module<B>
16        where
17            B: Backend<ZnxWord = i64> + HalVmpImpl<B>,
18        {
19            $($body)+
20        }
21    };
22}
23
24impl<B: Backend> VmpPMatAlloc<B> for Module<B> {
25    fn vmp_pmat_alloc(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<B> {
26        VmpPMatOwned::<B>::alloc(self.n(), rows, cols_in, cols_out, size)
27    }
28}
29
30impl<B: Backend> VmpPMatBytesOf for Module<B> {
31    fn bytes_of_vmp_pmat(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
32        B::bytes_of_vmp_pmat(self.n(), rows, cols_in, cols_out, size)
33    }
34}
35
36impl_vmp_delegate!(
37    VmpPrepareTmpBytes,
38    fn vmp_prepare_tmp_bytes(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
39        B::vmp_prepare_tmp_bytes(self, rows, cols_in, cols_out, size)
40    }
41);
42
43impl_vmp_delegate!(
44    VmpPrepare<B>,
45    fn vmp_prepare(&self, res: &mut VmpPMatBackendMut<'_, B>, a: &MatZnxBackendRef<'_, B>, scratch: &mut ScratchArena<'_, B>) {
46        B::vmp_prepare(self, res, a, scratch);
47    }
48);
49
50impl_vmp_delegate!(
51    VmpApplyDftTmpBytes,
52    fn vmp_apply_dft_tmp_bytes(
53        &self,
54        res_size: usize,
55        a_size: usize,
56        b_rows: usize,
57        b_cols_in: usize,
58        b_cols_out: usize,
59        b_size: usize,
60    ) -> usize {
61        B::vmp_apply_dft_tmp_bytes(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
62    }
63);
64
65impl_vmp_delegate!(
66    VmpApplyDft<B>,
67    fn vmp_apply_dft<R>(
68        &self,
69        res: &mut R,
70        a: &VecZnxBackendRef<'_, B>,
71        b: &VmpPMatBackendRef<'_, B>,
72        scratch: &mut ScratchArena<'_, B>,
73    ) where
74        R: VecZnxDftToBackendMut<B>,
75    {
76        B::vmp_apply_dft(self, res, a, b, scratch)
77    }
78);
79
80impl_vmp_delegate!(
81    VmpApplyDftToDftTmpBytes,
82    fn vmp_apply_dft_to_dft_tmp_bytes(
83        &self,
84        res_size: usize,
85        a_size: usize,
86        b_rows: usize,
87        b_cols_in: usize,
88        b_cols_out: usize,
89        b_size: usize,
90    ) -> usize {
91        B::vmp_apply_dft_to_dft_tmp_bytes(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
92    }
93);
94
95impl_vmp_delegate!(
96    VmpApplyDftToDft<B>,
97    fn vmp_apply_dft_to_dft(
98        &self,
99        res: &mut VecZnxDftBackendMut<'_, B>,
100        a: &VecZnxDftBackendRef<'_, B>,
101        b: &VmpPMatBackendRef<'_, B>,
102        limb_offset: usize,
103        scratch: &mut ScratchArena<'_, B>,
104    ) {
105        B::vmp_apply_dft_to_dft(self, res, a, b, limb_offset, scratch)
106    }
107);
108
109impl_vmp_delegate!(
110    VmpApplyDftToDftAccumulateTmpBytes,
111    fn vmp_apply_dft_to_dft_accumulate_tmp_bytes(
112        &self,
113        res_size: usize,
114        a_size: usize,
115        b_rows: usize,
116        b_cols_in: usize,
117        b_cols_out: usize,
118        b_size: usize,
119    ) -> usize {
120        B::vmp_apply_dft_to_dft_accumulate_tmp_bytes(self, res_size, a_size, b_rows, b_cols_in, b_cols_out, b_size)
121    }
122);
123
124impl_vmp_delegate!(
125    VmpApplyDftToDftAccumulate<B>,
126    fn vmp_apply_dft_to_dft_accumulate(
127        &self,
128        res: &mut VecZnxDftBackendMut<'_, B>,
129        a: &VecZnxDftBackendRef<'_, B>,
130        b: &VmpPMatBackendRef<'_, B>,
131        limb_offset: usize,
132        scratch: &mut ScratchArena<'_, B>,
133    ) {
134        B::vmp_apply_dft_to_dft_accumulate(self, res, a, b, limb_offset, scratch);
135    }
136);
137
138impl_vmp_delegate!(
139    VmpExtractSelectedRows<B>,
140    fn vmp_extract_selected_rows(
141        &self,
142        res: &mut VmpPMatBackendMut<'_, B>,
143        a: &VmpPMatBackendRef<'_, B>,
144        first_row: usize,
145        row_step: usize,
146    ) {
147        assert_extractable(res, a, first_row, row_step);
148        B::vmp_extract_selected_rows(self, res, a, first_row, row_step);
149    }
150);
151
152/// Rejects a selection the backend kernel must not be handed: mismatched
153/// prepared shapes, a truncation that widens, or a last row that is outside
154/// `a` or whose index overflows.
155///
156/// Enforced here rather than per backend so a kernel may index without bounds
157/// checks in release, and so a new backend inherits the contract.
158fn assert_extractable<B: Backend>(
159    res: &VmpPMatBackendMut<'_, B>,
160    a: &VmpPMatBackendRef<'_, B>,
161    first_row: usize,
162    row_step: usize,
163) {
164    assert!(row_step > 0, "row_step must be positive");
165    assert_eq!(res.n(), a.n(), "res.n(): {} != a.n(): {}", res.n(), a.n());
166    assert_eq!(
167        res.cols_in(),
168        a.cols_in(),
169        "res.cols_in(): {} != a.cols_in(): {}",
170        res.cols_in(),
171        a.cols_in()
172    );
173    assert_eq!(
174        res.cols_out(),
175        a.cols_out(),
176        "res.cols_out(): {} != a.cols_out(): {}",
177        res.cols_out(),
178        a.cols_out()
179    );
180    assert!(res.size() <= a.size(), "res.size(): {} > a.size(): {}", res.size(), a.size());
181    let Some(rows) = res.rows().checked_sub(1) else {
182        return;
183    };
184    let last_row: Option<usize> = rows.checked_mul(row_step).and_then(|o| o.checked_add(first_row));
185    assert!(
186        last_row.is_some_and(|last| last < a.rows()),
187        "selected rows {first_row}..={:?} step {row_step} exceed a.rows(): {}",
188        last_row,
189        a.rows()
190    );
191}
192
193impl_vmp_delegate!(
194    VmpZero<B>,
195    fn vmp_zero(&self, res: &mut VmpPMatBackendMut<'_, B>) {
196        B::vmp_zero(self, res);
197    }
198);
poulpy_hal/delegates/vmp_pmat.rs

poulpy_hal/delegates/
vmp_pmat.rs
