1mod net;
2pub mod pfm;
3mod setup;
4mod tile;
5mod tza;
6mod wgsl;
7
8use bytemuck::{Pod, Zeroable};
9use pfx_gpu::Gpu;
10use wgpu::util::DeviceExt;
11
12use crate::net::Net;
13use crate::tile::Plan;
14
15pub const OIDN_VERSION: &str = "2.5.1";
16const MACS_PER_SUBMISSION: u64 = 24_000_000_000;
17const WEIGHTS: &[u8] = include_bytes!("../weights/rt_hdr_alb_nrm.tza");
18
19#[repr(C)]
20#[derive(Clone, Copy, Pod, Zeroable, Default)]
21struct Job {
22 src: [i32; 2],
23 dst: [i32; 2],
24 size: [i32; 2],
25 out_src: [i32; 2],
26 out_dst: [i32; 2],
27 out_size: [i32; 2],
28 dims: [u32; 2],
29 image: [u32; 2],
30 stride: u32,
31 norm: f32,
32 rcp_norm: f32,
33 pad: u32,
34}
35
36struct Step {
37 pipeline: usize,
38 bindings: wgpu::BindGroup,
39 groups: [u32; 3],
40 macs: u64,
41}
42
43struct Images {
44 color: wgpu::Buffer,
45 albedo: wgpu::Buffer,
46 normal: wgpu::Buffer,
47 output: wgpu::Buffer,
48}
49
50struct Setup {
51 width: u32,
52 height: u32,
53 stride: u32,
54 plan: Plan,
55 images: Images,
56 pipelines: Vec<wgpu::ComputePipeline>,
57 steps: Vec<Step>,
58}
59
60struct Model {
61 net: Net,
62 weights: Vec<wgpu::Buffer>,
63 biases: Vec<wgpu::Buffer>,
64}
65
66pub struct Denoiser {
67 model: Model,
68 setup: Option<Setup>,
69}
70
71fn norm_scales() -> (f32, f32) {
72 let max = 1.926_532_5e-1 * (65504.0f32 + 6.260_261e-3).ln() + 9.986_201_5e-1;
73 ((1.0f64 / f64::from(max)) as f32, max)
74}
75
76impl Denoiser {
77 pub fn new(gpu: &Gpu) -> Result<Self, String> {
78 let net = net::build(&tza::parse(WEIGHTS)?, 9)?;
79 let upload = |label: &str, data: &[f32]| {
80 gpu.device
81 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
82 label: Some(label),
83 contents: bytemuck::cast_slice(data),
84 usage: wgpu::BufferUsages::STORAGE,
85 })
86 };
87 let weights = net
88 .layers
89 .iter()
90 .map(|layer| upload(&layer.name, &layer.weights))
91 .collect();
92 let biases = net
93 .layers
94 .iter()
95 .map(|layer| upload(&layer.name, &layer.bias))
96 .collect();
97 Ok(Self {
98 model: Model {
99 net,
100 weights,
101 biases,
102 },
103 setup: None,
104 })
105 }
106
107 pub fn tiles(&self) -> Option<(u32, u32, usize)> {
108 self.setup
109 .as_ref()
110 .map(|setup| (setup.plan.tile_w, setup.plan.tile_h, setup.plan.tiles.len()))
111 }
112
113 pub fn denoise(
114 &mut self,
115 gpu: &Gpu,
116 width: u32,
117 height: u32,
118 color: &[[f32; 4]],
119 albedo: &[[f32; 4]],
120 normal: &[[f32; 4]],
121 ) -> Result<Vec<[f32; 4]>, String> {
122 let pixels = (width as usize)
123 .checked_mul(height as usize)
124 .ok_or("image dimensions overflow")?;
125 if color.len() != pixels || albedo.len() != pixels || normal.len() != pixels {
126 return Err("image size mismatch".into());
127 }
128 self.prepare(gpu, width, height)?;
129 let setup = self.setup.as_ref().unwrap();
130 for (buffer, pixels) in [
131 (&setup.images.color, color),
132 (&setup.images.albedo, albedo),
133 (&setup.images.normal, normal),
134 ] {
135 upload(gpu, buffer, width, setup.stride, pixels);
136 }
137 self.run(gpu)?;
138 self.read_output(gpu)
139 }
140
141 pub fn denoise_textures(
142 &mut self,
143 gpu: &Gpu,
144 color: &wgpu::Texture,
145 albedo: &wgpu::Texture,
146 normal: &wgpu::Texture,
147 ) -> Result<Vec<[f32; 4]>, String> {
148 self.load_textures(gpu, color, albedo, normal)?;
149 self.run(gpu)?;
150 self.read_output(gpu)
151 }
152
153 pub fn denoise_textures_into(
154 &mut self,
155 gpu: &Gpu,
156 color: &wgpu::Texture,
157 albedo: &wgpu::Texture,
158 normal: &wgpu::Texture,
159 output: &wgpu::Texture,
160 ) -> Result<(), String> {
161 check_texture(output, color.width(), color.height())?;
162 if !output.usage().contains(wgpu::TextureUsages::COPY_DST) {
163 return Err("the output texture needs COPY_DST".into());
164 }
165 self.load_textures(gpu, color, albedo, normal)?;
166 self.run(gpu)?;
167 let setup = self.setup.as_ref().unwrap();
168 let mut encoder = gpu
169 .device
170 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
171 label: Some("denoise output"),
172 });
173 encoder.copy_buffer_to_texture(
174 wgpu::TexelCopyBufferInfo {
175 buffer: &setup.images.output,
176 layout: image_layout(setup.stride),
177 },
178 output.as_image_copy(),
179 extent(setup.width, setup.height),
180 );
181 gpu.queue.submit(Some(encoder.finish()));
182 Ok(())
183 }
184
185 fn load_textures(
186 &mut self,
187 gpu: &Gpu,
188 color: &wgpu::Texture,
189 albedo: &wgpu::Texture,
190 normal: &wgpu::Texture,
191 ) -> Result<(), String> {
192 let (width, height) = (color.width(), color.height());
193 for texture in [color, albedo, normal] {
194 check_texture(texture, width, height)?;
195 if !texture.usage().contains(wgpu::TextureUsages::COPY_SRC) {
196 return Err("the input textures need COPY_SRC".into());
197 }
198 }
199 self.prepare(gpu, width, height)?;
200 let setup = self.setup.as_ref().unwrap();
201 let mut encoder = gpu
202 .device
203 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
204 label: Some("denoise inputs"),
205 });
206 for (texture, buffer) in [
207 (color, &setup.images.color),
208 (albedo, &setup.images.albedo),
209 (normal, &setup.images.normal),
210 ] {
211 encoder.copy_texture_to_buffer(
212 texture.as_image_copy(),
213 wgpu::TexelCopyBufferInfo {
214 buffer,
215 layout: image_layout(setup.stride),
216 },
217 extent(width, height),
218 );
219 }
220 gpu.queue.submit(Some(encoder.finish()));
221 Ok(())
222 }
223
224 fn run(&self, gpu: &Gpu) -> Result<(), String> {
225 let setup = self.setup.as_ref().ok_or("the denoiser has no image")?;
226 let mut at = 0;
227 while at < setup.steps.len() {
228 let mut encoder = gpu
229 .device
230 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
231 label: Some("denoise"),
232 });
233 {
234 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
235 label: Some("denoise"),
236 timestamp_writes: None,
237 });
238 let mut spent = 0u64;
239 while at < setup.steps.len() {
240 let step = &setup.steps[at];
241 if spent > 0 && spent + step.macs > MACS_PER_SUBMISSION {
242 break;
243 }
244 pass.set_pipeline(&setup.pipelines[step.pipeline]);
245 pass.set_bind_group(0, &step.bindings, &[]);
246 pass.dispatch_workgroups(step.groups[0], step.groups[1], step.groups[2]);
247 spent += step.macs;
248 at += 1;
249 }
250 }
251 gpu.queue.submit(Some(encoder.finish()));
252 }
253 Ok(())
254 }
255
256 fn read_output(&self, gpu: &Gpu) -> Result<Vec<[f32; 4]>, String> {
257 let setup = self.setup.as_ref().ok_or("the denoiser has no image")?;
258 let size = setup.images.output.size();
259 let read = gpu.device.create_buffer(&wgpu::BufferDescriptor {
260 label: Some("denoise readback"),
261 size,
262 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
263 mapped_at_creation: false,
264 });
265 let mut encoder = gpu
266 .device
267 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
268 label: Some("denoise readback"),
269 });
270 encoder.copy_buffer_to_buffer(&setup.images.output, 0, &read, 0, size);
271 gpu.queue.submit(Some(encoder.finish()));
272 let (send, receive) = std::sync::mpsc::channel();
273 let slice = read.slice(..);
274 slice.map_async(wgpu::MapMode::Read, move |result| {
275 let _ = send.send(result);
276 });
277 gpu.device
278 .poll(wgpu::PollType::wait_indefinitely())
279 .map_err(|error| format!("{error:?}"))?;
280 receive
281 .recv()
282 .map_err(|error| error.to_string())?
283 .map_err(|error| error.to_string())?;
284 let pixels = {
285 let bytes = slice.get_mapped_range().expect("a mapped range");
286 let all: &[[f32; 4]] = bytemuck::cast_slice(&bytes);
287 let mut pixels = Vec::with_capacity(setup.width as usize * setup.height as usize);
288 for row in all.chunks_exact(setup.stride as usize) {
289 pixels.extend_from_slice(&row[..setup.width as usize]);
290 }
291 pixels
292 };
293 read.unmap();
294 Ok(pixels)
295 }
296
297 fn prepare(&mut self, gpu: &Gpu, width: u32, height: u32) -> Result<(), String> {
298 if width == 0 || height == 0 {
299 return Err("image dimensions must be nonzero".into());
300 }
301 if self
302 .setup
303 .as_ref()
304 .is_some_and(|setup| setup.width == width && setup.height == height)
305 {
306 return Ok(());
307 }
308 self.setup = None;
309 self.setup = Some(setup::build(gpu, &self.model, width, height)?);
310 Ok(())
311 }
312}
313
314fn check_texture(texture: &wgpu::Texture, width: u32, height: u32) -> Result<(), String> {
315 if texture.format() != wgpu::TextureFormat::Rgba32Float {
316 return Err("the denoiser reads and writes Rgba32Float textures".into());
317 }
318 if texture.width() != width || texture.height() != height {
319 return Err("image size mismatch".into());
320 }
321 Ok(())
322}
323
324fn image_layout(stride: u32) -> wgpu::TexelCopyBufferLayout {
325 wgpu::TexelCopyBufferLayout {
326 offset: 0,
327 bytes_per_row: Some(stride * 16),
328 rows_per_image: None,
329 }
330}
331
332fn extent(width: u32, height: u32) -> wgpu::Extent3d {
333 wgpu::Extent3d {
334 width,
335 height,
336 depth_or_array_layers: 1,
337 }
338}
339
340fn upload(gpu: &Gpu, buffer: &wgpu::Buffer, width: u32, stride: u32, pixels: &[[f32; 4]]) {
341 if width == stride {
342 gpu.queue
343 .write_buffer(buffer, 0, bytemuck::cast_slice(pixels));
344 return;
345 }
346 let mut padded = vec![[0.0f32; 4]; (buffer.size() / 16) as usize];
347 for (row, line) in pixels.chunks_exact(width as usize).enumerate() {
348 padded[row * stride as usize..][..width as usize].copy_from_slice(line);
349 }
350 gpu.queue
351 .write_buffer(buffer, 0, bytemuck::cast_slice(&padded));
352}