handwired_rs_gpu/
lib.rs

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

handwired_rs_gpu/
lib.rs
