wgpu_examples/ray_cube_fragment/
mod.rs

1use crate::utils;
2use bytemuck::{Pod, Zeroable};
3use glam::{Mat4, Quat, Vec3};
4use std::ops::IndexMut;
5use std::{borrow::Cow, iter, mem};
6use wgpu::util::DeviceExt;
7
8// from cube
9#[repr(C)]
10#[derive(Clone, Copy, Pod, Zeroable)]
11struct Vertex {
12    _pos: [f32; 4],
13    _tex_coord: [f32; 2],
14}
15
16fn vertex(pos: [i8; 3], tc: [i8; 2]) -> Vertex {
17    Vertex {
18        _pos: [pos[0] as f32, pos[1] as f32, pos[2] as f32, 1.0],
19        _tex_coord: [tc[0] as f32, tc[1] as f32],
20    }
21}
22
23fn create_vertices() -> (Vec<Vertex>, Vec<u16>) {
24    let vertex_data = [
25        // top (0, 0, 1)
26        vertex([-1, -1, 1], [0, 0]),
27        vertex([1, -1, 1], [1, 0]),
28        vertex([1, 1, 1], [1, 1]),
29        vertex([-1, 1, 1], [0, 1]),
30        // bottom (0, 0, -1)
31        vertex([-1, 1, -1], [1, 0]),
32        vertex([1, 1, -1], [0, 0]),
33        vertex([1, -1, -1], [0, 1]),
34        vertex([-1, -1, -1], [1, 1]),
35        // right (1, 0, 0)
36        vertex([1, -1, -1], [0, 0]),
37        vertex([1, 1, -1], [1, 0]),
38        vertex([1, 1, 1], [1, 1]),
39        vertex([1, -1, 1], [0, 1]),
40        // left (-1, 0, 0)
41        vertex([-1, -1, 1], [1, 0]),
42        vertex([-1, 1, 1], [0, 0]),
43        vertex([-1, 1, -1], [0, 1]),
44        vertex([-1, -1, -1], [1, 1]),
45        // front (0, 1, 0)
46        vertex([1, 1, -1], [1, 0]),
47        vertex([-1, 1, -1], [0, 0]),
48        vertex([-1, 1, 1], [0, 1]),
49        vertex([1, 1, 1], [1, 1]),
50        // back (0, -1, 0)
51        vertex([1, -1, 1], [0, 0]),
52        vertex([-1, -1, 1], [1, 0]),
53        vertex([-1, -1, -1], [1, 1]),
54        vertex([1, -1, -1], [0, 1]),
55    ];
56
57    let index_data: &[u16] = &[
58        0, 1, 2, 2, 3, 0, // top
59        4, 5, 6, 6, 7, 4, // bottom
60        8, 9, 10, 10, 11, 8, // right
61        12, 13, 14, 14, 15, 12, // left
62        16, 17, 18, 18, 19, 16, // front
63        20, 21, 22, 22, 23, 20, // back
64    ];
65
66    (vertex_data.to_vec(), index_data.to_vec())
67}
68
69#[repr(C)]
70#[derive(Clone, Copy, Pod, Zeroable)]
71struct Uniforms {
72    view_inverse: Mat4,
73    proj_inverse: Mat4,
74}
75
76struct Example {
77    uniforms: Uniforms,
78    uniform_buf: wgpu::Buffer,
79    blas: wgpu::Blas,
80    tlas: wgpu::Tlas,
81    pipeline: wgpu::RenderPipeline,
82    bind_group: wgpu::BindGroup,
83    animation_timer: utils::AnimationTimer,
84}
85
86impl crate::framework::Example for Example {
87    fn required_features() -> wgpu::Features {
88        wgpu::Features::EXPERIMENTAL_RAY_QUERY
89    }
90
91    fn required_downlevel_capabilities() -> wgpu::DownlevelCapabilities {
92        wgpu::DownlevelCapabilities {
93            flags: wgpu::DownlevelFlags::COMPUTE_SHADERS,
94            ..Default::default()
95        }
96    }
97
98    fn required_limits() -> wgpu::Limits {
99        wgpu::Limits::default().using_minimum_supported_acceleration_structure_values()
100    }
101
102    fn init(
103        config: &wgpu::SurfaceConfiguration,
104        _adapter: &wgpu::Adapter,
105        device: &wgpu::Device,
106        queue: &wgpu::Queue,
107    ) -> Self {
108        let side_count = 8;
109
110        let uniforms = {
111            let view = Mat4::look_at_rh(Vec3::new(0.0, 0.0, 2.5), Vec3::ZERO, Vec3::Y);
112            let proj = Mat4::perspective_rh(
113                59.0_f32.to_radians(),
114                config.width as f32 / config.height as f32,
115                0.001,
116                1000.0,
117            );
118
119            Uniforms {
120                view_inverse: view.inverse(),
121                proj_inverse: proj.inverse(),
122            }
123        };
124
125        let uniform_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
126            label: Some("Uniform Buffer"),
127            contents: bytemuck::cast_slice(&[uniforms]),
128            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
129        });
130
131        let (vertex_data, index_data) = create_vertices();
132
133        let vertex_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
134            label: Some("Vertex Buffer"),
135            contents: bytemuck::cast_slice(&vertex_data),
136            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::BLAS_INPUT,
137        });
138
139        let index_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
140            label: Some("Index Buffer"),
141            contents: bytemuck::cast_slice(&index_data),
142            usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::BLAS_INPUT,
143        });
144
145        let blas_geo_size_desc = wgpu::BlasTriangleGeometrySizeDescriptor {
146            vertex_format: wgpu::VertexFormat::Float32x3,
147            vertex_count: vertex_data.len() as u32,
148            index_format: Some(wgpu::IndexFormat::Uint16),
149            index_count: Some(index_data.len() as u32),
150            flags: wgpu::AccelerationStructureGeometryFlags::OPAQUE,
151        };
152
153        let blas = device.create_blas(
154            &wgpu::CreateBlasDescriptor {
155                label: None,
156                flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
157                update_mode: wgpu::AccelerationStructureUpdateMode::Build,
158            },
159            wgpu::BlasGeometrySizeDescriptors::Triangles {
160                descriptors: vec![blas_geo_size_desc.clone()],
161            },
162        );
163
164        let tlas = device.create_tlas(&wgpu::CreateTlasDescriptor {
165            label: None,
166            flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
167            update_mode: wgpu::AccelerationStructureUpdateMode::Build,
168            max_instances: side_count * side_count,
169        });
170
171        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
172            label: None,
173            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("shader.wgsl"))),
174        });
175
176        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
177            label: None,
178            layout: None,
179            vertex: wgpu::VertexState {
180                module: &shader,
181                entry_point: Some("vs_main"),
182                compilation_options: Default::default(),
183                buffers: &[],
184            },
185            fragment: Some(wgpu::FragmentState {
186                module: &shader,
187                entry_point: Some("fs_main"),
188                compilation_options: Default::default(),
189                targets: &[Some(config.format.into())],
190            }),
191            primitive: wgpu::PrimitiveState {
192                topology: wgpu::PrimitiveTopology::TriangleList,
193                ..Default::default()
194            },
195            depth_stencil: None,
196            multisample: wgpu::MultisampleState::default(),
197            multiview_mask: None,
198            cache: None,
199        });
200
201        let bind_group_layout = pipeline.get_bind_group_layout(0);
202
203        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
204            label: None,
205            layout: &bind_group_layout,
206            entries: &[
207                wgpu::BindGroupEntry {
208                    binding: 0,
209                    resource: uniform_buf.as_entire_binding(),
210                },
211                wgpu::BindGroupEntry {
212                    binding: 1,
213                    resource: wgpu::BindingResource::AccelerationStructure(&tlas),
214                },
215            ],
216        });
217
218        let mut encoder =
219            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
220
221        encoder.build_acceleration_structures(
222            iter::once(&wgpu::BlasBuildEntry {
223                blas: &blas,
224                geometry: wgpu::BlasGeometries::TriangleGeometries(vec![
225                    wgpu::BlasTriangleGeometry {
226                        size: &blas_geo_size_desc,
227                        vertex_buffer: &vertex_buf,
228                        first_vertex: 0,
229                        vertex_stride: mem::size_of::<Vertex>() as u64,
230                        index_buffer: Some(&index_buf),
231                        first_index: Some(0),
232                        transform_buffer: None,
233                        transform_buffer_offset: None,
234                    },
235                ]),
236            }),
237            // iter::empty(),
238            iter::once(&tlas),
239        );
240
241        queue.submit(Some(encoder.finish()));
242
243        Example {
244            uniforms,
245            uniform_buf,
246            blas,
247            tlas,
248            pipeline,
249            bind_group,
250            animation_timer: utils::AnimationTimer::default(),
251        }
252    }
253
254    fn update(&mut self, _event: winit::event::WindowEvent) {}
255
256    fn resize(
257        &mut self,
258        config: &wgpu::SurfaceConfiguration,
259        _device: &wgpu::Device,
260        queue: &wgpu::Queue,
261    ) {
262        let proj = Mat4::perspective_rh(
263            59.0_f32.to_radians(),
264            config.width as f32 / config.height as f32,
265            0.001,
266            1000.0,
267        );
268
269        self.uniforms.proj_inverse = proj.inverse();
270
271        queue.write_buffer(&self.uniform_buf, 0, bytemuck::cast_slice(&[self.uniforms]));
272    }
273
274    fn render(&mut self, view: &wgpu::TextureView, device: &wgpu::Device, queue: &wgpu::Queue) {
275        // scene update
276        {
277            let dist = 12.0;
278
279            let side_count = 8;
280
281            let anim_time = self.animation_timer.time();
282
283            for x in 0..side_count {
284                for y in 0..side_count {
285                    let instance = self.tlas.index_mut((x + y * side_count) as usize);
286
287                    let x = x as f32 / (side_count - 1) as f32;
288                    let y = y as f32 / (side_count - 1) as f32;
289                    let x = x * 2.0 - 1.0;
290                    let y = y * 2.0 - 1.0;
291
292                    let transform = Mat4::from_rotation_translation(
293                        Quat::from_euler(
294                            glam::EulerRot::XYZ,
295                            anim_time * 0.5 * 0.342,
296                            anim_time * 0.5 * 0.254,
297                            anim_time * 0.5 * 0.832,
298                        ),
299                        Vec3 {
300                            x: x * dist,
301                            y: y * dist,
302                            z: -24.0,
303                        },
304                    );
305                    let transform = transform.transpose().to_cols_array()[..12]
306                        .try_into()
307                        .unwrap();
308
309                    *instance = Some(wgpu::TlasInstance::new(&self.blas, transform, 0, 0xff));
310                }
311            }
312        }
313
314        let mut encoder =
315            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
316
317        encoder.build_acceleration_structures(iter::empty(), iter::once(&self.tlas));
318
319        {
320            let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
321                label: None,
322                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
323                    view,
324                    depth_slice: None,
325                    resolve_target: None,
326                    ops: wgpu::Operations {
327                        load: wgpu::LoadOp::Clear(wgpu::Color::GREEN),
328                        store: wgpu::StoreOp::Store,
329                    },
330                })],
331                depth_stencil_attachment: None,
332                timestamp_writes: None,
333                occlusion_query_set: None,
334                multiview_mask: None,
335            });
336
337            rpass.set_pipeline(&self.pipeline);
338            rpass.set_bind_group(0, Some(&self.bind_group), &[]);
339            rpass.draw(0..3, 0..1);
340        }
341
342        queue.submit(Some(encoder.finish()));
343    }
344}
345
346pub fn main() {
347    crate::framework::run::<Example>("ray-cube");
348}
349
350#[cfg(test)]
351#[wgpu_test::gpu_test]
352pub static TEST: crate::framework::ExampleTestParams = crate::framework::ExampleTestParams {
353    name: "ray_cube_fragment",
354    image_path: "/examples/features/src/ray_cube_fragment/screenshot.png",
355    width: 1024,
356    height: 768,
357    optional_features: wgpu::Features::default(),
358    base_test_parameters: wgpu_test::TestParameters::default()
359        // https://github.com/gfx-rs/wgpu/issues/9100
360        .expect_fail(wgpu_test::FailureCase::backend(wgpu::Backends::METAL)),
361    comparisons: &[wgpu_test::ComparisonType::Mean(0.02)],
362    _phantom: std::marker::PhantomData::<Example>,
363};