Skip to main content

poulpy_hal/layouts/
svp_ppol.rs

1use std::{
2    fmt,
3    hash::{DefaultHasher, Hasher},
4    marker::PhantomData,
5};
6
7use crate::layouts::{Backend, Data, DataView, DataViewMut, DigestU64, HostDataRef, ScalarZnxShape, ZnxInfos, ZnxView};
8
9/// Prepared (DFT-domain) scalar polynomial for scalar-vector products.
10///
11/// An `SvpPPol` holds a single polynomial in the backend's prepared
12/// representation ([`Backend::ScalarPrep`]). It is used as the left
13/// operand in [`SvpApplyDft`](crate::api::SvpApplyDft) to efficiently
14/// multiply a scalar polynomial by each column of a [`VecZnxDft`](crate::layouts::VecZnxDft).
15///
16/// Create via [`SvpPrepare`](crate::api::SvpPrepare) from a
17/// coefficient-domain [`ScalarZnx`](crate::layouts::ScalarZnx).
18///
19/// Ring degree `n` is always a power of two, so the DFT-domain layout has a
20/// coefficient count that matches vector lane widths relative to buffer alignment.
21#[repr(C)]
22#[derive(PartialEq, Eq, Hash)]
23pub struct SvpPPol<D: Data, B: Backend> {
24    pub data: D,
25    shape: ScalarZnxShape,
26    pub _phantom: PhantomData<B>,
27}
28
29impl<D: HostDataRef, B: Backend> DigestU64 for SvpPPol<D, B> {
30    fn digest_u64(&self) -> u64 {
31        let mut h: DefaultHasher = DefaultHasher::new();
32        h.write(self.data.as_ref());
33        h.write_usize(self.n());
34        h.write_usize(self.cols());
35        h.finish()
36    }
37}
38
39impl<D: HostDataRef, B: Backend> ZnxView for SvpPPol<D, B> {
40    type Scalar = B::ScalarPrep;
41}
42
43impl<D: Data, B: Backend> ZnxInfos for SvpPPol<D, B> {
44    fn cols(&self) -> usize {
45        self.shape.cols()
46    }
47
48    fn rows(&self) -> usize {
49        1
50    }
51
52    fn n(&self) -> usize {
53        self.shape.n()
54    }
55
56    fn size(&self) -> usize {
57        1
58    }
59}
60
61impl<D: Data, B: Backend> SvpPPol<D, B> {
62    pub fn n(&self) -> usize {
63        self.shape.n()
64    }
65
66    pub fn cols(&self) -> usize {
67        self.shape.cols()
68    }
69
70    pub fn shape(&self) -> ScalarZnxShape {
71        self.shape
72    }
73}
74
75impl<D: Data, B: Backend> DataView for SvpPPol<D, B> {
76    type D = D;
77    fn data(&self) -> &Self::D {
78        &self.data
79    }
80}
81
82impl<D: Data, B: Backend> DataViewMut for SvpPPol<D, B> {
83    fn data_mut(&mut self) -> &mut Self::D {
84        &mut self.data
85    }
86}
87
88impl<B: Backend> SvpPPol<<B as Backend>::OwnedBuf, B> {
89    pub fn alloc(n: usize, cols: usize) -> Self {
90        let data: <B as Backend>::OwnedBuf = B::alloc_zeroed_bytes(B::bytes_of_svp_ppol(n, cols));
91        Self {
92            data,
93            shape: ScalarZnxShape::new(n, cols),
94            _phantom: PhantomData,
95        }
96    }
97}
98
99/// Owned `SvpPPol` backed by a backend-owned buffer.
100pub type SvpPPolOwned<B> = SvpPPol<<B as Backend>::OwnedBuf, B>;
101/// Shared backend-native borrow of an `SvpPPol`.
102pub type SvpPPolBackendRef<'a, B> = SvpPPol<<B as Backend>::BufRef<'a>, B>;
103/// Mutable backend-native borrow of an `SvpPPol`.
104pub type SvpPPolBackendMut<'a, B> = SvpPPol<<B as Backend>::BufMut<'a>, B>;
105
106/// Reborrow a mutable backend-native `SvpPPol` view as a shared backend-native view.
107pub fn svp_ppol_backend_ref_from_mut<'a, 'b, B: Backend>(ppol: &'a SvpPPolBackendMut<'b, B>) -> SvpPPolBackendRef<'a, B> {
108    SvpPPol {
109        data: B::view_ref_mut(&ppol.data),
110        shape: ppol.shape,
111        _phantom: PhantomData,
112    }
113}
114
115/// Borrow a backend-owned `SvpPPol` using the backend's native view type.
116pub trait SvpPPolToBackendRef<B: Backend> {
117    fn to_backend_ref(&self) -> SvpPPolBackendRef<'_, B>;
118}
119
120impl<B: Backend> SvpPPolToBackendRef<B> for SvpPPol<B::OwnedBuf, B> {
121    fn to_backend_ref(&self) -> SvpPPolBackendRef<'_, B> {
122        SvpPPol {
123            data: B::view(&self.data),
124            shape: self.shape,
125            _phantom: PhantomData,
126        }
127    }
128}
129
130impl<'b, B: Backend + 'b> SvpPPolToBackendRef<B> for &SvpPPol<B::BufRef<'b>, B> {
131    fn to_backend_ref(&self) -> SvpPPolBackendRef<'_, B> {
132        SvpPPol {
133            data: B::view_ref(&self.data),
134            shape: self.shape,
135            _phantom: PhantomData,
136        }
137    }
138}
139
140/// Reborrow an already backend-borrowed `SvpPPol` as a shared backend-native view.
141pub trait SvpPPolReborrowBackendRef<B: Backend> {
142    fn reborrow_backend_ref(&self) -> SvpPPolBackendRef<'_, B>;
143}
144
145impl<'b, B: Backend + 'b> SvpPPolReborrowBackendRef<B> for SvpPPol<B::BufMut<'b>, B> {
146    fn reborrow_backend_ref(&self) -> SvpPPolBackendRef<'_, B> {
147        svp_ppol_backend_ref_from_mut::<B>(self)
148    }
149}
150
151/// Mutably borrow a backend-owned `SvpPPol` using the backend's native view type.
152pub trait SvpPPolToBackendMut<B: Backend> {
153    fn to_backend_mut(&mut self) -> SvpPPolBackendMut<'_, B>;
154}
155
156impl<B: Backend> SvpPPolToBackendMut<B> for SvpPPol<B::OwnedBuf, B> {
157    fn to_backend_mut(&mut self) -> SvpPPolBackendMut<'_, B> {
158        SvpPPol {
159            data: B::view_mut(&mut self.data),
160            shape: self.shape,
161            _phantom: PhantomData,
162        }
163    }
164}
165
166impl<'b, B: Backend + 'b> SvpPPolToBackendMut<B> for &mut SvpPPol<B::BufMut<'b>, B> {
167    fn to_backend_mut(&mut self) -> SvpPPolBackendMut<'_, B> {
168        SvpPPol {
169            data: B::view_mut_ref(&mut self.data),
170            shape: self.shape,
171            _phantom: PhantomData,
172        }
173    }
174}
175
176/// Reborrow an already backend-borrowed `SvpPPol` as a mutable backend-native view.
177pub trait SvpPPolReborrowBackendMut<B: Backend> {
178    fn reborrow_backend_mut(&mut self) -> SvpPPolBackendMut<'_, B>;
179}
180
181impl<'b, B: Backend + 'b> SvpPPolReborrowBackendMut<B> for SvpPPol<B::BufMut<'b>, B> {
182    fn reborrow_backend_mut(&mut self) -> SvpPPolBackendMut<'_, B> {
183        SvpPPol {
184            data: B::view_mut_ref(&mut self.data),
185            shape: self.shape,
186            _phantom: PhantomData,
187        }
188    }
189}
190
191impl<D: Data, B: Backend> SvpPPol<D, B> {
192    pub fn from_data(data: D, n: usize, cols: usize) -> Self {
193        Self {
194            data,
195            shape: ScalarZnxShape::new(n, cols),
196            _phantom: PhantomData,
197        }
198    }
199}
200
201impl<D: HostDataRef, B: Backend> fmt::Display for SvpPPol<D, B> {
202    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
203        writeln!(f, "SvpPPol(n={}, cols={})", self.n(), self.cols())?;
204
205        for col in 0..self.cols() {
206            writeln!(f, "Column {col}:")?;
207            let coeffs = self.at(col, 0);
208            write!(f, "[")?;
209
210            let max_show = 100;
211            let show_count = coeffs.len().min(max_show);
212
213            for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
214                if i > 0 {
215                    write!(f, ", ")?;
216                }
217                write!(f, "{coeff}")?;
218            }
219
220            if coeffs.len() > max_show {
221                write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
222            }
223
224            writeln!(f, "]")?;
225        }
226        Ok(())
227    }
228}