Skip to main content

fidget_wgpu/
lib.rs

1//! Shader generation and WGPU-based image rendering
2#![warn(missing_docs)]
3
4use heck::ToShoutySnakeCase;
5use zerocopy::{FromBytes, Immutable};
6
7pub mod buf;
8pub mod effects;
9pub mod voxel;
10
11/// Re-export the `wgpu` module
12pub use wgpu;
13
14const COMMON_SHADER: &str = include_str!("shaders/common.wgsl");
15
16////////////////////////////////////////////////////////////////////////////////
17
18/// Returns a set of constant definitions for each opcode
19fn opcode_constants() -> String {
20    let mut out = String::new();
21    for (op, i) in fidget_bytecode::iter_ops() {
22        out += &format!("const OP_{}: u32 = {i};\n", op.to_shouty_snake_case());
23    }
24    out
25}
26
27////////////////////////////////////////////////////////////////////////////////
28
29/// Error type for GPU initialization
30#[derive(Debug, thiserror::Error)]
31pub enum InitError {
32    /// Error when requesting an adapter
33    #[error(transparent)]
34    Adapter(#[from] wgpu::RequestAdapterError),
35
36    /// Error when requesting a device
37    #[error(transparent)]
38    Device(#[from] wgpu::RequestDeviceError),
39}
40
41/// Handle to a GPU device
42#[derive(Clone)]
43pub struct Gpu {
44    /// GPU device
45    pub device: wgpu::Device,
46    /// GPU queue
47    pub queue: wgpu::Queue,
48}
49
50impl Gpu {
51    /// Returns a [`Gpu`] object with customized settings
52    ///
53    /// Non-default settings are as follows:
54    /// - We request a [`wgpu::PowerPreference::HighPerformance`] adapter
55    /// - We enable the [`wgpu::Features::TIMESTAMP_QUERY`] feature
56    ///
57    /// This is a helper function for simplicity; more sophisticated systems
58    /// will likely construct the adapter, device, and queue themselves.
59    pub async fn init() -> Result<Gpu, InitError> {
60        let instance = wgpu::Instance::default();
61        let adapter = instance
62            .request_adapter(&wgpu::RequestAdapterOptions {
63                power_preference: wgpu::PowerPreference::HighPerformance,
64                ..wgpu::RequestAdapterOptions::default()
65            })
66            .await?;
67        let (device, queue) = adapter
68            .request_device(&wgpu::DeviceDescriptor {
69                required_features: wgpu::Features::TIMESTAMP_QUERY,
70                ..wgpu::DeviceDescriptor::default()
71            })
72            .await?;
73        Ok(Gpu { device, queue })
74    }
75
76    /// Returns a [`Gpu`] object with default settings
77    ///
78    /// This is useful for CI, where `TIMESTAMP_QUERY` is unsupported
79    #[doc(hidden)]
80    pub async fn init_basic() -> Result<Gpu, InitError> {
81        let instance = wgpu::Instance::default();
82        let adapter = instance
83            .request_adapter(&wgpu::RequestAdapterOptions::default())
84            .await?;
85        let (device, queue) = adapter
86            .request_device(&wgpu::DeviceDescriptor::default())
87            .await?;
88        Ok(Gpu { device, queue })
89    }
90
91    /// Returns a readable buffer for the given image buffer
92    pub fn read_buffer_for<T: buf::BufferTag>(
93        &self,
94        buf: &buf::ImageBuffer<T>,
95    ) -> buf::ImageReadBuffer<T> {
96        buf::ImageBuffer::new(
97            &self.device,
98            format!("{} (read)", buf.name()),
99            buf.size(),
100        )
101        .expect("buf.size should always be a valid size for ImageBuffer::new")
102    }
103
104    /// Maps a readable image buffer, returning a mapped image
105    pub fn map<'a, T: buf::BufferTag>(
106        &self,
107        buf: &'a mut buf::ImageReadBuffer<T>,
108    ) -> buf::MappedImage<'a, T> {
109        buf::MappedImage::map(&self.device, buf)
110    }
111
112    /// Debug function to read from a buffer to a `Vec<T>`
113    pub fn read_vec<T: FromBytes + Immutable + Clone + Copy>(
114        &self,
115        buf: &wgpu::Buffer,
116    ) -> Vec<T> {
117        let scratch = self.device.create_buffer(&wgpu::BufferDescriptor {
118            label: None,
119            size: buf.size(),
120            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
121            mapped_at_creation: false,
122        });
123        let mut encoder = self.device.create_command_encoder(
124            &wgpu::CommandEncoderDescriptor {
125                label: Some("read_buffer"),
126            },
127        );
128        encoder.copy_buffer_to_buffer(buf, 0, &scratch, 0, buf.size());
129        self.queue.submit(Some(encoder.finish()));
130
131        let buffer_slice = scratch.slice(..);
132        buffer_slice.map_async(wgpu::MapMode::Read, |_| {});
133        self.device
134            .poll(wgpu::PollType::wait_indefinitely())
135            .unwrap();
136
137        let result = <[T]>::ref_from_bytes(&buffer_slice.get_mapped_range())
138            .unwrap()
139            .to_vec();
140        scratch.unmap();
141        result
142    }
143}
144
145////////////////////////////////////////////////////////////////////////////////
146
147/// Helper function for use in unit tests
148#[cfg(test)]
149fn compile_shader(src: &str, desc: &str) {
150    let mut v = naga::valid::Validator::new(
151        naga::valid::ValidationFlags::all(),
152        naga::valid::Capabilities::all(),
153    );
154
155    // This isn't the best formatting, but it will at least include the
156    // relevant text.
157    let m = naga::front::wgsl::parse_str(src).unwrap_or_else(|e| {
158        if let Some(i) = e.location(src) {
159            let pos = i.offset as usize..(i.offset + i.length) as usize;
160            panic!(
161                "shader compilation failed\n{src}\n{}",
162                e.emit_to_string_with_path(&src[pos], desc)
163            );
164        } else {
165            panic!(
166                "shader compilation failed\n{src}\n{}",
167                e.emit_to_string(desc)
168            );
169        }
170    });
171    if let Err(e) = v.validate(&m) {
172        let (pos, desc) = e.spans().next().unwrap();
173        panic!(
174            "shader compilation failed\n{src}\n{}",
175            e.emit_to_string_with_path(&src[pos.to_range().unwrap()], desc)
176        );
177    }
178}
179
180#[cfg(test)]
181mod test {
182    use super::*;
183    use fidget_core::{context::Tree, vm::VmShape};
184    use fidget_raster::voxel::RenderSize;
185
186    #[test]
187    fn render_and_merge() {
188        // We only run in CI if we're on MacOS (because other runners don't have
189        // GPUs and will fail to build the context).
190        #[cfg(not(target_os = "macos"))]
191        if std::env::var("CI").is_ok() {
192            return;
193        }
194
195        let gpu = pollster::block_on(Gpu::init_basic()).unwrap();
196        let voxel_ctx = voxel::Context::new(&gpu);
197        let effects_ctx = effects::Context::new(&gpu);
198
199        let size = 128;
200        let image_size = RenderSize::from(size);
201        let mut buf = voxel_ctx.buffers(image_size).unwrap();
202        let mut merge_buf = effects_ctx.merge_buffers(size.into()).unwrap();
203        let mut shade_buf = effects_ctx.shade_buffers(size.into()).unwrap();
204        let mut shade_out = gpu.read_buffer_for(shade_buf.output());
205
206        let (x, y, z) = Tree::axes();
207        let x_ = x.clone() - 0.2;
208        let sphere1 = (x_.square() + y.square() + z.square()).sqrt()
209            - Tree::constant(0.5);
210        let x_ = x + 0.2;
211        let sphere2 = (x_.square() + y.square() + z.square()).sqrt()
212            - Tree::constant(0.5);
213        let spheres = sphere1.min(sphere2);
214        let shape = voxel_ctx.shape(&VmShape::from(spheres)).unwrap();
215
216        voxel_ctx
217            .submit(
218                &shape,
219                &mut buf,
220                None,
221                &voxel::RenderConfig {
222                    world_to_model: nalgebra::Matrix4::identity(),
223                },
224            )
225            .unwrap();
226        effects_ctx
227            .submit_merge(&[buf.image_storage_buffer()], true, &mut merge_buf)
228            .unwrap();
229        let mut ssao_buf = effects_ctx.ssao_buffers(size.into()).unwrap();
230        effects_ctx.submit_ssao(&merge_buf, &mut ssao_buf).unwrap();
231        effects_ctx
232            .submit_shade(
233                &merge_buf,
234                Some(&ssao_buf),
235                &mut shade_buf,
236                Some(&mut shade_out),
237            )
238            .unwrap();
239        let img = gpu.map(&mut shade_out);
240        let (_out, img_size) = img.image().take();
241        assert_eq!(img_size.width(), size);
242        assert_eq!(img_size.height(), size);
243    }
244}