adele_ring/
batch.rs

1//! `RnsBatch` — the single flat buffer format shared by both backends.
2//!
3//! Having one canonical layout means there is **no reformatting** when switching
4//! between the CPU (rayon) and GPU (wgpu) backends: residues are stored natively
5//! as `u32` (every basis prime is `< 2^16`), so the GPU storage buffer *is* this
6//! struct's `data` and the bytemuck cast is genuinely zero-copy.
7//!
8//! Layout is row-major `[batch_size × n_channels]`: element
9//! `data[b * n_channels + c]` is the residue of item `b` in channel `c`.
10
11use crate::basis::Basis;
12use crate::rns::RnsInt;
13
14/// A batch of RNS values in a flat, backend-agnostic `u32` buffer.
15#[derive(Clone, Debug)]
16pub struct RnsBatch {
17    /// Flat row-major residues: length `batch_size * basis.len()`.
18    pub data: Vec<u32>,
19    /// Number of items `B`.
20    pub batch_size: usize,
21    /// The `K` channels shared by every item.
22    pub basis: Basis,
23}
24
25impl RnsBatch {
26    /// Allocate a zeroed batch (alias of [`RnsBatch::zeros`]).
27    pub fn new(batch_size: usize, basis: Basis) -> Self {
28        Self::zeros(batch_size, basis)
29    }
30
31    /// Allocate a batch with all residues set to zero.
32    pub fn zeros(batch_size: usize, basis: Basis) -> Self {
33        let k = basis.len();
34        RnsBatch { data: vec![0; batch_size * k], batch_size, basis }
35    }
36
37    /// Number of channels `K`.
38    #[inline]
39    pub fn channels_len(&self) -> usize {
40        self.basis.len()
41    }
42
43    /// Residue of item `b` in channel `c`.
44    #[inline]
45    pub fn get(&self, b: usize, c: usize) -> u32 {
46        self.data[b * self.basis.len() + c]
47    }
48
49    /// Set the residue of item `b` in channel `c`.
50    #[inline]
51    pub fn set(&mut self, b: usize, c: usize, val: u32) {
52        let k = self.basis.len();
53        self.data[b * k + c] = val;
54    }
55
56    /// Pack a slice of [`RnsInt`] values into a batch.
57    ///
58    /// Panics if `items` is empty (no basis to infer) or if the items do not all
59    /// share the same basis.
60    pub fn from_rns_ints(items: &[RnsInt]) -> Self {
61        assert!(!items.is_empty(), "cannot build an RnsBatch from zero items");
62        let basis = items[0].basis.clone();
63        let k = basis.len();
64        let mut data = Vec::with_capacity(items.len() * k);
65        for item in items {
66            debug_assert_eq!(item.basis, basis, "all items must share the basis");
67            debug_assert_eq!(item.residues.len(), k);
68            data.extend_from_slice(&item.residues);
69        }
70        RnsBatch { data, batch_size: items.len(), basis }
71    }
72
73    /// Unpack the batch back into individual [`RnsInt`] values.
74    pub fn to_rns_ints(&self) -> Vec<RnsInt> {
75        let k = self.basis.len();
76        (0..self.batch_size)
77            .map(|b| {
78                let residues = self.data[b * k..(b + 1) * k].to_vec();
79                RnsInt::from_residues(residues, self.basis.clone())
80            })
81            .collect()
82    }
83
84    /// Pack for GPU upload: residues as little-endian `u32` bytes (zero-copy).
85    pub fn as_u32_bytes(&self) -> Vec<u8> {
86        bytemuck::cast_slice(&self.data).to_vec()
87    }
88
89    /// Rebuild a batch from a flat `u32` slice downloaded from the GPU.
90    pub fn from_u32(values: &[u32], batch_size: usize, basis: Basis) -> Self {
91        RnsBatch { data: values.to_vec(), batch_size, basis }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn get_set_roundtrip() {
101        let b = Basis::standard();
102        let mut batch = RnsBatch::zeros(4, b);
103        batch.set(2, 3, 7);
104        assert_eq!(batch.get(2, 3), 7);
105        assert_eq!(batch.get(0, 0), 0);
106    }
107
108    #[test]
109    fn pack_unpack_roundtrip() {
110        let b = Basis::standard();
111        let items = vec![
112            RnsInt::from_i64(123, b.clone()),
113            RnsInt::from_i64(456, b.clone()),
114            RnsInt::from_i64(789, b.clone()),
115        ];
116        let batch = RnsBatch::from_rns_ints(&items);
117        assert_eq!(batch.batch_size, 3);
118        let back = batch.to_rns_ints();
119        for (a, bb) in items.iter().zip(back.iter()) {
120            assert_eq!(a.to_bigint(), bb.to_bigint());
121        }
122    }
123
124    #[test]
125    fn u32_byte_layout() {
126        let b = Basis::standard();
127        let mut batch = RnsBatch::zeros(1, b);
128        batch.set(0, 0, 1);
129        batch.set(0, 1, 2);
130        let bytes = batch.as_u32_bytes();
131        assert_eq!(bytes[0], 1);
132        assert_eq!(bytes[4], 2);
133    }
134}
adele_ring/batch.rs

adele_ring/
batch.rs
