Skip to main content

poulpy_hal/layouts/
mod.rs

1//! Data layout types and trait definitions for the hardware abstraction layer.
2//!
3//! This module aggregates all layout-related types and re-exports them from
4//! their respective sub-modules, including convolution kernels, matrix and
5//! vector representations over polynomial rings, serialization support,
6//! statistical utilities, and scratch-space management.
7//!
8//! It also defines the shared storage trait aliases used throughout the crate.
9//! `Data` models backend-owned storage in the abstract, while
10//! `HostDataRef`/`HostDataMut` capture host-byte-readable buffers for the
11//! portions of the API that still require direct byte access.
12
13mod convolution;
14mod encoding;
15mod mat_znx;
16mod module;
17mod scalar_znx;
18mod scratch;
19mod scratch_views;
20mod serialization;
21mod stats;
22mod svp_ppol;
23mod vec_znx;
24mod vec_znx_big;
25mod vec_znx_dft;
26mod vmp_pmat;
27mod znx_base;
28
29pub use convolution::*;
30pub use mat_znx::*;
31pub use module::*;
32pub use scalar_znx::*;
33pub use scratch::*;
34pub use scratch_views::*;
35pub use serialization::*;
36pub use stats::*;
37pub use svp_ppol::*;
38pub use vec_znx::*;
39pub use vec_znx_big::*;
40pub use vec_znx_dft::*;
41pub use vmp_pmat::*;
42pub use znx_base::*;
43
44use anyhow::Result;
45use std::ptr::NonNull;
46
47use crate::oep::HalModuleImpl;
48
49/// Base trait alias for all data containers.
50///
51/// Requires equality comparison ([`PartialEq`], [`Eq`]), a known size at
52/// compile time ([`Sized`]), and a default value ([`Default`]). Every
53/// layout type that holds raw data must satisfy at least this bound.
54pub trait Data = PartialEq + Eq + Sized + Default;
55
56/// Trait alias for read-only host-byte-accessible containers.
57///
58/// Extends [`Data`] with byte-level shared access via [`AsRef<[u8]>`] and
59/// thread-safe sharing via [`Sync`]. Types satisfying this bound can be
60/// borrowed immutably and read across threads.
61pub trait HostDataRef = Data + AsRef<[u8]> + Sync;
62
63/// Trait alias for mutable host-byte-accessible containers.
64///
65/// Extends [`HostDataRef`] with byte-level mutable access via [`AsMut<[u8]>`]
66/// and cross-thread transfer via [`Send`]. Types satisfying this bound
67/// support in-place modification and can be moved between threads.
68pub trait HostDataMut = HostDataRef + AsMut<[u8]> + Send;
69
70mod private {
71    pub trait Sealed {}
72}
73
74/// Sealed trait identifying the residency of a [`Backend`]'s buffers.
75///
76/// Implemented only by [`Host`] and [`Device`]. Each [`Backend`] declares
77/// its residency via its [`Backend::Location`] associated type, which lets
78/// generic code discriminate host- and device-resident backends at the
79/// type level.
80pub trait Location: private::Sealed {}
81
82/// Marker type for host-resident buffers.
83#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
84pub struct Host;
85
86/// Marker type for device-resident buffers.
87#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
88pub struct Device;
89
90impl private::Sealed for Host {}
91impl private::Sealed for Device {}
92impl Location for Host {}
93impl Location for Device {}
94
95/// Convenience marker for host-resident backends.
96pub trait HostBackend: Backend<Location = Host> {}
97impl<BE: Backend<Location = Host>> HostBackend for BE {}
98
99/// Convenience marker for host-resident backends whose borrowed views are directly readable and writable as host bytes.
100pub trait HostVisibleBackend: HostBackend
101where
102    for<'a> Self::BufRef<'a>: AsRef<[u8]>,
103    for<'a> Self::BufMut<'a>: AsRef<[u8]> + AsMut<[u8]>,
104{
105}
106
107impl<BE> HostVisibleBackend for BE
108where
109    BE: HostBackend,
110    for<'a> BE::BufRef<'a>: AsRef<[u8]>,
111    for<'a> BE::BufMut<'a>: AsRef<[u8]> + AsMut<[u8]>,
112{
113}
114
115/// Minimal host-resident backend used as the default backend adapter for
116/// host-visible byte-slice views in generic helper code.
117#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
118pub struct HostBytesBackend;
119
120impl Backend for HostBytesBackend {
121    type ScalarBig = i128;
122    type ScalarPrep = i64;
123    type OwnedBuf = Vec<u8>;
124    type BufRef<'a> = &'a [u8];
125    type BufMut<'a> = &'a mut [u8];
126    type Handle = ();
127    type Location = Host;
128
129    fn alloc_bytes(len: usize) -> Self::OwnedBuf {
130        crate::alloc_aligned::<u8>(len)
131    }
132
133    fn alloc_zeroed_bytes(len: usize) -> Self::OwnedBuf {
134        crate::alloc_aligned::<u8>(len)
135    }
136
137    fn from_host_bytes(bytes: &[u8]) -> Self::OwnedBuf {
138        let mut out = crate::alloc_aligned::<u8>(bytes.len());
139        out.copy_from_slice(bytes);
140        out
141    }
142
143    fn from_bytes(bytes: Vec<u8>) -> Self::OwnedBuf {
144        if crate::is_aligned(bytes.as_ptr()) {
145            bytes
146        } else {
147            let mut out = crate::alloc_aligned::<u8>(bytes.len());
148            out.copy_from_slice(&bytes);
149            out
150        }
151    }
152
153    fn to_host_bytes(buf: &Self::OwnedBuf) -> Vec<u8> {
154        buf.clone()
155    }
156
157    fn copy_to_host(buf: &Self::OwnedBuf, dst: &mut [u8]) {
158        assert!(
159            buf.len() >= dst.len(),
160            "backend buffer length {} is smaller than destination host slice length {}",
161            buf.len(),
162            dst.len()
163        );
164        dst.copy_from_slice(&buf[..dst.len()]);
165    }
166
167    fn copy_from_host(buf: &mut Self::OwnedBuf, src: &[u8]) {
168        assert!(
169            buf.len() >= src.len(),
170            "backend buffer length {} is smaller than source host slice length {}",
171            buf.len(),
172            src.len()
173        );
174        let src_len = src.len();
175        buf[..src_len].copy_from_slice(src);
176        buf[src_len..].fill(0);
177    }
178
179    fn len_bytes(buf: &Self::OwnedBuf) -> usize {
180        buf.len()
181    }
182
183    fn view(buf: &Self::OwnedBuf) -> Self::BufRef<'_> {
184        buf.as_slice()
185    }
186
187    fn view_ref<'a, 'b>(buf: &'a Self::BufRef<'b>) -> Self::BufRef<'a>
188    where
189        Self: 'b,
190    {
191        buf
192    }
193
194    fn view_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>) -> Self::BufRef<'a>
195    where
196        Self: 'b,
197    {
198        buf
199    }
200
201    fn view_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>) -> Self::BufMut<'a>
202    where
203        Self: 'b,
204    {
205        buf
206    }
207
208    fn view_mut(buf: &mut Self::OwnedBuf) -> Self::BufMut<'_> {
209        buf.as_mut_slice()
210    }
211
212    fn region(buf: &Self::OwnedBuf, offset: usize, len: usize) -> Self::BufRef<'_> {
213        &buf[offset..offset + len]
214    }
215
216    fn region_mut(buf: &mut Self::OwnedBuf, offset: usize, len: usize) -> Self::BufMut<'_> {
217        &mut buf[offset..offset + len]
218    }
219
220    fn region_ref<'a, 'b>(buf: &'a Self::BufRef<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
221    where
222        Self: 'b,
223    {
224        &buf[offset..offset + len]
225    }
226
227    fn region_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
228    where
229        Self: 'b,
230    {
231        &buf[offset..offset + len]
232    }
233
234    fn region_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufMut<'a>
235    where
236        Self: 'b,
237    {
238        &mut buf[offset..offset + len]
239    }
240
241    unsafe fn destroy(_handle: NonNull<Self::Handle>) {}
242}
243
244unsafe impl HalModuleImpl<HostBytesBackend> for HostBytesBackend {
245    fn new(n: u64) -> crate::layouts::Module<Self> {
246        assert!(n.is_power_of_two(), "n must be a power of two, got {n}");
247        unsafe { crate::layouts::Module::from_nonnull(NonNull::dangling(), n) }
248    }
249}
250
251/// Convenience marker for device-resident backends.
252pub trait DeviceBackend: Backend<Location = Device> {}
253impl<BE: Backend<Location = Device>> DeviceBackend for BE {}
254
255/// Deep-clone a borrowed layout into a fully owned variant.
256///
257/// Unlike the standard [`Clone`] trait, `ToOwnedDeep` is intended for
258/// types that may borrow their underlying storage. Calling
259/// [`to_owned_deep`](ToOwnedDeep::to_owned_deep) produces an independent
260/// copy whose lifetime is not tied to the original.
261pub trait ToOwnedDeep {
262    type Owned;
263    fn to_owned_deep(&self) -> Self::Owned;
264}
265
266/// Compute a `u64` hash digest of a layout's contents.
267///
268/// Provides a lightweight fingerprint suitable for fast equality checks
269/// and debugging. This is **not** cryptographically secure; it is a
270/// convenience mechanism for detecting whether two values hold identical
271/// data without performing a full byte-by-byte comparison.
272pub trait DigestU64 {
273    fn digest_u64(&self) -> u64;
274}
275
276/// Backend-owned byte buffer type alias.
277pub type OwnedBuf<BE> = <BE as Backend>::OwnedBuf;
278
279/// Cross-backend buffer transfer into the destination backend `Self`.
280///
281/// This is intentionally destination-owned so the canonical public API can
282/// hang off `Module<To>` as `upload_*` / `download_*`.
283///
284/// Each concrete backend pair must provide an explicit impl. Two restricted
285/// blankets are provided for [`HostBytesBackend`] so that test/bench helpers
286/// that use it as a staging type continue to work without boilerplate:
287/// - any host `Vec<u8>` backend → `HostBytesBackend`
288/// - `HostBytesBackend` → any host `Vec<u8>` backend
289///
290/// All other backend-to-backend transfers (e.g. `FFT64Ref` ↔ `NTT120Ref`,
291/// `FFT64Ref` → `FFT64Avx`) must be implemented explicitly in the respective
292/// backend crates.
293pub trait TransferFrom<From: Backend>: Backend {
294    /// Transfers a buffer owned by `From` into `Self`.
295    fn transfer_buf(src: &From::OwnedBuf) -> Self::OwnedBuf;
296}
297
298impl<T: Backend<Location = Host, OwnedBuf = Vec<u8>>> TransferFrom<HostBytesBackend> for T {
299    fn transfer_buf(src: &Vec<u8>) -> Self::OwnedBuf {
300        T::from_host_bytes(src)
301    }
302}
303
304/// Implement a backend marker by forwarding all storage- and handle-level
305/// behavior to an existing backend.
306///
307/// This is useful for proof or delegating backends that want to remain a
308/// distinct backend type while reusing the same owned buffer, borrowed views,
309/// scalar types, and handle representation as a source backend.
310#[macro_export]
311macro_rules! impl_backend_from {
312    ($be:ty, $from:ty) => {
313        impl poulpy_hal::layouts::Backend for $be {
314            type ScalarBig = <$from as poulpy_hal::layouts::Backend>::ScalarBig;
315            type ScalarPrep = <$from as poulpy_hal::layouts::Backend>::ScalarPrep;
316            type OwnedBuf = <$from as poulpy_hal::layouts::Backend>::OwnedBuf;
317            type BufRef<'a> = <$from as poulpy_hal::layouts::Backend>::BufRef<'a>;
318            type BufMut<'a> = <$from as poulpy_hal::layouts::Backend>::BufMut<'a>;
319            type Handle = <$from as poulpy_hal::layouts::Backend>::Handle;
320            type Location = <$from as poulpy_hal::layouts::Backend>::Location;
321
322            fn alloc_bytes(len: usize) -> Self::OwnedBuf {
323                <$from as poulpy_hal::layouts::Backend>::alloc_bytes(len)
324            }
325
326            fn alloc_zeroed_bytes(len: usize) -> Self::OwnedBuf {
327                <$from as poulpy_hal::layouts::Backend>::alloc_zeroed_bytes(len)
328            }
329
330            fn from_host_bytes(bytes: &[u8]) -> Self::OwnedBuf {
331                <$from as poulpy_hal::layouts::Backend>::from_host_bytes(bytes)
332            }
333
334            fn from_bytes(bytes: Vec<u8>) -> Self::OwnedBuf {
335                <$from as poulpy_hal::layouts::Backend>::from_bytes(bytes)
336            }
337
338            fn to_host_bytes(buf: &Self::OwnedBuf) -> Vec<u8> {
339                <$from as poulpy_hal::layouts::Backend>::to_host_bytes(buf)
340            }
341
342            fn copy_to_host(buf: &Self::OwnedBuf, dst: &mut [u8]) {
343                <$from as poulpy_hal::layouts::Backend>::copy_to_host(buf, dst)
344            }
345
346            fn copy_from_host(buf: &mut Self::OwnedBuf, src: &[u8]) {
347                <$from as poulpy_hal::layouts::Backend>::copy_from_host(buf, src)
348            }
349
350            fn len_bytes(buf: &Self::OwnedBuf) -> usize {
351                <$from as poulpy_hal::layouts::Backend>::len_bytes(buf)
352            }
353
354            fn view(buf: &Self::OwnedBuf) -> Self::BufRef<'_> {
355                <$from as poulpy_hal::layouts::Backend>::view(buf)
356            }
357
358            fn view_ref<'a, 'b>(buf: &'a Self::BufRef<'b>) -> Self::BufRef<'a>
359            where
360                Self: 'b,
361            {
362                <$from as poulpy_hal::layouts::Backend>::view_ref(buf)
363            }
364
365            fn view_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>) -> Self::BufRef<'a>
366            where
367                Self: 'b,
368            {
369                <$from as poulpy_hal::layouts::Backend>::view_ref_mut(buf)
370            }
371
372            fn view_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>) -> Self::BufMut<'a>
373            where
374                Self: 'b,
375            {
376                <$from as poulpy_hal::layouts::Backend>::view_mut_ref(buf)
377            }
378
379            fn view_mut(buf: &mut Self::OwnedBuf) -> Self::BufMut<'_> {
380                <$from as poulpy_hal::layouts::Backend>::view_mut(buf)
381            }
382
383            fn region(buf: &Self::OwnedBuf, offset: usize, len: usize) -> Self::BufRef<'_> {
384                <$from as poulpy_hal::layouts::Backend>::region(buf, offset, len)
385            }
386
387            fn region_mut(buf: &mut Self::OwnedBuf, offset: usize, len: usize) -> Self::BufMut<'_> {
388                <$from as poulpy_hal::layouts::Backend>::region_mut(buf, offset, len)
389            }
390
391            fn region_ref<'a, 'b>(buf: &'a Self::BufRef<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
392            where
393                Self: 'b,
394            {
395                <$from as poulpy_hal::layouts::Backend>::region_ref(buf, offset, len)
396            }
397
398            fn region_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
399            where
400                Self: 'b,
401            {
402                <$from as poulpy_hal::layouts::Backend>::region_ref_mut(buf, offset, len)
403            }
404
405            fn region_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufMut<'a>
406            where
407                Self: 'b,
408            {
409                <$from as poulpy_hal::layouts::Backend>::region_mut_ref(buf, offset, len)
410            }
411
412            unsafe fn destroy(handle: std::ptr::NonNull<Self::Handle>) {
413                <$from as poulpy_hal::layouts::Backend>::destroy(handle)
414            }
415        }
416    };
417}
418
419#[derive(Clone, Copy, Debug)]
420pub struct NoiseInfos {
421    pub k: usize,
422    pub sigma: f64,
423    pub bound: f64,
424}
425
426impl NoiseInfos {
427    pub fn new(k: usize, sigma: f64, bound: f64) -> Result<Self> {
428        anyhow::ensure!(sigma.is_sign_positive(), "sigma must be positive");
429        anyhow::ensure!(sigma >= 1.0, "sigma must be greater or equal to 1");
430        anyhow::ensure!(bound >= sigma, "bound: {bound} must be greater or equal to sigma: {sigma}");
431        Ok(Self { k, sigma, bound })
432    }
433
434    pub fn target_limb_and_scale(&self, base2k: usize) -> (usize, f64) {
435        let limb: usize = self.k.div_ceil(base2k) - 1;
436        let scale: f64 = (((limb + 1) * base2k - self.k) as f64).exp2();
437        (limb, scale)
438    }
439}