adele_ring/
cpu.rs

1//! `CpuBackend` — the always-available backend, using rayon to parallelize over
2//! batch items (and, for single values, over channels above a threshold).
3
4use num_bigint::BigInt;
5use rayon::prelude::*;
6
7use crate::backend::ArithmeticBackend;
8use crate::batch::RnsBatch;
9use crate::rns::{add_channel, crt_balanced, mul_channel, sub_channel, RnsInt};
10use crate::RAYON_CHANNEL_THRESHOLD;
11
12/// CPU backend backed by a dedicated rayon thread pool.
13pub struct CpuBackend {
14    pool: rayon::ThreadPool,
15}
16
17impl Default for CpuBackend {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl CpuBackend {
24    /// Build a backend with a thread pool spanning all logical cores.
25    pub fn new() -> Self {
26        Self {
27            pool: rayon::ThreadPoolBuilder::new()
28                .num_threads(0) // 0 = use all logical cores
29                .thread_name(|i| format!("adele-ring-cpu-{i}"))
30                .build()
31                .expect("rayon pool init failed"),
32        }
33    }
34
35    /// Single-value addition: parallel over channels only when `k` is large.
36    ///
37    /// Below [`RAYON_CHANNEL_THRESHOLD`] channels, rayon's per-task overhead
38    /// (~50ns) exceeds the cost of a channel op (~1ns), so we stay sequential.
39    pub fn rns_add_single(&self, a: &RnsInt, b: &RnsInt) -> RnsInt {
40        let moduli = a.basis.moduli();
41        let k = a.basis.len();
42        let residues: Vec<u32> = if k >= RAYON_CHANNEL_THRESHOLD {
43            self.pool.install(|| {
44                a.residues
45                    .par_iter()
46                    .zip(b.residues.par_iter())
47                    .zip(moduli.par_iter())
48                    .map(|((&av, &bv), &m)| add_channel(av, bv, m))
49                    .collect()
50            })
51        } else {
52            a.residues
53                .iter()
54                .zip(b.residues.iter())
55                .zip(moduli.iter())
56                .map(|((&av, &bv), &m)| add_channel(av, bv, m))
57                .collect()
58        };
59        RnsInt::from_residues(residues, a.basis.clone())
60    }
61
62    fn elementwise(
63        &self,
64        a: &RnsBatch,
65        b: &RnsBatch,
66        f: impl Fn(u32, u32, u32) -> u32 + Sync + Send,
67    ) -> RnsBatch {
68        let k = a.basis.len();
69        let moduli = a.basis.moduli();
70        let mut result = RnsBatch::zeros(a.batch_size, a.basis.clone());
71        self.pool.install(|| {
72            result
73                .data
74                .par_chunks_mut(k)
75                .enumerate()
76                .for_each(|(b_idx, out_row)| {
77                    let base = b_idx * k;
78                    for c in 0..k {
79                        out_row[c] = f(a.data[base + c], b.data[base + c], moduli[c]);
80                    }
81                });
82        });
83        result
84    }
85}
86
87impl ArithmeticBackend for CpuBackend {
88    fn batch_add(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
89        self.elementwise(a, b, add_channel)
90    }
91
92    fn batch_sub(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
93        self.elementwise(a, b, sub_channel)
94    }
95
96    fn batch_mul(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
97        self.elementwise(a, b, mul_channel)
98    }
99
100    fn batch_crt(&self, batch: &RnsBatch) -> Vec<BigInt> {
101        let k = batch.basis.len();
102        let moduli = batch.basis.moduli();
103        self.pool.install(|| {
104            (0..batch.batch_size)
105                .into_par_iter()
106                .map(|b| crt_balanced(&batch.data[b * k..(b + 1) * k], moduli))
107                .collect()
108        })
109    }
110
111    fn name(&self) -> &'static str {
112        "cpu-rayon"
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::basis::Basis;
120
121    #[test]
122    fn batch_add_matches_scalar() {
123        let b = Basis::standard();
124        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(123, b.clone()); 64]);
125        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(456, b.clone()); 64]);
126        let cpu = CpuBackend::new();
127        for item in cpu.batch_add(&x, &y).to_rns_ints() {
128            assert_eq!(item.to_bigint(), BigInt::from(579));
129        }
130    }
131
132    #[test]
133    fn batch_sub_matches_scalar() {
134        let b = Basis::standard();
135        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(123, b.clone()); 64]);
136        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(456, b.clone()); 64]);
137        let cpu = CpuBackend::new();
138        for item in cpu.batch_sub(&x, &y).to_rns_ints() {
139            assert_eq!(item.to_bigint(), BigInt::from(-333));
140        }
141    }
142
143    #[test]
144    fn batch_mul_matches_scalar() {
145        let b = Basis::standard();
146        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(123, b.clone()); 64]);
147        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(456, b.clone()); 64]);
148        let cpu = CpuBackend::new();
149        for item in cpu.batch_mul(&x, &y).to_rns_ints() {
150            assert_eq!(item.to_bigint(), BigInt::from(123 * 456));
151        }
152    }
153
154    #[test]
155    fn single_add() {
156        let b = Basis::standard();
157        let cpu = CpuBackend::new();
158        let r = cpu.rns_add_single(&RnsInt::from_i64(10, b.clone()), &RnsInt::from_i64(32, b));
159        assert_eq!(r.to_bigint(), BigInt::from(42));
160    }
161
162    #[test]
163    fn batch_crt_is_balanced() {
164        let b = Basis::standard();
165        let x = RnsBatch::from_rns_ints(&[RnsInt::from_i64(-5, b.clone()), RnsInt::from_i64(9, b)]);
166        let cpu = CpuBackend::new();
167        let crt = cpu.batch_crt(&x);
168        assert_eq!(crt, vec![BigInt::from(-5), BigInt::from(9)]);
169    }
170}
adele_ring/cpu.rs

adele_ring/
cpu.rs
