1use pfx_gpu::{GpuProfiler, wgpu};
2
3pub const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba16Float;
4const WORKGROUP: u32 = 8;
5
6pub const SHARPEN_WGSL: &str = r#"
7struct Params {
8 width: u32,
9 height: u32,
10 strength: f32,
11 exposure: f32,
12}
13
14@group(0) @binding(0) var<uniform> params: Params;
15@group(0) @binding(1) var source: texture_2d<f32>;
16@group(0) @binding(2) var display: texture_storage_2d<rgba16float, write>;
17
18fn tap(p: vec2<i32>) -> vec3<f32> {
19 let size = vec2<i32>(i32(params.width), i32(params.height));
20 let c = max(textureLoad(source, clamp(p, vec2<i32>(0), size - vec2<i32>(1)), 0).rgb * params.exposure, vec3<f32>(0.0));
21 return c / (1.0 + max(c.r, max(c.g, c.b)));
22}
23
24@compute @workgroup_size(8, 8)
25fn sharpen(@builtin(global_invocation_id) gid: vec3<u32>) {
26 if (gid.x >= params.width || gid.y >= params.height) {
27 return;
28 }
29 let p = vec2<i32>(gid.xy);
30 let centre = textureLoad(source, p, 0);
31 let a = tap(p + vec2<i32>(-1, -1));
32 let b = tap(p + vec2<i32>(0, -1));
33 let c = tap(p + vec2<i32>(1, -1));
34 let d = tap(p + vec2<i32>(-1, 0));
35 let e = tap(p);
36 let f = tap(p + vec2<i32>(1, 0));
37 let g = tap(p + vec2<i32>(-1, 1));
38 let h = tap(p + vec2<i32>(0, 1));
39 let i = tap(p + vec2<i32>(1, 1));
40 let cross_lo = min(min(min(d, e), min(f, b)), h);
41 let cross_hi = max(max(max(d, e), max(f, b)), h);
42 let lo = cross_lo + min(cross_lo, min(min(a, c), min(g, i)));
43 let hi = cross_hi + max(cross_hi, max(max(a, c), max(g, i)));
44 let amp = sqrt(clamp(min(lo, vec3<f32>(2.0) - hi) / max(hi, vec3<f32>(1.0e-5)), vec3<f32>(0.0), vec3<f32>(1.0)));
45 let w = amp * (-params.strength / mix(8.0, 5.0, params.strength));
46 let t = clamp((b * w + d * w + f * w + h * w + e) / (vec3<f32>(1.0) + 4.0 * w), vec3<f32>(0.0), vec3<f32>(0.999));
47 let back = t / (1.0 - max(t.r, max(t.g, t.b))) / params.exposure;
48 textureStore(display, p, vec4<f32>(back, centre.a));
49}
50
51@compute @workgroup_size(8, 8)
52fn copy(@builtin(global_invocation_id) gid: vec3<u32>) {
53 if (gid.x >= params.width || gid.y >= params.height) {
54 return;
55 }
56 let p = vec2<i32>(gid.xy);
57 textureStore(display, p, textureLoad(source, p, 0));
58}
59"#;
60
61pub fn check_strength(strength: f32) -> Result<f32, String> {
62 if strength.is_finite() && (0.0..=1.0).contains(&strength) {
63 Ok(strength)
64 } else {
65 Err("sharpen strength must be between 0 and 1".into())
66 }
67}
68
69fn squash(c: [f32; 3], exposure: f32) -> [f32; 3] {
70 let c = c.map(|v| (v * exposure).max(0.0));
71 let peak = c[0].max(c[1]).max(c[2]);
72 c.map(|v| v / (1.0 + peak))
73}
74
75pub fn sharpen(pixels: &[[f32; 4]], width: usize, strength: f32, exposure: f32) -> Vec<[f32; 4]> {
76 if strength == 0.0 || width == 0 {
77 return pixels.to_vec();
78 }
79 let height = pixels.len() / width;
80 let tap = |x: isize, y: isize| {
81 let x = x.clamp(0, width as isize - 1) as usize;
82 let y = y.clamp(0, height as isize - 1) as usize;
83 let p = pixels[y * width + x];
84 squash([p[0], p[1], p[2]], exposure)
85 };
86 let peak = -strength / (8.0 + (5.0 - 8.0) * strength);
87 let mut out = Vec::with_capacity(pixels.len());
88 for y in 0..height as isize {
89 for x in 0..width as isize {
90 let taps: [[f32; 3]; 9] =
91 std::array::from_fn(|k| tap(x + k as isize % 3 - 1, y + k as isize / 3 - 1));
92 let mut back = [0.0; 3];
93 for (channel, slot) in back.iter_mut().enumerate() {
94 let v: [f32; 9] = std::array::from_fn(|k| taps[k][channel]);
95 let cross = [v[1], v[3], v[4], v[5], v[7]];
96 let corners = [v[0], v[2], v[6], v[8]];
97 let cross_lo = cross.iter().copied().fold(f32::INFINITY, f32::min);
98 let cross_hi = cross.iter().copied().fold(f32::NEG_INFINITY, f32::max);
99 let lo = cross_lo + corners.iter().copied().fold(cross_lo, f32::min);
100 let hi = cross_hi + corners.iter().copied().fold(cross_hi, f32::max);
101 let amp = ((lo.min(2.0 - hi) / hi.max(1.0e-5)).clamp(0.0, 1.0)).sqrt();
102 let w = amp * peak;
103 *slot = ((v[1] * w + v[3] * w + v[5] * w + v[7] * w + v[4]) / (1.0 + 4.0 * w))
104 .clamp(0.0, 0.999);
105 }
106 let top = back[0].max(back[1]).max(back[2]);
107 let centre = pixels[y as usize * width + x as usize];
108 out.push([
109 back[0] / (1.0 - top) / exposure,
110 back[1] / (1.0 - top) / exposure,
111 back[2] / (1.0 - top) / exposure,
112 centre[3],
113 ]);
114 }
115 }
116 out
117}
118
119pub struct Pass {
120 layout: wgpu::BindGroupLayout,
121 sharpen: wgpu::ComputePipeline,
122 copy: wgpu::ComputePipeline,
123 params: wgpu::Buffer,
124 target: Option<(wgpu::Texture, wgpu::TextureView)>,
125 width: u32,
126 height: u32,
127 viewport: (u32, u32),
128}
129
130impl Pass {
131 pub fn new(device: &wgpu::Device, width: u32, height: u32) -> Self {
132 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
133 label: Some("crisp"),
134 entries: &[
135 wgpu::BindGroupLayoutEntry {
136 binding: 0,
137 visibility: wgpu::ShaderStages::COMPUTE,
138 ty: wgpu::BindingType::Buffer {
139 ty: wgpu::BufferBindingType::Uniform,
140 has_dynamic_offset: false,
141 min_binding_size: None,
142 },
143 count: None,
144 },
145 wgpu::BindGroupLayoutEntry {
146 binding: 1,
147 visibility: wgpu::ShaderStages::COMPUTE,
148 ty: wgpu::BindingType::Texture {
149 sample_type: wgpu::TextureSampleType::Float { filterable: false },
150 view_dimension: wgpu::TextureViewDimension::D2,
151 multisampled: false,
152 },
153 count: None,
154 },
155 wgpu::BindGroupLayoutEntry {
156 binding: 2,
157 visibility: wgpu::ShaderStages::COMPUTE,
158 ty: wgpu::BindingType::StorageTexture {
159 access: wgpu::StorageTextureAccess::WriteOnly,
160 format: FORMAT,
161 view_dimension: wgpu::TextureViewDimension::D2,
162 },
163 count: None,
164 },
165 ],
166 });
167 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
168 label: Some("crisp"),
169 bind_group_layouts: &[Some(&layout)],
170 immediate_size: 0,
171 });
172 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
173 label: Some("crisp"),
174 source: wgpu::ShaderSource::Wgsl(SHARPEN_WGSL.into()),
175 });
176 let pipeline = |entry: &str| {
177 device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
178 label: Some(entry),
179 layout: Some(&pipeline_layout),
180 module: &module,
181 entry_point: Some(entry),
182 compilation_options: Default::default(),
183 cache: None,
184 })
185 };
186 Self {
187 sharpen: pipeline("sharpen"),
188 copy: pipeline("copy"),
189 layout,
190 params: device.create_buffer(&wgpu::BufferDescriptor {
191 label: Some("crisp params"),
192 size: 16,
193 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
194 mapped_at_creation: false,
195 }),
196 target: None,
197 width,
198 height,
199 viewport: (width, height),
200 }
201 }
202
203 pub fn resize(&mut self, width: u32, height: u32) {
204 if (width, height) != (self.width, self.height) {
205 self.target = None;
206 self.width = width;
207 self.height = height;
208 }
209 self.viewport = (width, height);
210 }
211
212 pub fn set_viewport(&mut self, width: u32, height: u32) -> Result<(), String> {
213 crate::viewport::check([width, height], [self.width, self.height])?;
214 self.viewport = (width, height);
215 Ok(())
216 }
217
218 pub fn target(&self) -> Option<&wgpu::Texture> {
219 self.target.as_ref().map(|(texture, _)| texture)
220 }
221
222 #[allow(clippy::too_many_arguments)]
223 pub fn encode(
224 &mut self,
225 device: &wgpu::Device,
226 queue: &wgpu::Queue,
227 encoder: &mut wgpu::CommandEncoder,
228 source: &wgpu::TextureView,
229 strength: f32,
230 exposure: f32,
231 profiler: Option<&mut GpuProfiler>,
232 ) -> &wgpu::TextureView {
233 let (capacity_width, capacity_height) = (self.width, self.height);
234 let (width, height) = self.viewport;
235 let (_, target) = self.target.get_or_insert_with(|| {
236 let texture = device.create_texture(&wgpu::TextureDescriptor {
237 label: Some("crisp display"),
238 size: wgpu::Extent3d {
239 width: capacity_width,
240 height: capacity_height,
241 depth_or_array_layers: 1,
242 },
243 mip_level_count: 1,
244 sample_count: 1,
245 dimension: wgpu::TextureDimension::D2,
246 format: FORMAT,
247 usage: wgpu::TextureUsages::STORAGE_BINDING
248 | wgpu::TextureUsages::TEXTURE_BINDING
249 | wgpu::TextureUsages::RENDER_ATTACHMENT
250 | wgpu::TextureUsages::COPY_SRC,
251 view_formats: &[],
252 });
253 let view = texture.create_view(&Default::default());
254 (texture, view)
255 });
256 let mut bytes = Vec::with_capacity(16);
257 bytes.extend_from_slice(&width.to_le_bytes());
258 bytes.extend_from_slice(&height.to_le_bytes());
259 bytes.extend_from_slice(&strength.to_le_bytes());
260 bytes.extend_from_slice(&exposure.max(1.0e-8).to_le_bytes());
261 queue.write_buffer(&self.params, 0, &bytes);
262 let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
263 label: Some("crisp"),
264 layout: &self.layout,
265 entries: &[
266 wgpu::BindGroupEntry {
267 binding: 0,
268 resource: self.params.as_entire_binding(),
269 },
270 wgpu::BindGroupEntry {
271 binding: 1,
272 resource: wgpu::BindingResource::TextureView(source),
273 },
274 wgpu::BindGroupEntry {
275 binding: 2,
276 resource: wgpu::BindingResource::TextureView(target),
277 },
278 ],
279 });
280 let sharpening = strength > 0.0;
281 let mut profiler = profiler;
282 let timing = profiler
283 .as_deref_mut()
284 .and_then(|timer| timer.pass(if sharpening { "sharpen" } else { "crisp copy" }));
285 {
286 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
287 label: Some("crisp"),
288 timestamp_writes: profiler
289 .as_deref()
290 .and_then(|timer| timer.compute_writes(timing)),
291 });
292 pass.set_pipeline(if sharpening {
293 &self.sharpen
294 } else {
295 &self.copy
296 });
297 pass.set_bind_group(0, &bind, &[]);
298 pass.dispatch_workgroups(width.div_ceil(WORKGROUP), height.div_ceil(WORKGROUP), 1);
299 }
300 target
301 }
302}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307
308 fn edge(width: usize, height: usize) -> Vec<[f32; 4]> {
309 (0..width * height)
310 .map(|k| {
311 let x = k % width;
312 let v = match x {
313 0..=3 => 0.05,
314 4 => 0.3,
315 5 => 0.6,
316 _ => 0.85,
317 };
318 [v, v * 0.9, v * 0.8, 1.0]
319 })
320 .collect()
321 }
322
323 #[test]
324 fn strength_zero_is_the_input() {
325 let pixels = edge(10, 4);
326 assert_eq!(sharpen(&pixels, 10, 0.0, 1.0), pixels);
327 }
328
329 #[test]
330 fn a_flat_image_stays_flat() {
331 let pixels = vec![[0.4, 0.2, 0.1, 0.7]; 36];
332 for strength in [0.25, 1.0] {
333 for p in sharpen(&pixels, 6, strength, 1.3) {
334 for c in 0..3 {
335 assert!((p[c] - pixels[0][c]).abs() < 1e-5, "{p:?}");
336 }
337 assert_eq!(p[3], 0.7);
338 }
339 }
340 }
341
342 #[test]
343 fn an_edge_gets_steeper_with_strength_and_stays_finite() {
344 let pixels = edge(10, 4);
345 let slope = |image: &[[f32; 4]]| image[10 + 5][0] - image[10 + 4][0];
346 let soft = slope(&pixels);
347 let half = sharpen(&pixels, 10, 0.5, 1.0);
348 let full = sharpen(&pixels, 10, 1.0, 1.0);
349 assert!(slope(&half) > soft, "{} against {soft}", slope(&half));
350 assert!(slope(&full) > slope(&half));
351 for p in full.iter().chain(&half) {
352 assert!(p.iter().all(|v| v.is_finite() && *v >= 0.0), "{p:?}");
353 }
354 }
355
356 #[test]
357 fn hdr_values_round_trip_through_the_squash() {
358 let pixels = vec![[12.0, 3.0, 0.5, 1.0]; 25];
359 for p in sharpen(&pixels, 5, 1.0, 0.5) {
360 for c in 0..3 {
361 assert!((p[c] - pixels[0][c]).abs() / pixels[0][c] < 1e-3, "{p:?}");
362 }
363 }
364 }
365
366 #[test]
367 fn strength_outside_zero_to_one_is_refused() {
368 assert!(check_strength(-0.1).is_err());
369 assert!(check_strength(1.5).is_err());
370 assert!(check_strength(f32::NAN).is_err());
371 assert_eq!(check_strength(0.3), Ok(0.3));
372 }
373
374 #[test]
375 fn the_shader_validates() {
376 let module = naga::front::wgsl::parse_str(SHARPEN_WGSL).unwrap();
377 naga::valid::Validator::new(
378 naga::valid::ValidationFlags::all(),
379 naga::valid::Capabilities::all(),
380 )
381 .validate(&module)
382 .unwrap();
383 }
384}