Skip to main content

pebble/wgpu/
mipmap.rs

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
30/// Formats [`MipmapGenerator`] knows how to downsample — matches the formats
31/// `GPUTexture` can decode into (see `wgpu::textures::bytes_per_pixel`).
32const 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
44/// GPU-side mipmap generator: a shared render pipeline (per texture format)
45/// that downsamples one mip level into the next via a fullscreen-triangle
46/// blit, plus the bind group layouts/samplers it needs.
47///
48/// Built once as a [`LazyResource`] and reused by every [`GPUTexture`](crate::wgpu::textures::GPUTexture)
49/// upload that requests mips.
50pub 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            // Rgba32Float has no filterable-sampling support without an extra
159            // device feature, so it gets a nearest-sampled, non-filtering path.
160            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    /// Downsamples `texture` (already holding valid data at mip level 0 of
180    /// every layer) into its remaining `mip_count - 1` levels, entirely on
181    /// the GPU. `layer_count` is 1 for a plain 2D texture, N for a texture
182    /// array, or 6 for a cubemap — each array layer/face is downsampled
183    /// independently via its own single-layer 2D view.
184    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}