renderling/cubemap/
cpu.rs

1//! CPU side of the cubemap module.
2use std::sync::Arc;
3
4use glam::{Mat4, UVec2, Vec3, Vec4};
5use image::GenericImageView;
6
7use crate::{
8    stage::{Stage, StageRendering},
9    texture::Texture,
10};
11
12use super::shader::{CubemapDescriptor, CubemapFaceDirection};
13
14pub fn cpu_sample_cubemap(cubemap: &[image::DynamicImage; 6], coord: Vec3) -> Vec4 {
15    let coord = coord.normalize_or(Vec3::X);
16    let (face_index, uv) = CubemapDescriptor::get_face_index_and_uv(coord);
17
18    // Get the selected image
19    let image = &cubemap[face_index];
20
21    // Convert 2D UV to pixel coordinates
22    let (width, height) = image.dimensions();
23    let px = uv.x * (width as f32 - 1.0);
24    let py = uv.y * (height as f32 - 1.0);
25
26    // Sample using the nearest neighbor for simplicity
27    let image::Rgba([r, g, b, a]) = image.get_pixel(px.round() as u32, py.round() as u32);
28
29    // Convert the sampled pixel to Vec4
30    Vec4::new(
31        r as f32 / 255.0,
32        g as f32 / 255.0,
33        b as f32 / 255.0,
34        a as f32 / 255.0,
35    )
36}
37
38/// A cubemap that acts as a render target for an entire scene.
39///
40/// Use this to create and update a skybox with scene geometry.
41pub struct SceneCubemap {
42    pipeline: Arc<wgpu::RenderPipeline>,
43    cubemap_texture: wgpu::Texture,
44    depth_texture: crate::texture::Texture,
45    clear_color: wgpu::Color,
46}
47
48impl SceneCubemap {
49    pub fn new(
50        device: &wgpu::Device,
51        size: UVec2,
52        format: wgpu::TextureFormat,
53        clear_color: Vec4,
54    ) -> Self {
55        let label = Some("scene-to-cubemap");
56        let cubemap_texture = device.create_texture(&wgpu::TextureDescriptor {
57            label,
58            size: wgpu::Extent3d {
59                width: size.x,
60                height: size.y,
61                depth_or_array_layers: 6,
62            },
63            mip_level_count: 1,
64            sample_count: 1,
65            dimension: wgpu::TextureDimension::D2,
66            format,
67            usage: wgpu::TextureUsages::RENDER_ATTACHMENT
68                | wgpu::TextureUsages::TEXTURE_BINDING
69                | wgpu::TextureUsages::COPY_DST
70                | wgpu::TextureUsages::COPY_SRC,
71            view_formats: &[],
72        });
73        let depth_texture = Texture::create_depth_texture(device, size.x, size.y, 1, label);
74        let pipeline = Arc::new(Stage::create_primitive_pipeline(device, format, 1));
75        Self {
76            pipeline,
77            cubemap_texture,
78            depth_texture,
79            clear_color: wgpu::Color {
80                r: clear_color.x as f64,
81                g: clear_color.y as f64,
82                b: clear_color.z as f64,
83                a: clear_color.w as f64,
84            },
85        }
86    }
87
88    pub fn run(&self, stage: &Stage) {
89        let previous_camera_id = stage.used_camera_id();
90
91        // create a new camera for our cube, and use it to render with
92        let camera = stage.geometry.new_camera();
93        stage.use_camera(&camera);
94
95        // By setting this to 90 degrees (PI/2 radians) we make sure the viewing field
96        // is exactly large enough to fill a single face of the cubemap such that all
97        // faces align correctly to each other at the edges.
98        let fovy = std::f32::consts::FRAC_PI_2;
99        let aspect = self.cubemap_texture.width() as f32 / self.cubemap_texture.height() as f32;
100        let projection = Mat4::perspective_lh(fovy, aspect, 1.0, 25.0);
101        // Render each face by rendering the scene from each camera angle into the cubemap
102        for (i, face) in CubemapFaceDirection::FACES.iter().enumerate() {
103            // Update the camera angle, no need to sync as calling `Stage::render` does this
104            // implicitly
105            camera.set_projection_and_view(projection, face.view());
106            let label_s = format!("scene-to-cubemap-{i}");
107            let view = self
108                .cubemap_texture
109                .create_view(&wgpu::TextureViewDescriptor {
110                    label: Some(&label_s),
111                    base_array_layer: i as u32,
112                    array_layer_count: Some(1),
113                    dimension: Some(wgpu::TextureViewDimension::D2),
114                    ..Default::default()
115                });
116            let color_attachment = wgpu::RenderPassColorAttachment {
117                view: &view,
118                resolve_target: None,
119                ops: wgpu::Operations {
120                    load: wgpu::LoadOp::Clear(self.clear_color),
121                    store: wgpu::StoreOp::Store,
122                },
123                depth_slice: None,
124            };
125            let depth_stencil_attachment = wgpu::RenderPassDepthStencilAttachment {
126                view: &self.depth_texture.view,
127                depth_ops: Some(wgpu::Operations {
128                    load: wgpu::LoadOp::Clear(1.0),
129                    store: wgpu::StoreOp::Store,
130                }),
131                stencil_ops: None,
132            };
133            let (_, _) = StageRendering {
134                pipeline: &self.pipeline,
135                stage,
136                color_attachment,
137                depth_stencil_attachment,
138            }
139            .run();
140        }
141
142        stage.use_camera_id(previous_camera_id);
143    }
144}
145
146/// A render pipeline for blitting an equirectangular image as a cubemap.
147pub struct EquirectangularImageToCubemapBlitter(pub wgpu::RenderPipeline);
148
149impl EquirectangularImageToCubemapBlitter {
150    pub fn create_bindgroup_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
151        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
152            label: Some("cubemap-making bindgroup"),
153            entries: &[
154                wgpu::BindGroupLayoutEntry {
155                    binding: 0,
156                    visibility: wgpu::ShaderStages::VERTEX,
157                    ty: wgpu::BindingType::Buffer {
158                        ty: wgpu::BufferBindingType::Storage { read_only: true },
159                        has_dynamic_offset: false,
160                        min_binding_size: None,
161                    },
162                    count: None,
163                },
164                wgpu::BindGroupLayoutEntry {
165                    binding: 1,
166                    visibility: wgpu::ShaderStages::FRAGMENT,
167                    ty: wgpu::BindingType::Texture {
168                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
169                        view_dimension: wgpu::TextureViewDimension::D2,
170                        multisampled: false,
171                    },
172                    count: None,
173                },
174                wgpu::BindGroupLayoutEntry {
175                    binding: 2,
176                    visibility: wgpu::ShaderStages::FRAGMENT,
177                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
178                    count: None,
179                },
180            ],
181        })
182    }
183
184    pub fn create_bindgroup(
185        device: &wgpu::Device,
186        label: Option<&str>,
187        buffer: &wgpu::Buffer,
188        // The texture to sample the environment from
189        texture: &Texture,
190    ) -> wgpu::BindGroup {
191        device.create_bind_group(&wgpu::BindGroupDescriptor {
192            label,
193            layout: &Self::create_bindgroup_layout(device),
194            entries: &[
195                wgpu::BindGroupEntry {
196                    binding: 0,
197                    resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
198                },
199                wgpu::BindGroupEntry {
200                    binding: 1,
201                    resource: wgpu::BindingResource::TextureView(&texture.view),
202                },
203                wgpu::BindGroupEntry {
204                    binding: 2,
205                    resource: wgpu::BindingResource::Sampler(&texture.sampler),
206                },
207            ],
208        })
209    }
210
211    /// Create the rendering pipeline that creates cubemaps from equirectangular
212    /// images.
213    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
214        log::trace!("creating cubemap-making render pipeline with format '{format:?}'");
215        let vertex_linkage = crate::linkage::skybox_cubemap_vertex::linkage(device);
216        let fragment_linkage = crate::linkage::skybox_equirectangular_fragment::linkage(device);
217        let bg_layout = Self::create_bindgroup_layout(device);
218        let pp_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
219            label: Some("cubemap-making pipeline layout"),
220            bind_group_layouts: &[&bg_layout],
221            push_constant_ranges: &[],
222        });
223        EquirectangularImageToCubemapBlitter(device.create_render_pipeline(
224            &wgpu::RenderPipelineDescriptor {
225                label: Some("cubemap-making pipeline"),
226                layout: Some(&pp_layout),
227                vertex: wgpu::VertexState {
228                    module: &vertex_linkage.module,
229                    entry_point: Some(vertex_linkage.entry_point),
230                    buffers: &[],
231                    compilation_options: Default::default(),
232                },
233                primitive: wgpu::PrimitiveState {
234                    topology: wgpu::PrimitiveTopology::TriangleList,
235                    strip_index_format: None,
236                    front_face: wgpu::FrontFace::Ccw,
237                    cull_mode: None,
238                    unclipped_depth: false,
239                    polygon_mode: wgpu::PolygonMode::Fill,
240                    conservative: false,
241                },
242                depth_stencil: None,
243                multisample: wgpu::MultisampleState {
244                    mask: !0,
245                    alpha_to_coverage_enabled: false,
246                    count: 1,
247                },
248                fragment: Some(wgpu::FragmentState {
249                    module: &fragment_linkage.module,
250                    entry_point: Some(fragment_linkage.entry_point),
251                    targets: &[Some(wgpu::ColorTargetState {
252                        format,
253                        blend: Some(wgpu::BlendState::ALPHA_BLENDING),
254                        write_mask: wgpu::ColorWrites::ALL,
255                    })],
256                    compilation_options: Default::default(),
257                }),
258                multiview: None,
259                cache: None,
260            },
261        ))
262    }
263}
264
265#[cfg(test)]
266mod test {
267    use craballoc::slab::SlabAllocator;
268    use glam::Vec4;
269    use image::GenericImageView;
270
271    use crate::{
272        context::Context,
273        geometry::Vertex,
274        math::{UNIT_INDICES, UNIT_POINTS},
275        test::BlockOnFuture,
276        texture::CopiedTextureBuffer,
277    };
278
279    use super::*;
280
281    #[test]
282    fn hand_rolled_cubemap_sampling() {
283        let width = 256;
284        let height = 256;
285        let ctx = Context::headless(width, height).block();
286        let stage = ctx
287            .new_stage()
288            .with_background_color(Vec4::ZERO)
289            .with_lighting(false)
290            .with_msaa_sample_count(4);
291        let projection = crate::camera::perspective(width as f32, height as f32);
292        let view = Mat4::look_at_rh(Vec3::splat(3.0), Vec3::ZERO, Vec3::Y);
293        let _camera = stage
294            .new_camera()
295            .with_projection_and_view(projection, view);
296        // geometry is the "clip cube" where colors are normalized 3d space coords
297        let _rez = stage
298            .new_primitive()
299            .with_vertices(stage.new_vertices(UNIT_POINTS.map(|unit_cube_point| {
300                Vertex::default()
301                    // multiply by 2.0 because the unit cube's AABB bounds are at 0.5, and we want 1.0
302                    .with_position(unit_cube_point * 2.0)
303                    // "normalize" (really "shift") the space coord from [-0.5, 0.5] to [0.0, 1.0]
304                    .with_color((unit_cube_point + 0.5).extend(1.0))
305            })))
306            .with_indices(stage.new_indices(UNIT_INDICES.map(|u| u as u32)));
307
308        let frame = ctx.get_next_frame().unwrap();
309        stage.render(&frame.view());
310        let img = frame.read_image().block().unwrap();
311        img_diff::assert_img_eq("cubemap/hand_rolled_cubemap_sampling/cube.png", img);
312        frame.present();
313
314        let scene_cubemap = SceneCubemap::new(
315            ctx.get_device(),
316            UVec2::new(width, height),
317            wgpu::TextureFormat::Rgba8Unorm,
318            Vec4::ZERO,
319        );
320        scene_cubemap.run(&stage);
321
322        let slab = SlabAllocator::new(&ctx, "cubemap-sampling-test", wgpu::BufferUsages::empty());
323        let uv = slab.new_value(Vec3::ZERO);
324        let buffer = slab.commit();
325        let label = Some("cubemap-sampling-test");
326        let bind_group_layout =
327            ctx.get_device()
328                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
329                    label,
330                    entries: &[
331                        wgpu::BindGroupLayoutEntry {
332                            binding: 0,
333                            visibility: wgpu::ShaderStages::VERTEX,
334                            ty: wgpu::BindingType::Buffer {
335                                ty: wgpu::BufferBindingType::Storage { read_only: true },
336                                has_dynamic_offset: false,
337                                min_binding_size: None,
338                            },
339                            count: None,
340                        },
341                        wgpu::BindGroupLayoutEntry {
342                            binding: 1,
343                            visibility: wgpu::ShaderStages::FRAGMENT,
344                            ty: wgpu::BindingType::Texture {
345                                sample_type: wgpu::TextureSampleType::Float { filterable: true },
346                                view_dimension: wgpu::TextureViewDimension::Cube,
347                                multisampled: false,
348                            },
349                            count: None,
350                        },
351                        wgpu::BindGroupLayoutEntry {
352                            binding: 2,
353                            visibility: wgpu::ShaderStages::FRAGMENT,
354                            ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
355                            count: None,
356                        },
357                    ],
358                });
359        let cubemap_sampling_pipeline_layout =
360            ctx.get_device()
361                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
362                    label,
363                    bind_group_layouts: &[&bind_group_layout],
364                    push_constant_ranges: &[],
365                });
366        let vertex = crate::linkage::cubemap_sampling_test_vertex::linkage(ctx.get_device());
367        let fragment = crate::linkage::cubemap_sampling_test_fragment::linkage(ctx.get_device());
368        let cubemap_sampling_pipeline =
369            ctx.get_device()
370                .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
371                    label,
372                    layout: Some(&cubemap_sampling_pipeline_layout),
373                    vertex: wgpu::VertexState {
374                        module: &vertex.module,
375                        entry_point: Some(vertex.entry_point),
376                        compilation_options: wgpu::PipelineCompilationOptions::default(),
377                        buffers: &[],
378                    },
379                    primitive: wgpu::PrimitiveState {
380                        topology: wgpu::PrimitiveTopology::TriangleList,
381                        strip_index_format: None,
382                        front_face: wgpu::FrontFace::Ccw,
383                        cull_mode: None,
384                        unclipped_depth: false,
385                        polygon_mode: wgpu::PolygonMode::Fill,
386                        conservative: false,
387                    },
388                    depth_stencil: None,
389                    multisample: wgpu::MultisampleState::default(),
390                    fragment: Some(wgpu::FragmentState {
391                        module: &fragment.module,
392                        entry_point: Some(fragment.entry_point),
393                        compilation_options: Default::default(),
394                        targets: &[Some(wgpu::ColorTargetState {
395                            format: wgpu::TextureFormat::Rgba8Unorm,
396                            blend: None,
397                            write_mask: wgpu::ColorWrites::all(),
398                        })],
399                    }),
400                    multiview: None,
401                    cache: None,
402                });
403
404        let cubemap_view =
405            scene_cubemap
406                .cubemap_texture
407                .create_view(&wgpu::TextureViewDescriptor {
408                    label,
409                    dimension: Some(wgpu::TextureViewDimension::Cube),
410                    ..Default::default()
411                });
412        let cubemap_sampler = ctx.get_device().create_sampler(&wgpu::SamplerDescriptor {
413            label,
414            ..Default::default()
415        });
416        let bind_group = ctx
417            .get_device()
418            .create_bind_group(&wgpu::BindGroupDescriptor {
419                label,
420                layout: &bind_group_layout,
421                entries: &[
422                    wgpu::BindGroupEntry {
423                        binding: 0,
424                        resource: buffer.as_entire_binding(),
425                    },
426                    wgpu::BindGroupEntry {
427                        binding: 1,
428                        resource: wgpu::BindingResource::TextureView(&cubemap_view),
429                    },
430                    wgpu::BindGroupEntry {
431                        binding: 2,
432                        resource: wgpu::BindingResource::Sampler(&cubemap_sampler),
433                    },
434                ],
435            });
436        let render_target = ctx.get_device().create_texture(&wgpu::TextureDescriptor {
437            label,
438            size: wgpu::Extent3d {
439                width: 1,
440                height: 1,
441                depth_or_array_layers: 1,
442            },
443            mip_level_count: 1,
444            sample_count: 1,
445            dimension: wgpu::TextureDimension::D2,
446            format: wgpu::TextureFormat::Rgba8Unorm,
447            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
448            view_formats: &[],
449        });
450        let render_target_view = render_target.create_view(&wgpu::TextureViewDescriptor::default());
451
452        let sample = |dir: Vec3| -> Vec4 {
453            uv.set(dir.normalize_or(Vec3::ZERO));
454            slab.commit();
455
456            let mut encoder = ctx
457                .get_device()
458                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label });
459            {
460                let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
461                    label,
462                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
463                        view: &render_target_view,
464                        resolve_target: None,
465                        ops: wgpu::Operations {
466                            load: wgpu::LoadOp::Clear(wgpu::Color {
467                                r: 0.0,
468                                g: 0.0,
469                                b: 0.0,
470                                a: 0.0,
471                            }),
472                            store: wgpu::StoreOp::Store,
473                        },
474                        depth_slice: None,
475                    })],
476                    depth_stencil_attachment: None,
477                    timestamp_writes: None,
478                    occlusion_query_set: None,
479                });
480                render_pass.set_pipeline(&cubemap_sampling_pipeline);
481                render_pass.set_bind_group(0, &bind_group, &[]);
482                render_pass.draw(0..6, 0..1);
483            }
484            let submission_index = ctx.get_queue().submit(Some(encoder.finish()));
485            ctx.get_device()
486                .poll(wgpu::PollType::WaitForSubmissionIndex(submission_index))
487                .unwrap();
488
489            let img = Texture::read(&ctx, &render_target, 1, 1, 4, 1)
490                .into_image::<u8, image::Rgba<u8>>(ctx.get_device())
491                .block()
492                .unwrap();
493            let image::Rgba([r, g, b, a]) = img.get_pixel(0, 0);
494            Vec4::new(
495                r as f32 / 255.0,
496                g as f32 / 255.0,
497                b as f32 / 255.0,
498                a as f32 / 255.0,
499            )
500        };
501
502        fn index_to_face_string(index: usize) -> &'static str {
503            match index {
504                0 => "+X",
505                1 => "-X",
506                2 => "+Y",
507                3 => "-Y",
508                4 => "+Z",
509                5 => "-Z",
510                _ => "?",
511            }
512        }
513
514        let mut cpu_cubemap = vec![];
515        for i in 0..6 {
516            let img = CopiedTextureBuffer::read_from(
517                &ctx,
518                &scene_cubemap.cubemap_texture,
519                width as usize,
520                height as usize,
521                4,
522                1,
523                0,
524                Some(wgpu::Origin3d { x: 0, y: 0, z: i }),
525            )
526            .into_image::<u8, image::Rgba<u8>>(ctx.get_device())
527            .block()
528            .unwrap();
529
530            img_diff::assert_img_eq(
531                &format!(
532                    "cubemap/hand_rolled_cubemap_sampling/face_{}.png",
533                    index_to_face_string(i as usize)
534                ),
535                img.clone(),
536            );
537
538            cpu_cubemap.push(img);
539        }
540        let cpu_cubemap = [
541            cpu_cubemap.remove(0),
542            cpu_cubemap.remove(0),
543            cpu_cubemap.remove(0),
544            cpu_cubemap.remove(0),
545            cpu_cubemap.remove(0),
546            cpu_cubemap.remove(0),
547        ];
548
549        {
550            // assert a few sanity checks on the cpu cubemap
551            println!("x samples sanity");
552            let x_samples_uv = [
553                UVec2::ZERO,
554                UVec2::new(255, 0),
555                UVec2::new(127, 127),
556                UVec2::new(255, 255),
557                UVec2::new(0, 255),
558            ];
559
560            for uv in x_samples_uv {
561                let image::Rgba([r, g, b, a]) = cpu_cubemap[0].get_pixel(uv.x, uv.y);
562                println!("uv: {uv}");
563                println!("rgba: {r} {g} {b} {a}");
564            }
565        }
566
567        let mut uvs = vec![
568            // start with cardinal directions
569            Vec3::X,
570            Vec3::NEG_X,
571            Vec3::Y,
572            Vec3::NEG_Y,
573            Vec3::Z,
574            Vec3::NEG_Z,
575        ];
576
577        // add corners to the uvs to sample
578        for x in [-1.0, 1.0] {
579            for y in [-1.0, 1.0] {
580                for z in [-1.0, 1.0] {
581                    let uv = Vec3::new(x, y, z);
582                    uvs.push(uv);
583                }
584            }
585        }
586
587        // add in some deterministic pseudo-randomn points
588        {
589            let mut prng = crate::math::GpuRng::new(666);
590            let mut rf32 = move || prng.gen_f32(0.0, 1.0);
591            let mut rxvec3 = { || Vec3::new(f32::MAX, rf32(), rf32()).normalize_or(Vec3::X) };
592            // let mut rvec3 = || Vec3::new(rf32(), rf32(), rf32());
593            uvs.extend((0..20).map(|_| rxvec3()));
594        }
595
596        const THRESHOLD: f32 = 0.005;
597        for uv in uvs.into_iter() {
598            let nuv = uv.normalize_or(Vec3::X);
599            let color = sample(uv);
600            let (face_index, uv2d) =
601                CubemapDescriptor::get_face_index_and_uv(uv.normalize_or(Vec3::X));
602            let px = (uv2d.x * (width as f32 - 1.0)).round() as u32;
603            let py = (uv2d.y * (height as f32 - 1.0)).round() as u32;
604            let puv = UVec2::new(px, py);
605            let cpu_color = cpu_sample_cubemap(&cpu_cubemap, uv);
606            let dir_string = index_to_face_string(face_index);
607            println!(
608                "__uv: {uv},\n\
609                 _nuv: {nuv},\n\
610                 _gpu: {color}\n\
611                 _cpu: {cpu_color}\n\
612                 from: {dir_string}({face_index}) {uv2d} {puv}\n"
613            );
614            let cmp = pretty_assertions::Comparison::new(&color, &cpu_color);
615            let distance = color.distance(cpu_color);
616            if distance > THRESHOLD {
617                println!("distance: {distance}");
618                println!("{cmp}");
619                panic!("distance {distance} greater than {THRESHOLD}");
620            }
621        }
622    }
623}
renderling/cubemap/cpu.rs

renderling/cubemap/
cpu.rs
