1use std::collections::HashMap;
2
3use crate::{assets::singleton_asset::LazyResource, wgpu::backend::WGPUBackend};
4
5const MIP_SHADER: &str = r#"
6struct VOut {
7 @builtin(position) pos: vec4<f32>,
8 @location(0) uv: vec2<f32>,
9};
10
11@vertex
12fn vs_main(@builtin(vertex_index) idx: u32) -> VOut {
13 var out: VOut;
14 let x = f32((idx << 1u) & 2u);
15 let y = f32(idx & 2u);
16 out.pos = vec4<f32>(x * 2.0 - 1.0, 1.0 - y * 2.0, 0.0, 1.0);
17 out.uv = vec2<f32>(x, y);
18 return out;
19}
20
21@group(0) @binding(0) var src_texture: texture_2d<f32>;
22@group(0) @binding(1) var src_sampler: sampler;
23
24@fragment
25fn fs_main(in: VOut) -> @location(0) vec4<f32> {
26 return textureSample(src_texture, src_sampler, in.uv);
27}
28"#;
29
30const SUPPORTED_FORMATS: &[wgpu::TextureFormat] = &[
33 wgpu::TextureFormat::Rgba8Unorm,
34 wgpu::TextureFormat::Rgba8UnormSrgb,
35 wgpu::TextureFormat::Rgba16Float,
36 wgpu::TextureFormat::Rgba32Float,
37];
38
39struct MipPipeline {
40 pipeline: wgpu::RenderPipeline,
41 filterable: bool,
42}
43
44pub struct MipmapGenerator {
51 filtering_layout: wgpu::BindGroupLayout,
52 nonfiltering_layout: wgpu::BindGroupLayout,
53 linear_sampler: wgpu::Sampler,
54 nearest_sampler: wgpu::Sampler,
55 pipelines: HashMap<wgpu::TextureFormat, MipPipeline>,
56}
57
58impl MipmapGenerator {
59 fn bind_group_layout(device: &wgpu::Device, filterable: bool) -> wgpu::BindGroupLayout {
60 device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
61 label: Some("mipmap-blit-bind-group-layout"),
62 entries: &[
63 wgpu::BindGroupLayoutEntry {
64 binding: 0,
65 visibility: wgpu::ShaderStages::FRAGMENT,
66 ty: wgpu::BindingType::Texture {
67 sample_type: wgpu::TextureSampleType::Float { filterable },
68 view_dimension: wgpu::TextureViewDimension::D2,
69 multisampled: false,
70 },
71 count: None,
72 },
73 wgpu::BindGroupLayoutEntry {
74 binding: 1,
75 visibility: wgpu::ShaderStages::FRAGMENT,
76 ty: wgpu::BindingType::Sampler(if filterable {
77 wgpu::SamplerBindingType::Filtering
78 } else {
79 wgpu::SamplerBindingType::NonFiltering
80 }),
81 count: None,
82 },
83 ],
84 })
85 }
86
87 fn build_pipeline(
88 device: &wgpu::Device,
89 module: &wgpu::ShaderModule,
90 layout: &wgpu::BindGroupLayout,
91 format: wgpu::TextureFormat,
92 ) -> wgpu::RenderPipeline {
93 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
94 label: Some("mipmap-blit-pipeline-layout"),
95 bind_group_layouts: &[Some(layout)],
96 immediate_size: 0,
97 });
98
99 device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
100 label: Some("mipmap-blit-pipeline"),
101 layout: Some(&pipeline_layout),
102 vertex: wgpu::VertexState {
103 module,
104 entry_point: Some("vs_main"),
105 compilation_options: Default::default(),
106 buffers: &[],
107 },
108 primitive: wgpu::PrimitiveState {
109 topology: wgpu::PrimitiveTopology::TriangleList,
110 strip_index_format: None,
111 front_face: wgpu::FrontFace::Ccw,
112 cull_mode: None,
113 unclipped_depth: false,
114 polygon_mode: wgpu::PolygonMode::Fill,
115 conservative: false,
116 },
117 depth_stencil: None,
118 multisample: wgpu::MultisampleState::default(),
119 fragment: Some(wgpu::FragmentState {
120 module,
121 entry_point: Some("fs_main"),
122 compilation_options: Default::default(),
123 targets: &[Some(wgpu::ColorTargetState {
124 format,
125 blend: None,
126 write_mask: wgpu::ColorWrites::ALL,
127 })],
128 }),
129 multiview_mask: None,
130 cache: None,
131 })
132 }
133
134 fn new(device: &wgpu::Device) -> Self {
135 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
136 label: Some("mipmap-blit-shader"),
137 source: wgpu::ShaderSource::Wgsl(MIP_SHADER.into()),
138 });
139
140 let filtering_layout = Self::bind_group_layout(device, true);
141 let nonfiltering_layout = Self::bind_group_layout(device, false);
142
143 let linear_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
144 label: Some("mipmap-linear-sampler"),
145 mag_filter: wgpu::FilterMode::Linear,
146 min_filter: wgpu::FilterMode::Linear,
147 ..Default::default()
148 });
149 let nearest_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
150 label: Some("mipmap-nearest-sampler"),
151 mag_filter: wgpu::FilterMode::Nearest,
152 min_filter: wgpu::FilterMode::Nearest,
153 ..Default::default()
154 });
155
156 let mut pipelines = HashMap::new();
157 for &format in SUPPORTED_FORMATS {
158 let filterable = format != wgpu::TextureFormat::Rgba32Float;
161 let layout = if filterable {
162 &filtering_layout
163 } else {
164 &nonfiltering_layout
165 };
166 let pipeline = Self::build_pipeline(device, &module, layout, format);
167 pipelines.insert(format, MipPipeline { pipeline, filterable });
168 }
169
170 Self {
171 filtering_layout,
172 nonfiltering_layout,
173 linear_sampler,
174 nearest_sampler,
175 pipelines,
176 }
177 }
178
179 pub fn generate_mips(
185 &self,
186 device: &wgpu::Device,
187 queue: &wgpu::Queue,
188 texture: &wgpu::Texture,
189 format: wgpu::TextureFormat,
190 mip_count: u32,
191 layer_count: u32,
192 ) {
193 if mip_count <= 1 {
194 return;
195 }
196
197 let Some(entry) = self.pipelines.get(&format) else {
198 tracing::error!("MipmapGenerator: no blit pipeline for format {format:?}, skipping mip generation");
199 return;
200 };
201 let layout = if entry.filterable {
202 &self.filtering_layout
203 } else {
204 &self.nonfiltering_layout
205 };
206 let sampler = if entry.filterable {
207 &self.linear_sampler
208 } else {
209 &self.nearest_sampler
210 };
211
212 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
213 label: Some("mipmap-blit-encoder"),
214 });
215
216 for layer in 0..layer_count {
217 for level in 1..mip_count {
218 let src_view = texture.create_view(&wgpu::TextureViewDescriptor {
219 dimension: Some(wgpu::TextureViewDimension::D2),
220 base_mip_level: level - 1,
221 mip_level_count: Some(1),
222 base_array_layer: layer,
223 array_layer_count: Some(1),
224 ..Default::default()
225 });
226 let dst_view = texture.create_view(&wgpu::TextureViewDescriptor {
227 dimension: Some(wgpu::TextureViewDimension::D2),
228 base_mip_level: level,
229 mip_level_count: Some(1),
230 base_array_layer: layer,
231 array_layer_count: Some(1),
232 ..Default::default()
233 });
234
235 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
236 label: Some("mipmap-blit-bind-group"),
237 layout,
238 entries: &[
239 wgpu::BindGroupEntry {
240 binding: 0,
241 resource: wgpu::BindingResource::TextureView(&src_view),
242 },
243 wgpu::BindGroupEntry {
244 binding: 1,
245 resource: wgpu::BindingResource::Sampler(sampler),
246 },
247 ],
248 });
249
250 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
251 label: Some("mipmap-blit-pass"),
252 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
253 view: &dst_view,
254 depth_slice: None,
255 resolve_target: None,
256 ops: wgpu::Operations {
257 load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
258 store: wgpu::StoreOp::Store,
259 },
260 })],
261 depth_stencil_attachment: None,
262 timestamp_writes: None,
263 occlusion_query_set: None,
264 multiview_mask: None,
265 });
266 pass.set_pipeline(&entry.pipeline);
267 pass.set_bind_group(0, &bind_group, &[]);
268 pass.draw(0..3, 0..1);
269 drop(pass);
270 }
271 }
272
273 queue.submit(std::iter::once(encoder.finish()));
274 }
275}
276
277impl LazyResource<WGPUBackend> for MipmapGenerator {
278 type Deps<'a> = ();
279
280 fn construct<'a>(backend: &WGPUBackend, _deps: &()) -> Option<Self> {
281 Some(Self::new(&backend.device))
282 }
283}