blade_render/
env_map.rs

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

blade_render/
env_map.rs
