renderling/pbr/ibl/
cpu.rs

1//! CPU side of IBL
2
3use core::sync::atomic::AtomicBool;
4use std::sync::Arc;
5
6use craballoc::{runtime::WgpuRuntime, slab::SlabAllocator, value::Hybrid};
7use crabslab::Id;
8use glam::{Mat4, Vec3};
9
10use crate::{
11    camera::Camera, convolution::shader::VertexPrefilterEnvironmentCubemapIds, skybox::Skybox,
12    texture,
13};
14
15/// Image based lighting resources.
16#[derive(Clone)]
17pub struct Ibl {
18    is_empty: Arc<AtomicBool>,
19    // Cubemap texture of the pre-computed irradiance cubemap
20    pub(crate) irradiance_cubemap: texture::Texture,
21    // Cubemap texture and mip maps of the specular highlights,
22    // where each mip level is a different roughness.
23    pub(crate) prefiltered_environment_cubemap: texture::Texture,
24}
25
26impl Ibl {
27    /// Create a new [`Ibl`] resource.
28    pub fn new(runtime: impl AsRef<WgpuRuntime>, skybox: &Skybox) -> Self {
29        log::trace!("creating new IBL");
30        let runtime = runtime.as_ref();
31        let slab = SlabAllocator::new(runtime, "ibl", wgpu::BufferUsages::VERTEX);
32        let proj = Mat4::perspective_rh(std::f32::consts::FRAC_PI_2, 1.0, 0.1, 10.0);
33        let camera = Camera::new(&slab).with_projection(proj);
34        let roughness = slab.new_value(0.0f32);
35        let prefilter_ids = slab.new_value(VertexPrefilterEnvironmentCubemapIds {
36            camera: camera.id(),
37            roughness: roughness.id(),
38        });
39
40        let buffer = slab.commit();
41        let mut buffer_upkeep = || {
42            let possibly_new_buffer = slab.commit();
43            debug_assert!(!possibly_new_buffer.is_new_this_commit());
44        };
45
46        let views = [
47            Mat4::look_at_rh(
48                Vec3::new(0.0, 0.0, 0.0),
49                Vec3::new(1.0, 0.0, 0.0),
50                Vec3::new(0.0, -1.0, 0.0),
51            ),
52            Mat4::look_at_rh(
53                Vec3::new(0.0, 0.0, 0.0),
54                Vec3::new(-1.0, 0.0, 0.0),
55                Vec3::new(0.0, -1.0, 0.0),
56            ),
57            Mat4::look_at_rh(
58                Vec3::new(0.0, 0.0, 0.0),
59                Vec3::new(0.0, -1.0, 0.0),
60                Vec3::new(0.0, 0.0, -1.0),
61            ),
62            Mat4::look_at_rh(
63                Vec3::new(0.0, 0.0, 0.0),
64                Vec3::new(0.0, 1.0, 0.0),
65                Vec3::new(0.0, 0.0, 1.0),
66            ),
67            Mat4::look_at_rh(
68                Vec3::new(0.0, 0.0, 0.0),
69                Vec3::new(0.0, 0.0, 1.0),
70                Vec3::new(0.0, -1.0, 0.0),
71            ),
72            Mat4::look_at_rh(
73                Vec3::new(0.0, 0.0, 0.0),
74                Vec3::new(0.0, 0.0, -1.0),
75                Vec3::new(0.0, -1.0, 0.0),
76            ),
77        ];
78
79        let environment_cubemap = skybox.environment_cubemap_texture();
80
81        // Convolve the environment map.
82        let irradiance_cubemap = create_irradiance_map(
83            runtime,
84            &buffer,
85            &mut buffer_upkeep,
86            environment_cubemap,
87            &camera,
88            views,
89        );
90
91        // Generate specular IBL pre-filtered environment map.
92        let prefiltered_environment_cubemap = create_prefiltered_environment_map(
93            runtime,
94            &buffer,
95            &mut buffer_upkeep,
96            &camera,
97            &roughness,
98            prefilter_ids.id(),
99            environment_cubemap,
100            views,
101        );
102
103        Self {
104            is_empty: Arc::new(skybox.is_empty().into()),
105            irradiance_cubemap,
106            prefiltered_environment_cubemap,
107        }
108    }
109
110    /// Returns whether this [`Ibl`] is empty.
111    ///
112    /// An [`Ibl`] is empty if it was created from an empty [`Skybox`].
113    pub fn is_empty(&self) -> bool {
114        self.is_empty.load(std::sync::atomic::Ordering::Relaxed)
115    }
116}
117
118fn create_irradiance_map(
119    runtime: impl AsRef<WgpuRuntime>,
120    buffer: &wgpu::Buffer,
121    buffer_upkeep: impl FnMut(),
122    environment_texture: &texture::Texture,
123    camera: &Camera,
124    views: [Mat4; 6],
125) -> texture::Texture {
126    let runtime = runtime.as_ref();
127    let device = &runtime.device;
128    let pipeline = crate::pbr::ibl::DiffuseIrradianceConvolutionRenderPipeline::new(
129        device,
130        wgpu::TextureFormat::Rgba16Float,
131    );
132
133    let bindgroup = crate::pbr::ibl::diffuse_irradiance_convolution_bindgroup(
134        device,
135        Some("irradiance"),
136        buffer,
137        environment_texture,
138    );
139
140    texture::Texture::render_cubemap(
141        runtime,
142        "diffuse-irradiance",
143        &pipeline.0,
144        buffer_upkeep,
145        camera,
146        &bindgroup,
147        views,
148        32,
149        None,
150    )
151}
152
153/// Pipeline for creating a prefiltered environment map from an existing
154/// environment cubemap.
155pub(crate) fn create_prefiltered_environment_pipeline_and_bindgroup(
156    device: &wgpu::Device,
157    buffer: &wgpu::Buffer,
158    environment_texture: &crate::texture::Texture,
159) -> (wgpu::RenderPipeline, wgpu::BindGroup) {
160    let label = Some("prefiltered environment");
161    let bindgroup_layout_desc = wgpu::BindGroupLayoutDescriptor {
162        label,
163        entries: &[
164            wgpu::BindGroupLayoutEntry {
165                binding: 0,
166                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
167                ty: wgpu::BindingType::Buffer {
168                    ty: wgpu::BufferBindingType::Storage { read_only: true },
169                    has_dynamic_offset: false,
170                    min_binding_size: None,
171                },
172                count: None,
173            },
174            wgpu::BindGroupLayoutEntry {
175                binding: 1,
176                visibility: wgpu::ShaderStages::FRAGMENT,
177                ty: wgpu::BindingType::Texture {
178                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
179                    view_dimension: wgpu::TextureViewDimension::Cube,
180                    multisampled: false,
181                },
182                count: None,
183            },
184            wgpu::BindGroupLayoutEntry {
185                binding: 2,
186                visibility: wgpu::ShaderStages::FRAGMENT,
187                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
188                count: None,
189            },
190        ],
191    };
192    let bg_layout = device.create_bind_group_layout(&bindgroup_layout_desc);
193    let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor {
194        label,
195        layout: &bg_layout,
196        entries: &[
197            wgpu::BindGroupEntry {
198                binding: 0,
199                resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
200            },
201            wgpu::BindGroupEntry {
202                binding: 1,
203                resource: wgpu::BindingResource::TextureView(&environment_texture.view),
204            },
205            wgpu::BindGroupEntry {
206                binding: 2,
207                resource: wgpu::BindingResource::Sampler(&environment_texture.sampler),
208            },
209        ],
210    });
211    let pp_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
212        label,
213        bind_group_layouts: &[&bg_layout],
214        push_constant_ranges: &[],
215    });
216    let vertex_linkage = crate::linkage::prefilter_environment_cubemap_vertex::linkage(device);
217    let fragment_linkage = crate::linkage::prefilter_environment_cubemap_fragment::linkage(device);
218    let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
219        label: Some("prefiltered environment"),
220        layout: Some(&pp_layout),
221        vertex: wgpu::VertexState {
222            module: &vertex_linkage.module,
223            entry_point: Some(vertex_linkage.entry_point),
224            buffers: &[],
225            compilation_options: Default::default(),
226        },
227        primitive: wgpu::PrimitiveState {
228            topology: wgpu::PrimitiveTopology::TriangleList,
229            strip_index_format: None,
230            front_face: wgpu::FrontFace::Ccw,
231            cull_mode: None,
232            unclipped_depth: false,
233            polygon_mode: wgpu::PolygonMode::Fill,
234            conservative: false,
235        },
236        depth_stencil: None,
237        multisample: wgpu::MultisampleState {
238            mask: !0,
239            alpha_to_coverage_enabled: false,
240            count: 1,
241        },
242        fragment: Some(wgpu::FragmentState {
243            module: &fragment_linkage.module,
244            entry_point: Some(fragment_linkage.entry_point),
245            targets: &[Some(wgpu::ColorTargetState {
246                format: wgpu::TextureFormat::Rgba16Float,
247                blend: Some(wgpu::BlendState {
248                    color: wgpu::BlendComponent::REPLACE,
249                    alpha: wgpu::BlendComponent::REPLACE,
250                }),
251                write_mask: wgpu::ColorWrites::ALL,
252            })],
253            compilation_options: Default::default(),
254        }),
255        multiview: None,
256        cache: None,
257    });
258    (pipeline, bindgroup)
259}
260
261#[allow(clippy::too_many_arguments)]
262fn create_prefiltered_environment_map(
263    runtime: impl AsRef<WgpuRuntime>,
264    buffer: &wgpu::Buffer,
265    mut buffer_upkeep: impl FnMut(),
266    camera: &Camera,
267    roughness: &Hybrid<f32>,
268    prefilter_id: Id<VertexPrefilterEnvironmentCubemapIds>,
269    environment_texture: &texture::Texture,
270    views: [Mat4; 6],
271) -> texture::Texture {
272    let (pipeline, bindgroup) =
273        crate::pbr::ibl::create_prefiltered_environment_pipeline_and_bindgroup(
274            &runtime.as_ref().device,
275            buffer,
276            environment_texture,
277        );
278    let mut cubemap_faces = Vec::new();
279
280    for (i, view) in views.iter().enumerate() {
281        for mip_level in 0..5 {
282            let mip_width: u32 = 128 >> mip_level;
283            let mip_height: u32 = 128 >> mip_level;
284
285            let mut encoder =
286                runtime
287                    .as_ref()
288                    .device
289                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
290                        label: Some("specular convolution"),
291                    });
292
293            let cubemap_face = texture::Texture::new_with(
294                runtime.as_ref(),
295                Some(&format!("cubemap{i}{mip_level}prefiltered_environment")),
296                Some(wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC),
297                None,
298                wgpu::TextureFormat::Rgba16Float,
299                4,
300                2,
301                mip_width,
302                mip_height,
303                1,
304                &[],
305            );
306
307            // update the roughness for these mips
308            roughness.set(mip_level as f32 / 4.0);
309            // update the view to point at one of the cube faces
310            camera.set_view(*view);
311            buffer_upkeep();
312            {
313                let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
314                    label: Some(&format!("cubemap{i}")),
315                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
316                        view: &cubemap_face.view,
317                        resolve_target: None,
318                        ops: wgpu::Operations {
319                            load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
320                            store: wgpu::StoreOp::Store,
321                        },
322                        depth_slice: None,
323                    })],
324                    depth_stencil_attachment: None,
325                    ..Default::default()
326                });
327
328                render_pass.set_pipeline(&pipeline);
329                render_pass.set_bind_group(0, Some(&bindgroup), &[]);
330                render_pass.draw(0..36, prefilter_id.inner()..prefilter_id.inner() + 1);
331            }
332
333            runtime.as_ref().queue.submit([encoder.finish()]);
334            cubemap_faces.push(cubemap_face);
335        }
336    }
337
338    texture::Texture::new_cubemap_texture(
339        runtime,
340        Some("prefiltered-environment-cubemap"),
341        128,
342        cubemap_faces.as_slice(),
343        wgpu::TextureFormat::Rgba16Float,
344        5,
345    )
346}
347
348pub fn diffuse_irradiance_convolution_bindgroup_layout(
349    device: &wgpu::Device,
350) -> wgpu::BindGroupLayout {
351    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
352        label: Some("convolution bindgroup"),
353        entries: &[
354            wgpu::BindGroupLayoutEntry {
355                binding: 0,
356                visibility: wgpu::ShaderStages::VERTEX,
357                ty: wgpu::BindingType::Buffer {
358                    ty: wgpu::BufferBindingType::Storage { read_only: true },
359                    has_dynamic_offset: false,
360                    min_binding_size: None,
361                },
362                count: None,
363            },
364            wgpu::BindGroupLayoutEntry {
365                binding: 1,
366                visibility: wgpu::ShaderStages::FRAGMENT,
367                ty: wgpu::BindingType::Texture {
368                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
369                    view_dimension: wgpu::TextureViewDimension::Cube,
370                    multisampled: false,
371                },
372                count: None,
373            },
374            wgpu::BindGroupLayoutEntry {
375                binding: 2,
376                visibility: wgpu::ShaderStages::FRAGMENT,
377                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
378                count: None,
379            },
380        ],
381    })
382}
383
384pub fn diffuse_irradiance_convolution_bindgroup(
385    device: &wgpu::Device,
386    label: Option<&str>,
387    buffer: &wgpu::Buffer,
388    // The texture to sample the environment from
389    texture: &crate::texture::Texture,
390) -> wgpu::BindGroup {
391    device.create_bind_group(&wgpu::BindGroupDescriptor {
392        label,
393        layout: &diffuse_irradiance_convolution_bindgroup_layout(device),
394        entries: &[
395            wgpu::BindGroupEntry {
396                binding: 0,
397                resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
398            },
399            wgpu::BindGroupEntry {
400                binding: 1,
401                resource: wgpu::BindingResource::TextureView(&texture.view),
402            },
403            wgpu::BindGroupEntry {
404                binding: 2,
405                resource: wgpu::BindingResource::Sampler(&texture.sampler),
406            },
407        ],
408    })
409}
410
411pub struct DiffuseIrradianceConvolutionRenderPipeline(pub wgpu::RenderPipeline);
412
413impl DiffuseIrradianceConvolutionRenderPipeline {
414    /// Create the rendering pipeline that performs a convolution.
415    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
416        let vertex_linkage = crate::linkage::skybox_cubemap_vertex::linkage(device);
417        let fragment_linkage = crate::linkage::di_convolution_fragment::linkage(device);
418        let bg_layout = diffuse_irradiance_convolution_bindgroup_layout(device);
419        let pp_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
420            label: Some("convolution pipeline layout"),
421            bind_group_layouts: &[&bg_layout],
422            push_constant_ranges: &[],
423        });
424
425        DiffuseIrradianceConvolutionRenderPipeline(device.create_render_pipeline(
426            &wgpu::RenderPipelineDescriptor {
427                label: Some("convolution pipeline"),
428                layout: Some(&pp_layout),
429                vertex: wgpu::VertexState {
430                    module: &vertex_linkage.module,
431                    entry_point: Some(vertex_linkage.entry_point),
432                    buffers: &[],
433                    compilation_options: Default::default(),
434                },
435                primitive: wgpu::PrimitiveState {
436                    topology: wgpu::PrimitiveTopology::TriangleList,
437                    strip_index_format: None,
438                    front_face: wgpu::FrontFace::Ccw,
439                    cull_mode: None,
440                    unclipped_depth: false,
441                    polygon_mode: wgpu::PolygonMode::Fill,
442                    conservative: false,
443                },
444                depth_stencil: None,
445                multisample: wgpu::MultisampleState {
446                    mask: !0,
447                    alpha_to_coverage_enabled: false,
448                    count: 1,
449                },
450                fragment: Some(wgpu::FragmentState {
451                    module: &fragment_linkage.module,
452                    entry_point: Some(fragment_linkage.entry_point),
453                    targets: &[Some(wgpu::ColorTargetState {
454                        format,
455                        blend: Some(wgpu::BlendState::ALPHA_BLENDING),
456                        write_mask: wgpu::ColorWrites::ALL,
457                    })],
458                    compilation_options: Default::default(),
459                }),
460                multiview: None,
461                cache: None,
462            },
463        ))
464    }
465}
466
467#[cfg(test)]
468mod test {
469    use glam::{Mat4, Vec3};
470
471    use crate::{
472        context::Context,
473        test::{workspace_dir, BlockOnFuture},
474        texture::CopiedTextureBuffer,
475    };
476
477    #[test]
478    /// Creates an Ibl and reads out its diffuse irradiance and prefiltered
479    /// environment cubemap mips and compares them against known images to
480    /// ensure creation is valid.
481    fn creates_valid_cubemaps() {
482        let ctx = Context::headless(600, 400).block();
483        let proj = crate::camera::perspective(600.0, 400.0);
484        let view = crate::camera::look_at(Vec3::new(0.0, 0.0, 2.0), Vec3::ZERO, Vec3::Y);
485        let stage = ctx.new_stage();
486        let _camera = stage.new_camera().with_projection_and_view(proj, view);
487        let skybox = stage
488            .new_skybox_from_path(workspace_dir().join("img/hdr/resting_place.hdr"))
489            .unwrap();
490        let ibl = stage.new_ibl(&skybox);
491        stage.use_ibl(&ibl);
492        assert_eq!(
493            wgpu::TextureFormat::Rgba16Float,
494            ibl.irradiance_cubemap.texture.format()
495        );
496        assert_eq!(
497            wgpu::TextureFormat::Rgba16Float,
498            ibl.prefiltered_environment_cubemap.texture.format()
499        );
500        for i in 0..6 {
501            // save out the irradiance face
502            let copied_buffer = CopiedTextureBuffer::read_from(
503                &ctx,
504                &ibl.irradiance_cubemap.texture,
505                32,
506                32,
507                4,
508                2,
509                0,
510                Some(wgpu::Origin3d { x: 0, y: 0, z: i }),
511            );
512            let pixels = copied_buffer.pixels(ctx.get_device()).block().unwrap();
513            let pixels = bytemuck::cast_slice::<u8, u16>(pixels.as_slice())
514                .iter()
515                .map(|p| half::f16::from_bits(*p).to_f32())
516                .collect::<Vec<_>>();
517            assert_eq!(32 * 32 * 4, pixels.len());
518            let img: image::Rgba32FImage = image::ImageBuffer::from_vec(32, 32, pixels).unwrap();
519            let img = image::DynamicImage::from(img);
520            let img = img.to_rgba8();
521            img_diff::assert_img_eq(&format!("skybox/irradiance{i}.png"), img);
522            for mip_level in 0..5 {
523                let mip_size = 128u32 >> mip_level;
524                // save out the prefiltered environment faces' mips
525                let copied_buffer = CopiedTextureBuffer::read_from(
526                    &ctx,
527                    &ibl.prefiltered_environment_cubemap.texture,
528                    mip_size as usize,
529                    mip_size as usize,
530                    4,
531                    2,
532                    mip_level,
533                    Some(wgpu::Origin3d { x: 0, y: 0, z: i }),
534                );
535                let pixels = copied_buffer.pixels(ctx.get_device()).block().unwrap();
536                let pixels = bytemuck::cast_slice::<u8, u16>(pixels.as_slice())
537                    .iter()
538                    .map(|p| half::f16::from_bits(*p).to_f32())
539                    .collect::<Vec<_>>();
540                assert_eq!((mip_size * mip_size * 4) as usize, pixels.len());
541                let img: image::Rgba32FImage =
542                    image::ImageBuffer::from_vec(mip_size, mip_size, pixels).unwrap();
543                let img = image::DynamicImage::from(img);
544                let img = img.to_rgba8();
545                img_diff::assert_img_eq(
546                    &format!("skybox/prefiltered_environment_face{i}_mip{mip_level}.png"),
547                    img,
548                );
549            }
550        }
551    }
552
553    #[test]
554    /// Creates a Skybox, Ibl, and uses the Ibl to light a mirror cube.
555    fn mirror_cube_is_lit_by_environment() {
556        let ctx = Context::headless(256, 256).block();
557        let stage = ctx.new_stage();
558
559        let _camera = stage
560            .new_camera()
561            .with_default_perspective(256.0, 256.0)
562            .with_view(Mat4::look_at_rh(Vec3::ONE * 1.5, Vec3::ZERO, Vec3::Y));
563        let _model = stage.new_primitive().with_material(
564            stage
565                .new_material()
566                .with_metallic_factor(0.9)
567                .with_roughness_factor(0.1),
568        );
569
570        let skybox = stage
571            .new_skybox_from_path(workspace_dir().join("img/hdr/helipad.hdr"))
572            .unwrap();
573        stage.use_skybox(&skybox);
574
575        // Render once here because we found a bug where rendering before setting
576        // ibl would cause the primitive bindgroup to *not* be invalidated when
577        // ibl was set.
578        //
579        // This essentially just ensures that `Stage::use_ibl` is invalidating the
580        // primitive bindgroup.
581        let frame = ctx.get_next_frame().unwrap();
582        stage.render(&frame.view());
583        frame.present();
584
585        let ibl = stage.new_ibl(&skybox);
586        stage.use_ibl(&ibl);
587        stage.remove_skybox();
588
589        let frame = ctx.get_next_frame().unwrap();
590        stage.render(&frame.view());
591        let img = frame.read_image().block().unwrap();
592        img_diff::save("pbr/ibl/mirror_cube_is_lit_by_environment.png", img);
593        frame.present();
594    }
595}
renderling/pbr/ibl/cpu.rs

renderling/pbr/ibl/
cpu.rs
