adele_ring/
gpu.rs

1//! `GpuBackend` — wgpu compute-shader implementation of [`ArithmeticBackend`].
2//!
3//! wgpu is always compiled in (not feature-gated). At startup [`GpuBackend::try_init`]
4//! probes for a compatible adapter; if none is found it returns `Err` and the
5//! [`crate::backend::Executor`] transparently falls back to the CPU backend.
6//!
7//! Each shader thread handles one `(batch_item × channel)` pair. The buffer
8//! layout is identical to [`crate::batch::RnsBatch`], so there is no reformatting
9//! on upload beyond the `u64 -> u32` narrowing (safe because all moduli `< 2^16`).
10
11use bytemuck::{Pod, Zeroable};
12use num_bigint::BigInt;
13use wgpu::util::DeviceExt;
14
15use crate::backend::ArithmeticBackend;
16use crate::batch::RnsBatch;
17use crate::error::GpuError;
18use crate::rns::crt_balanced;
19
20#[repr(C)]
21#[derive(Clone, Copy, Pod, Zeroable)]
22struct Params {
23    batch_size: u32,
24    n_channels: u32,
25    _pad: [u32; 2], // pad to 16 bytes for std140 uniform layout
26}
27
28/// GPU backend holding a device, queue, and the pre-built compute pipelines.
29pub struct GpuBackend {
30    device: wgpu::Device,
31    queue: wgpu::Queue,
32    bind_group_layout: wgpu::BindGroupLayout,
33    add_pipeline: wgpu::ComputePipeline,
34    sub_pipeline: wgpu::ComputePipeline,
35    mul_pipeline: wgpu::ComputePipeline,
36    adapter_info: wgpu::AdapterInfo,
37}
38
39impl GpuBackend {
40    /// Probe for a GPU and build the pipelines. Blocks on async init.
41    pub fn try_init() -> Result<Self, GpuError> {
42        pollster::block_on(Self::try_init_async())
43    }
44
45    async fn try_init_async() -> Result<Self, GpuError> {
46        let instance = wgpu::Instance::default();
47        let adapter = instance
48            .request_adapter(&wgpu::RequestAdapterOptions {
49                power_preference: wgpu::PowerPreference::HighPerformance,
50                force_fallback_adapter: false,
51                compatible_surface: None,
52            })
53            .await
54            .ok_or(GpuError::NoAdapter)?;
55
56        let adapter_info = adapter.get_info();
57
58        let (device, queue) = adapter
59            .request_device(
60                &wgpu::DeviceDescriptor {
61                    label: Some("adele-ring-device"),
62                    required_features: wgpu::Features::empty(),
63                    required_limits: wgpu::Limits::downlevel_defaults(),
64                    memory_hints: wgpu::MemoryHints::Performance,
65                },
66                None,
67            )
68            .await?;
69
70        let add_shader =
71            device.create_shader_module(wgpu::include_wgsl!("../shaders/rns_add.wgsl"));
72        let sub_shader =
73            device.create_shader_module(wgpu::include_wgsl!("../shaders/rns_sub.wgsl"));
74        let mul_shader =
75            device.create_shader_module(wgpu::include_wgsl!("../shaders/rns_mul.wgsl"));
76
77        let bind_group_layout = Self::make_bind_group_layout(&device);
78        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
79            label: Some("adele-ring-pipeline-layout"),
80            bind_group_layouts: &[&bind_group_layout],
81            push_constant_ranges: &[],
82        });
83
84        let add_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
85            label: Some("rns-add"),
86            layout: Some(&pipeline_layout),
87            module: &add_shader,
88            entry_point: "main",
89            compilation_options: Default::default(),
90            cache: None,
91        });
92        let sub_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
93            label: Some("rns-sub"),
94            layout: Some(&pipeline_layout),
95            module: &sub_shader,
96            entry_point: "main",
97            compilation_options: Default::default(),
98            cache: None,
99        });
100        let mul_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
101            label: Some("rns-mul"),
102            layout: Some(&pipeline_layout),
103            module: &mul_shader,
104            entry_point: "main",
105            compilation_options: Default::default(),
106            cache: None,
107        });
108
109        Ok(Self {
110            device,
111            queue,
112            bind_group_layout,
113            add_pipeline,
114            sub_pipeline,
115            mul_pipeline,
116            adapter_info,
117        })
118    }
119
120    /// Human-readable adapter name (e.g. "NVIDIA GeForce RTX 4080").
121    pub fn adapter_name(&self) -> &str {
122        &self.adapter_info.name
123    }
124
125    fn make_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
126        let storage = |read_only: bool| wgpu::BindingType::Buffer {
127            ty: wgpu::BufferBindingType::Storage { read_only },
128            has_dynamic_offset: false,
129            min_binding_size: None,
130        };
131        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
132            label: Some("adele-ring-bgl"),
133            entries: &[
134                wgpu::BindGroupLayoutEntry {
135                    binding: 0,
136                    visibility: wgpu::ShaderStages::COMPUTE,
137                    ty: wgpu::BindingType::Buffer {
138                        ty: wgpu::BufferBindingType::Uniform,
139                        has_dynamic_offset: false,
140                        min_binding_size: None,
141                    },
142                    count: None,
143                },
144                wgpu::BindGroupLayoutEntry {
145                    binding: 1,
146                    visibility: wgpu::ShaderStages::COMPUTE,
147                    ty: storage(true),
148                    count: None,
149                },
150                wgpu::BindGroupLayoutEntry {
151                    binding: 2,
152                    visibility: wgpu::ShaderStages::COMPUTE,
153                    ty: storage(true),
154                    count: None,
155                },
156                wgpu::BindGroupLayoutEntry {
157                    binding: 3,
158                    visibility: wgpu::ShaderStages::COMPUTE,
159                    ty: storage(true),
160                    count: None,
161                },
162                wgpu::BindGroupLayoutEntry {
163                    binding: 4,
164                    visibility: wgpu::ShaderStages::COMPUTE,
165                    ty: storage(false),
166                    count: None,
167                },
168            ],
169        })
170    }
171
172    fn run_pipeline(
173        &self,
174        pipeline: &wgpu::ComputePipeline,
175        a: &RnsBatch,
176        b: &RnsBatch,
177    ) -> RnsBatch {
178        let k = a.basis.len();
179        let b_size = a.batch_size;
180        let n_elems = b_size * k;
181        let byte_len = (n_elems * std::mem::size_of::<u32>()) as u64;
182
183        let params = Params {
184            batch_size: b_size as u32,
185            n_channels: k as u32,
186            _pad: [0, 0],
187        };
188        let params_buf = self
189            .device
190            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
191                label: Some("params"),
192                contents: bytemuck::bytes_of(&params),
193                usage: wgpu::BufferUsages::UNIFORM,
194            });
195
196        let moduli_buf = self
197            .device
198            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
199                label: Some("moduli"),
200                contents: bytemuck::cast_slice(a.basis.moduli()),
201                usage: wgpu::BufferUsages::STORAGE,
202            });
203
204        let a_buf = self
205            .device
206            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
207                label: Some("a"),
208                contents: &a.as_u32_bytes(),
209                usage: wgpu::BufferUsages::STORAGE,
210            });
211        let b_buf = self
212            .device
213            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
214                label: Some("b"),
215                contents: &b.as_u32_bytes(),
216                usage: wgpu::BufferUsages::STORAGE,
217            });
218
219        let out_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
220            label: Some("out"),
221            size: byte_len,
222            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
223            mapped_at_creation: false,
224        });
225        let staging_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
226            label: Some("staging"),
227            size: byte_len,
228            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
229            mapped_at_creation: false,
230        });
231
232        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
233            label: Some("adele-ring-bg"),
234            layout: &self.bind_group_layout,
235            entries: &[
236                wgpu::BindGroupEntry {
237                    binding: 0,
238                    resource: params_buf.as_entire_binding(),
239                },
240                wgpu::BindGroupEntry {
241                    binding: 1,
242                    resource: moduli_buf.as_entire_binding(),
243                },
244                wgpu::BindGroupEntry {
245                    binding: 2,
246                    resource: a_buf.as_entire_binding(),
247                },
248                wgpu::BindGroupEntry {
249                    binding: 3,
250                    resource: b_buf.as_entire_binding(),
251                },
252                wgpu::BindGroupEntry {
253                    binding: 4,
254                    resource: out_buf.as_entire_binding(),
255                },
256            ],
257        });
258
259        let mut encoder = self
260            .device
261            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
262        {
263            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
264                label: Some("rns-pass"),
265                timestamp_writes: None,
266            });
267            pass.set_pipeline(pipeline);
268            pass.set_bind_group(0, &bind_group, &[]);
269            pass.dispatch_workgroups(
270                (b_size as u32).div_ceil(16),
271                (k as u32).div_ceil(16),
272                1,
273            );
274        }
275        encoder.copy_buffer_to_buffer(&out_buf, 0, &staging_buf, 0, byte_len);
276        self.queue.submit([encoder.finish()]);
277
278        let slice = staging_buf.slice(..);
279        let (tx, rx) = std::sync::mpsc::channel();
280        slice.map_async(wgpu::MapMode::Read, move |res| {
281            let _ = tx.send(res);
282        });
283        self.device.poll(wgpu::Maintain::Wait);
284        rx.recv()
285            .expect("map_async channel closed")
286            .expect("buffer map failed");
287
288        let data = slice.get_mapped_range();
289        let values: &[u32] = bytemuck::cast_slice(&data);
290        let result = RnsBatch::from_u32(values, b_size, a.basis.clone());
291        drop(data);
292        staging_buf.unmap();
293        result
294    }
295}
296
297impl ArithmeticBackend for GpuBackend {
298    fn batch_add(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
299        self.run_pipeline(&self.add_pipeline, a, b)
300    }
301
302    fn batch_sub(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
303        self.run_pipeline(&self.sub_pipeline, a, b)
304    }
305
306    fn batch_mul(&self, a: &RnsBatch, b: &RnsBatch) -> RnsBatch {
307        self.run_pipeline(&self.mul_pipeline, a, b)
308    }
309
310    fn batch_crt(&self, batch: &RnsBatch) -> Vec<BigInt> {
311        // CRT (Garner) is sequential; do it on the CPU regardless of backend.
312        let k = batch.basis.len();
313        let moduli = batch.basis.moduli();
314        (0..batch.batch_size)
315            .map(|b| crt_balanced(&batch.data[b * k..(b + 1) * k], moduli))
316            .collect()
317    }
318
319    fn name(&self) -> &'static str {
320        "gpu-wgpu"
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use super::*;
327    use crate::basis::Basis;
328    use crate::rns::RnsInt;
329
330    #[test]
331    fn gpu_matches_cpu_when_available() {
332        let gpu = match GpuBackend::try_init() {
333            Ok(g) => g,
334            Err(_) => return, // no GPU on this machine; skip
335        };
336        let b = Basis::standard();
337        let x = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(123, b.clone()); 256]);
338        let y = RnsBatch::from_rns_ints(&vec![RnsInt::from_i64(456, b.clone()); 256]);
339
340        let cpu = crate::cpu::CpuBackend::new();
341        assert_eq!(cpu.batch_add(&x, &y).data, gpu.batch_add(&x, &y).data);
342        assert_eq!(cpu.batch_sub(&x, &y).data, gpu.batch_sub(&x, &y).data);
343        assert_eq!(cpu.batch_mul(&x, &y).data, gpu.batch_mul(&x, &y).data);
344    }
345}
adele_ring/gpu.rs

adele_ring/
gpu.rs
