blade_render/render/
env_map.rs

1use std::num::NonZeroU32;
2
3use crate::DummyResources;
4
5#[repr(C)]
6#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod)]
7struct EnvPreprocParams {
8    target_level: u32,
9}
10
11#[derive(blade_macros::ShaderData)]
12struct EnvPreprocData {
13    source: blade_graphics::TextureView,
14    destination: blade_graphics::TextureView,
15    params: EnvPreprocParams,
16}
17
18pub struct EnvironmentMap {
19    pub main_view: blade_graphics::TextureView,
20    pub size: blade_graphics::Extent,
21    pub weight_texture: blade_graphics::Texture,
22    pub weight_view: blade_graphics::TextureView,
23    pub weight_mips: Vec<blade_graphics::TextureView>,
24    pub prepare_pipeline: blade_graphics::ComputePipeline,
25}
26
27impl EnvironmentMap {
28    pub fn init_pipeline(
29        shader: &blade_graphics::Shader,
30        gpu: &blade_graphics::Context,
31    ) -> Result<blade_graphics::ComputePipeline, &'static str> {
32        let layout = <EnvPreprocData as blade_graphics::ShaderData>::layout();
33        shader.check_struct_size::<EnvPreprocParams>();
34
35        Ok(
36            gpu.create_compute_pipeline(blade_graphics::ComputePipelineDesc {
37                name: "env-prepare",
38                data_layouts: &[&layout],
39                compute: shader.at("downsample"),
40            }),
41        )
42    }
43
44    pub fn with_pipeline(
45        dummy: &DummyResources,
46        prepare_pipeline: blade_graphics::ComputePipeline,
47    ) -> Self {
48        Self {
49            main_view: dummy.white_view,
50            size: blade_graphics::Extent::default(),
51            weight_texture: blade_graphics::Texture::default(),
52            weight_view: dummy.red_view,
53            weight_mips: Vec::new(),
54            prepare_pipeline,
55        }
56    }
57
58    pub fn new(
59        shader: &blade_graphics::Shader,
60        dummy: &DummyResources,
61        gpu: &blade_graphics::Context,
62    ) -> Self {
63        Self::with_pipeline(dummy, Self::init_pipeline(shader, gpu).unwrap())
64    }
65
66    fn weight_size(&self) -> blade_graphics::Extent {
67        // The weight texture has to include all of the edge pixels, starting at mip 1
68        blade_graphics::Extent {
69            width: self.size.width.next_power_of_two() / 2,
70            height: self.size.height.next_power_of_two() / 2,
71            depth: 1,
72        }
73    }
74
75    pub fn destroy(&mut self, gpu: &blade_graphics::Context) {
76        if self.weight_texture != blade_graphics::Texture::default() {
77            gpu.destroy_texture(self.weight_texture);
78            gpu.destroy_texture_view(self.weight_view);
79        }
80        for view in self.weight_mips.drain(..) {
81            gpu.destroy_texture_view(view);
82        }
83        gpu.destroy_compute_pipeline(&mut self.prepare_pipeline);
84    }
85
86    pub fn assign(
87        &mut self,
88        view: blade_graphics::TextureView,
89        extent: blade_graphics::Extent,
90        encoder: &mut blade_graphics::CommandEncoder,
91        gpu: &blade_graphics::Context,
92    ) {
93        if self.main_view == view {
94            return;
95        }
96        self.main_view = view;
97        self.size = extent;
98        self.destroy(gpu);
99
100        let mip_level_count = extent
101            .width
102            .max(extent.height)
103            .next_power_of_two()
104            .trailing_zeros();
105        let weight_extent = self.weight_size();
106        let format = blade_graphics::TextureFormat::Rgba16Float;
107        self.weight_texture = gpu.create_texture(blade_graphics::TextureDesc {
108            name: "env-weight",
109            format,
110            size: weight_extent,
111            dimension: blade_graphics::TextureDimension::D2,
112            array_layer_count: 1,
113            mip_level_count,
114            usage: blade_graphics::TextureUsage::RESOURCE | blade_graphics::TextureUsage::STORAGE,
115            sample_count: 1,
116        });
117        self.weight_view = gpu.create_texture_view(
118            self.weight_texture,
119            blade_graphics::TextureViewDesc {
120                name: "env-weight",
121                format,
122                dimension: blade_graphics::ViewDimension::D2,
123                subresources: &Default::default(),
124            },
125        );
126        for base_mip_level in 0..mip_level_count {
127            let view = gpu.create_texture_view(
128                self.weight_texture,
129                blade_graphics::TextureViewDesc {
130                    name: &format!("env-weight-mip{}", base_mip_level),
131                    format,
132                    dimension: blade_graphics::ViewDimension::D2,
133                    subresources: &blade_graphics::TextureSubresources {
134                        base_mip_level,
135                        mip_level_count: NonZeroU32::new(1),
136                        ..Default::default()
137                    },
138                },
139            );
140            self.weight_mips.push(view);
141        }
142
143        encoder.init_texture(self.weight_texture);
144        for target_level in 0..mip_level_count {
145            let groups = self
146                .prepare_pipeline
147                .get_dispatch_for(weight_extent.at_mip_level(target_level));
148            let mut compute = encoder.compute("pre-process env map");
149            let mut pass = compute.with(&self.prepare_pipeline);
150            pass.bind(
151                0,
152                &EnvPreprocData {
153                    source: if target_level == 0 {
154                        view
155                    } else {
156                        self.weight_mips[target_level as usize - 1]
157                    },
158                    destination: self.weight_mips[target_level as usize],
159                    params: EnvPreprocParams { target_level },
160                },
161            );
162            pass.dispatch(groups);
163        }
164    }
165}
blade_render/render/env_map.rs

blade_render/render/
env_map.rs
