adele_ring/
backend.rs

1//! The [`ArithmeticBackend`] trait and the [`Executor`] that selects between CPU
2//! and GPU at runtime. All batch math in the crate flows through the Executor —
3//! it never hard-codes a backend.
4
5use std::sync::{Arc, OnceLock};
6
7use num_bigint::BigInt;
8
9use crate::batch::RnsBatch;
10use crate::cpu::CpuBackend;
11use crate::gpu::GpuBackend;
12
13/// A backend that can perform elementwise RNS arithmetic over a [`RnsBatch`].
14pub trait ArithmeticBackend: Send + Sync {
15    /// Elementwise add: `result[b][c] = (a[b][c] + b[b][c]) % m[c]`.
16    fn batch_add(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch;
17
18    /// Elementwise subtract: `result[b][c] = (a[b][c] + m[c] - b[b][c]) % m[c]`.
19    fn batch_sub(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch;
20
21    /// Elementwise multiply: `result[b][c] = (a[b][c] * b[b][c]) % m[c]`.
22    fn batch_mul(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch;
23
24    /// CRT-reconstruct every item in the batch to its **signed** (balanced) value.
25    fn batch_crt(&self, batch: &RnsBatch) -> Vec<BigInt>;
26
27    /// Backend name for diagnostics.
28    fn name(&self) -> &'static str;
29}
30
31/// Runtime dispatcher between the CPU and (optional) GPU backends.
32pub struct Executor {
33    cpu: Arc<CpuBackend>,
34    gpu: Option<Arc<GpuBackend>>,
35    /// Batches smaller than this use the CPU even when a GPU is present, because
36    /// the GPU's upload/dispatch/download round-trip (~100µs) dominates for small
37    /// inputs. Public so callers can tune it for their hardware.
38    pub gpu_threshold: usize,
39}
40
41impl Executor {
42    /// Probe for a GPU and build the executor. This is the only constructor.
43    pub fn init() -> Self {
44        let cpu = Arc::new(CpuBackend::new());
45        let gpu = GpuBackend::try_init().ok().map(Arc::new);
46        match &gpu {
47            Some(g) => log::info!("adele-ring: GPU backend active ({})", g.adapter_name()),
48            None => log::info!("adele-ring: no GPU found, using CPU backend"),
49        }
50        Self {
51            cpu,
52            gpu,
53            gpu_threshold: 128,
54        }
55    }
56
57    /// Whether a GPU backend is available.
58    pub fn has_gpu(&self) -> bool {
59        self.gpu.is_some()
60    }
61
62    /// Borrow the CPU backend directly (used by benchmarks).
63    pub fn cpu(&self) -> &CpuBackend {
64        &self.cpu
65    }
66
67    /// Borrow the GPU backend if present (used by benchmarks).
68    pub fn gpu(&self) -> Option<&GpuBackend> {
69        self.gpu.as_deref()
70    }
71
72    /// Pick the best backend for a given batch size.
73    fn select(&self, batch_size: usize) -> &dyn ArithmeticBackend {
74        match &self.gpu {
75            Some(g) if batch_size >= self.gpu_threshold => g.as_ref(),
76            _ => self.cpu.as_ref(),
77        }
78    }
79
80    /// Elementwise batch addition.
81    pub fn add(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
82        self.select(a.batch_size).batch_add(a, b)
83    }
84
85    /// Elementwise batch subtraction.
86    pub fn sub(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
87        self.select(a.batch_size).batch_sub(a, b)
88    }
89
90    /// Elementwise batch multiplication.
91    pub fn mul(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
92        self.select(a.batch_size).batch_mul(a, b)
93    }
94
95    /// CRT reconstruction — always CPU-side (Garner's algorithm is sequential).
96    pub fn crt(&self, batch: &RnsBatch) -> Vec<BigInt> {
97        self.cpu.batch_crt(batch)
98    }
99}
100
101impl Default for Executor {
102    fn default() -> Self {
103        Self::init()
104    }
105}
106
107static EXECUTOR: OnceLock<Executor> = OnceLock::new();
108
109/// The lazily-initialized, crate-wide executor.
110pub fn executor() -> &'static Executor {
111    EXECUTOR.get_or_init(Executor::init)
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::basis::Basis;
118    use crate::rns::RnsInt;
119
120    #[test]
121    fn executor_adds_batches() {
122        let b = Basis::standard();
123        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(7, b.clone()); 200]);
124        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(35, b.clone()); 200]);
125        let sum = executor().add(&x, &y);
126        for item in sum.to_rns_ints() {
127            assert_eq!(item.to_bigint(), BigInt::from(42));
128        }
129    }
130
131    #[test]
132    fn executor_subtracts_batches() {
133        let b = Basis::standard();
134        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(35, b.clone()); 200]);
135        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(42, b.clone()); 200]);
136        let diff = executor().sub(&x, &y);
137        for item in diff.to_rns_ints() {
138            assert_eq!(item.to_bigint(), BigInt::from(-7));
139        }
140    }
141
142    #[test]
143    fn cpu_gpu_identical() {
144        let exec = Executor::init();
145        if !exec.has_gpu() {
146            return;
147        }
148        let b = Basis::standard();
149        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(123, b.clone()); 256]);
150        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(456, b.clone()); 256]);
151        let cpu = exec.cpu().batch_add(&x, &y);
152        let gpu = exec.gpu().unwrap().batch_add(&x, &y);
153        assert_eq!(cpu.data, gpu.data);
154    }
155}
adele_ring/backend.rs

adele_ring/
backend.rs
