handwired_rs_gpu/
lib.rs

1//! Run handwired circuits as WebGPU compute shaders.
2//!
3//! A program becomes WGSL compute shaders that evaluate one record per
4//! invocation, so a batch of records, e.g. 100,000 documents to score, runs in
5//! parallel. A long program is split into stages that run one after another
6//! (see [`wgsl`]). The shaders reproduce the CPU reference evaluator's values
7//! bit for bit, within the limits in [`wgsl`]'s documentation.
8//!
9//! ```no_run
10//! # use handwired_rs::handwired;
11//! # #[handwired] fn pick(a: u8) -> u8 { if a > 4 { 10 } else { 12 } }
12//! let gpu = handwired_rs_gpu::Gpu::new().unwrap();
13//! let batch = handwired_rs_gpu::GpuBatch::new(&gpu, &pick_circuit()).unwrap();
14//! let results = batch.run(&gpu, &[(1,), (5,), (200,)]).unwrap();
15//! assert_eq!(results, vec![12, 10, 10]);
16//! ```
17
18pub mod wgsl;
19
20use std::fmt;
21use std::marker::PhantomData;
22
23use handwired_rs::dsl::{Compiled, HwType, Params};
24use handwired_rs::{Circuit, Range, TensorProgram, Wire};
25use wgpu::util::DeviceExt;
26
27#[derive(Debug)]
28pub enum GpuError {
29    NoAdapter(String),
30    Device(String),
31    Shader(String),
32    Input(String),
33    Map(String),
34}
35
36impl fmt::Display for GpuError {
37    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
38        match self {
39            GpuError::NoAdapter(message) => write!(formatter, "no GPU adapter: {message}"),
40            GpuError::Device(message) => write!(formatter, "GPU device: {message}"),
41            GpuError::Shader(message) => write!(formatter, "shader: {message}"),
42            GpuError::Input(message) => write!(formatter, "input: {message}"),
43            GpuError::Map(message) => write!(formatter, "reading results: {message}"),
44        }
45    }
46}
47
48impl std::error::Error for GpuError {}
49
50/// A GPU device and its queue.
51pub struct Gpu {
52    pub device: wgpu::Device,
53    pub queue: wgpu::Queue,
54    pub info: wgpu::AdapterInfo,
55}
56
57impl Gpu {
58    /// The default adapter, preferring a discrete or integrated GPU. Set
59    /// `WGPU_ADAPTER_NAME` to pick one by name, e.g. `llvmpipe` for Mesa's
60    /// software Vulkan.
61    pub fn new() -> Result<Gpu, GpuError> {
62        pollster::block_on(Gpu::create())
63    }
64
65    async fn create() -> Result<Gpu, GpuError> {
66        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::from_env_or_default());
67        let adapter = match std::env::var("WGPU_ADAPTER_NAME") {
68            Ok(name) => {
69                let name = name.to_lowercase();
70                instance
71                    .enumerate_adapters(wgpu::Backends::all())
72                    .into_iter()
73                    .find(|adapter| adapter.get_info().name.to_lowercase().contains(&name))
74                    .ok_or_else(|| GpuError::NoAdapter(format!("no adapter named like {name:?}")))?
75            }
76            Err(_) => instance
77                .request_adapter(&wgpu::RequestAdapterOptions {
78                    power_preference: wgpu::PowerPreference::HighPerformance,
79                    force_fallback_adapter: false,
80                    compatible_surface: None,
81                })
82                .await
83                .map_err(|error| GpuError::NoAdapter(error.to_string()))?,
84        };
85        let info = adapter.get_info();
86        if info.device_type == wgpu::DeviceType::Cpu {
87            eprintln!("warning: using the software adapter {:?}; set WGPU_ADAPTER_NAME to choose another", info.name);
88        }
89        let (device, queue) = adapter
90            .request_device(&wgpu::DeviceDescriptor {
91                label: Some("handwired"),
92                required_features: wgpu::Features::empty(),
93                required_limits: adapter.limits(),
94                memory_hints: wgpu::MemoryHints::Performance,
95                trace: wgpu::Trace::Off,
96            })
97            .await
98            .map_err(|error| GpuError::Device(error.to_string()))?;
99        Ok(Gpu { device, queue, info })
100    }
101}
102
103/// A compiled program on the GPU: a pipeline per shader stage and the constant tables.
104pub struct GpuProgram {
105    pipelines: Vec<wgpu::ComputePipeline>,
106    bind_group_layout: wgpu::BindGroupLayout,
107    tables: wgpu::Buffer,
108    pub source_size: usize,
109    pub output_count: usize,
110    /// Operations the shaders evaluate per record.
111    pub operation_count: usize,
112    /// 32-bit words per record that stages pass to later stages.
113    pub carried: usize,
114    /// The WGSL of each stage, run in order.
115    pub stages: Vec<String>,
116    /// At most this many records per dispatch; `None` fits each dispatch's buffers in the device's storage binding
117    /// limit.
118    pub dispatch_records: Option<usize>,
119}
120
121impl GpuProgram {
122    /// Compile the operations that `roots` need into shaders.
123    pub fn new(gpu: &Gpu, circuit: &Circuit, roots: &[Wire]) -> Result<GpuProgram, GpuError> {
124        let program = circuit.compile(Some(roots));
125        GpuProgram::from_program(gpu, &program)
126    }
127
128    pub fn from_program(gpu: &Gpu, program: &TensorProgram) -> Result<GpuProgram, GpuError> {
129        GpuProgram::from_shader(gpu, wgsl::emit(program))
130    }
131
132    /// Like [`GpuProgram::from_program`], starting a new stage after `statements` statements whatever the
133    /// program's length (`usize::MAX` for one shader).
134    pub fn from_program_staged(gpu: &Gpu, program: &TensorProgram, statements: usize) -> Result<GpuProgram, GpuError> {
135        GpuProgram::from_shader(gpu, wgsl::emit_staged(program, statements))
136    }
137
138    fn from_shader(gpu: &Gpu, shader: wgsl::Shader) -> Result<GpuProgram, GpuError> {
139        let device = &gpu.device;
140        // An explicit layout, since an automatic one drops bindings the shader doesn't read, e.g. empty tables.
141        let entry = |binding: u32, ty: wgpu::BufferBindingType| wgpu::BindGroupLayoutEntry {
142            binding,
143            visibility: wgpu::ShaderStages::COMPUTE,
144            ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None },
145            count: None,
146        };
147        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
148            label: Some("handwired bindings"),
149            entries: &[
150                entry(0, wgpu::BufferBindingType::Uniform),
151                entry(1, wgpu::BufferBindingType::Storage { read_only: true }),
152                entry(2, wgpu::BufferBindingType::Storage { read_only: true }),
153                entry(3, wgpu::BufferBindingType::Storage { read_only: false }),
154                entry(4, wgpu::BufferBindingType::Storage { read_only: false }),
155            ],
156        });
157        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
158            label: Some("handwired program"),
159            bind_group_layouts: &[&bind_group_layout],
160            push_constant_ranges: &[],
161        });
162        // Stages compile on every core: the driver's time per stage is the setup time. wgpu 25's error scopes
163        // belong to the device, so this one collects the errors of every thread.
164        device.push_error_scope(wgpu::ErrorFilter::Validation);
165        let next = std::sync::atomic::AtomicUsize::new(0);
166        let threads = std::thread::available_parallelism().map_or(1, |count| count.get()).min(shader.stages.len());
167        let mut compiled: Vec<(usize, wgpu::ComputePipeline)> = std::thread::scope(|scope| {
168            let workers: Vec<_> = (0..threads)
169                .map(|_| {
170                    scope.spawn(|| {
171                        let mut pipelines = Vec::new();
172                        loop {
173                            let index = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
174                            let Some(code) = shader.stages.get(index) else { return pipelines };
175                            let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
176                                label: Some("handwired stage"),
177                                source: wgpu::ShaderSource::Wgsl(code.as_str().into()),
178                            });
179                            let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
180                                label: Some("handwired stage"),
181                                layout: Some(&pipeline_layout),
182                                module: &module,
183                                entry_point: Some("main"),
184                                compilation_options: Default::default(),
185                                cache: None,
186                            });
187                            pipelines.push((index, pipeline));
188                        }
189                    })
190                })
191                .collect();
192            workers.into_iter().flat_map(|worker| worker.join().unwrap()).collect()
193        });
194        compiled.sort_by_key(|&(index, _)| index);
195        let pipelines = compiled.into_iter().map(|(_, pipeline)| pipeline).collect();
196        if let Some(error) = pollster::block_on(device.pop_error_scope()) {
197            return Err(GpuError::Shader(error.to_string()));
198        }
199        // A storage binding can't be empty.
200        let tables = if shader.tables.is_empty() { vec![0.0f32] } else { shader.tables.clone() };
201        let tables = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
202            label: Some("handwired tables"),
203            contents: bytemuck::cast_slice(&tables),
204            usage: wgpu::BufferUsages::STORAGE,
205        });
206        Ok(GpuProgram {
207            pipelines,
208            bind_group_layout,
209            tables,
210            source_size: shader.source_size,
211            output_count: shader.output_count,
212            operation_count: shader.operation_count,
213            carried: shader.carried,
214            stages: shader.stages,
215            dispatch_records: None,
216        })
217    }
218
219    /// Evaluate a batch: `sources` holds one source table after another, and
220    /// the result holds each record's outputs in the order of the roots.
221    pub fn run(&self, gpu: &Gpu, sources: &[f32]) -> Result<Vec<f32>, GpuError> {
222        if self.source_size == 0 || !sources.len().is_multiple_of(self.source_size) {
223            return Err(GpuError::Input(format!(
224                "{} values isn't a whole number of {}-cell records",
225                sources.len(),
226                self.source_size
227            )));
228        }
229        // Each dispatch's buffers must fit in one storage binding; the carried cells are usually the largest.
230        let limits = gpu.device.limits();
231        let binding = (limits.max_storage_buffer_binding_size as u64).min(limits.max_buffer_size);
232        let cells = self.source_size.max(self.output_count).max(self.carried).max(1) as u64;
233        let fits = (binding / (cells * 4)).max(1) as usize;
234        let per_dispatch = self.dispatch_records.map_or(fits, |records| records.clamp(1, fits));
235        let mut values = Vec::with_capacity(sources.len() / self.source_size * self.output_count);
236        for part in sources.chunks(per_dispatch * self.source_size) {
237            values.extend(self.dispatch(gpu, part)?);
238        }
239        Ok(values)
240    }
241
242    fn dispatch(&self, gpu: &Gpu, sources: &[f32]) -> Result<Vec<f32>, GpuError> {
243        let records = sources.len() / self.source_size;
244        let device = &gpu.device;
245        let workgroups = (records as u32).div_ceil(wgsl::WORKGROUP_SIZE);
246        let limit = device.limits().max_compute_workgroups_per_dimension;
247        let (width, height) =
248            if workgroups <= limit { (workgroups.max(1), 1) } else { (limit, workgroups.div_ceil(limit)) };
249        let params: [u32; 4] = [f32::MIN.to_bits(), records as u32, width * wgsl::WORKGROUP_SIZE, 0];
250        let params = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
251            label: Some("handwired params"),
252            contents: bytemuck::cast_slice(&params),
253            usage: wgpu::BufferUsages::UNIFORM,
254        });
255        let source = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
256            label: Some("handwired sources"),
257            contents: bytemuck::cast_slice(sources),
258            usage: wgpu::BufferUsages::STORAGE,
259        });
260        let output_size = (records * self.output_count.max(1) * 4) as u64;
261        let outputs = device.create_buffer(&wgpu::BufferDescriptor {
262            label: Some("handwired outputs"),
263            size: output_size,
264            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
265            mapped_at_creation: false,
266        });
267        let carry = device.create_buffer(&wgpu::BufferDescriptor {
268            label: Some("handwired carried values"),
269            size: (records * self.carried.max(1) * 4) as u64,
270            usage: wgpu::BufferUsages::STORAGE,
271            mapped_at_creation: false,
272        });
273        let readback = device.create_buffer(&wgpu::BufferDescriptor {
274            label: Some("handwired readback"),
275            size: output_size,
276            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
277            mapped_at_creation: false,
278        });
279        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
280            label: Some("handwired bindings"),
281            layout: &self.bind_group_layout,
282            entries: &[
283                wgpu::BindGroupEntry { binding: 0, resource: params.as_entire_binding() },
284                wgpu::BindGroupEntry { binding: 1, resource: source.as_entire_binding() },
285                wgpu::BindGroupEntry { binding: 2, resource: self.tables.as_entire_binding() },
286                wgpu::BindGroupEntry { binding: 3, resource: outputs.as_entire_binding() },
287                wgpu::BindGroupEntry { binding: 4, resource: carry.as_entire_binding() },
288            ],
289        });
290        let mut encoder =
291            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("handwired run") });
292        {
293            // wgpu orders each dispatch's storage writes before the next dispatch's reads.
294            let mut pass = encoder
295                .begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("handwired"), timestamp_writes: None });
296            pass.set_bind_group(0, &bind_group, &[]);
297            for pipeline in &self.pipelines {
298                pass.set_pipeline(pipeline);
299                pass.dispatch_workgroups(width, height, 1);
300            }
301        }
302        encoder.copy_buffer_to_buffer(&outputs, 0, &readback, 0, output_size);
303        gpu.queue.submit([encoder.finish()]);
304
305        let slice = readback.slice(..);
306        let (sender, receiver) = std::sync::mpsc::channel();
307        slice.map_async(wgpu::MapMode::Read, move |result| {
308            let _ = sender.send(result);
309        });
310        device.poll(wgpu::PollType::Wait).map_err(|error| GpuError::Map(error.to_string()))?;
311        receiver
312            .recv()
313            .map_err(|error| GpuError::Map(error.to_string()))?
314            .map_err(|error| GpuError::Map(error.to_string()))?;
315        let values = bytemuck::cast_slice::<u8, f32>(&slice.get_mapped_range()).to_vec();
316        readback.unmap();
317        Ok(values)
318    }
319}
320
321/// A `#[handwired]` function's circuit on the GPU, run on typed arguments.
322pub struct GpuBatch<Args, R> {
323    pub program: GpuProgram,
324    /// The declared range of each source cell.
325    source_ranges: Vec<Range>,
326    marker: PhantomData<fn(Args) -> R>,
327}
328
329impl<Args: Params, R: HwType> GpuBatch<Args, R> {
330    pub fn new(gpu: &Gpu, compiled: &Compiled<Args, R>) -> Result<GpuBatch<Args, R>, GpuError> {
331        Ok(GpuBatch::with_program(compiled, GpuProgram::from_program(gpu, &compiled.program)?))
332    }
333
334    /// Like [`GpuBatch::new`], starting a new shader stage after `statements` statements (`usize::MAX` for one
335    /// shader).
336    pub fn staged(gpu: &Gpu, compiled: &Compiled<Args, R>, statements: usize) -> Result<GpuBatch<Args, R>, GpuError> {
337        Ok(GpuBatch::with_program(compiled, GpuProgram::from_program_staged(gpu, &compiled.program, statements)?))
338    }
339
340    fn with_program(compiled: &Compiled<Args, R>, program: GpuProgram) -> GpuBatch<Args, R> {
341        let circuit = &compiled.circuit;
342        let source_ranges = (0..circuit.source_size()).map(|cell| circuit.range(circuit.source(cell))).collect();
343        GpuBatch { program, source_ranges, marker: PhantomData }
344    }
345
346    /// The source tables for `records`, checked against the declared ranges.
347    pub fn encode(&self, records: &[Args]) -> Result<Vec<f32>, GpuError> {
348        let mut sources = Vec::with_capacity(records.len() * Args::CELLS);
349        for (record_index, record) in records.iter().enumerate() {
350            record.encode(&mut sources);
351            let start = record_index * Args::CELLS;
352            for (cell, (&value, range)) in sources[start..].iter().zip(&self.source_ranges).enumerate() {
353                if !(value as f64 >= range.min && value as f64 <= range.max) {
354                    return Err(GpuError::Input(format!(
355                        "record {record_index}: source cell {cell} is {value}, outside its declared range {}..{}",
356                        range.min, range.max
357                    )));
358                }
359            }
360        }
361        Ok(sources)
362    }
363
364    pub fn run(&self, gpu: &Gpu, records: &[Args]) -> Result<Vec<R>, GpuError> {
365        let outputs = self.program.run(gpu, &self.encode(records)?)?;
366        Ok(outputs
367            .chunks(self.program.output_count.max(1))
368            .take(records.len())
369            .map(|record| R::decode(&mut record.iter().copied()))
370            .collect())
371    }
372}
handwired_rs_gpu/lib.rs

handwired_rs_gpu/
lib.rs
