1mod constants;
2
3use crate::config::NtscShaderConfig;
4use crate::renderer::{PipelineShader, Shaders};
5use jgenesis_common::frontend::{CompositeParams, RenderFrameOptions, SamplesPerColorCycle};
6use std::sync::Arc;
7use wgpu::util::DeviceExt;
8
9const BACKDROP_PIXELS: u32 = 6;
10
11const FIR_BUFFER_LEN: usize = 84;
12
13struct NtscFilters {
14    luma_bsf: &'static [f32],
15    chroma_bpf: &'static [f32],
16    y_encode_lpf: &'static [f32],
17    y_decode_lpf: &'static [f32],
18    uv_lpf: &'static [f32],
19}
20
21impl NtscFilters {
22    fn from(samples_per_color_cycle: SamplesPerColorCycle) -> Self {
23        match samples_per_color_cycle {
24            SamplesPerColorCycle::Fifteen => Self {
25                luma_bsf: constants::LUMA_BSF_15_COEFFICIENTS,
26                chroma_bpf: constants::CHROMA_BPF_15_COEFFICIENTS,
27                y_encode_lpf: constants::Y_ENCODE_LPF_15_COEFFICIENTS,
28                y_decode_lpf: constants::Y_DECODE_LPF_15_COEFFICIENTS,
29                uv_lpf: constants::UV_LPF_15_COEFFICIENTS,
30            },
31            SamplesPerColorCycle::Twelve => Self {
32                luma_bsf: constants::LUMA_BSF_12_COEFFICIENTS,
33                chroma_bpf: constants::CHROMA_BPF_12_COEFFICIENTS,
34                y_encode_lpf: constants::Y_ENCODE_LPF_12_COEFFICIENTS,
35                y_decode_lpf: constants::Y_DECODE_LPF_12_COEFFICIENTS,
36                uv_lpf: constants::UV_LPF_12_COEFFICIENTS,
37            },
38        }
39    }
40}
41
42#[repr(C)]
43#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
44struct ImmediateParams {
45    frame_phase_offset: i32,
46    per_line_phase_offset: i32,
47}
48
49impl ImmediateParams {
50    const ZERO: Self = Self { frame_phase_offset: 0, per_line_phase_offset: 0 };
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum NtscShaderVariant {
55    // Input frame buffer contains RGB888 colors
56    // Encode from RGB to NTSC then decode back to RGB
57    Rgb,
58    // Input frame buffer contains 9-bit NES colors (2-bit luma, 4-bit hue, 3-bit color emphasis)
59    // Emulate the NES PPU's NTSC output, then decode to RGB
60    NesPpu,
61}
62
63pub struct NtscShader {
64    output: Arc<wgpu::Texture>,
65    ntsc_frame_size: wgpu::Extent3d,
66    samples_per_color_cycle: SamplesPerColorCycle,
67    immediates_bind_group_layout: wgpu::BindGroupLayout,
68    immediates_bind_group: wgpu::BindGroup,
69    rgb_to_ntsc_bind_group: wgpu::BindGroup,
70    rgb_to_ntsc_pipeline: wgpu::ComputePipeline,
71    separate_luma_chroma_bind_group: wgpu::BindGroup,
72    separate_luma_chroma_pipeline: wgpu::ComputePipeline,
73    luma_chroma_to_rgb_bind_group: wgpu::BindGroup,
74    luma_chroma_to_rgb_pipeline: wgpu::ComputePipeline,
75}
76
77impl NtscShader {
78    pub fn create(
79        device: &wgpu::Device,
80        shaders: &Shaders,
81        input: &wgpu::Texture,
82        params: CompositeParams,
83        config: NtscShaderConfig,
84        variant: NtscShaderVariant,
85    ) -> Self {
86        let ntsc_texture_descriptor = wgpu::TextureDescriptor {
87            label: None,
88            size: wgpu::Extent3d {
89                width: params.upscale_factor * (input.width() + 2 * BACKDROP_PIXELS),
90                height: input.height(),
91                depth_or_array_layers: 1,
92            },
93            mip_level_count: 1,
94            sample_count: 1,
95            dimension: wgpu::TextureDimension::D2,
96            format: wgpu::TextureFormat::R32Float,
97            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::STORAGE_BINDING,
98            view_formats: &[],
99        };
100
101        let ntsc_frame = device.create_texture(&ntsc_texture_descriptor);
102        let ntsc_pass = device.create_texture(&ntsc_texture_descriptor);
103        let ntsc_stop = device.create_texture(&ntsc_texture_descriptor);
104
105        let output_frame = device.create_texture(&wgpu::TextureDescriptor {
106            label: "ntsc_output_texture".into(),
107            size: wgpu::Extent3d {
108                width: params.upscale_factor * input.width(),
109                height: input.height(),
110                depth_or_array_layers: 1,
111            },
112            mip_level_count: 1,
113            sample_count: 1,
114            dimension: wgpu::TextureDimension::D2,
115            format: wgpu::TextureFormat::Rgba8Unorm,
116            usage: wgpu::TextureUsages::TEXTURE_BINDING
117                | wgpu::TextureUsages::STORAGE_BINDING
118                | wgpu::TextureUsages::COPY_SRC,
119            view_formats: &[wgpu::TextureFormat::Rgba8UnormSrgb],
120        });
121
122        let input_view = input.create_view(&wgpu::TextureViewDescriptor::default());
123        let ntsc_view = ntsc_frame.create_view(&wgpu::TextureViewDescriptor::default());
124        let ntsc_pass_view = ntsc_pass.create_view(&wgpu::TextureViewDescriptor::default());
125        let ntsc_stop_view = ntsc_stop.create_view(&wgpu::TextureViewDescriptor::default());
126        let output_view = output_frame.create_view(&wgpu::TextureViewDescriptor::default());
127
128        let filters = NtscFilters::from(params.samples_per_color_cycle);
129
130        let luma_bsf_fir_buffer = create_fir_buffer(device, filters.luma_bsf);
131        let chroma_bpf_fir_buffer = create_fir_buffer(device, filters.chroma_bpf);
132        let y_encode_lpf_fir_buffer = create_fir_buffer(device, filters.y_encode_lpf);
133        let y_decode_lpf_fir_buffer = create_fir_buffer(device, filters.y_decode_lpf);
134        let uv_lpf_fir_buffer = create_fir_buffer(device, filters.uv_lpf);
135
136        let fir_len = match params.samples_per_color_cycle {
137            SamplesPerColorCycle::Fifteen => constants::FIR_LEN_15,
138            SamplesPerColorCycle::Twelve => constants::FIR_LEN_12,
139        };
140        let decode_hue_offset = match variant {
141            NtscShaderVariant::Rgb => 0.0,
142            NtscShaderVariant::NesPpu => 2.9 / 12.0 * 2.0 * std::f64::consts::PI,
143        };
144        let pipeline_compilation_options = wgpu::PipelineCompilationOptions {
145            constants: &[
146                ("samples_per_color_cycle", u32::from(params.samples_per_color_cycle).into()),
147                ("fir_len", fir_len.into()),
148                ("upscale_factor", params.upscale_factor.into()),
149                ("decode_hue_offset", decode_hue_offset),
150                ("decode_brightness", config.brightness),
151                ("decode_saturation", config.saturation),
152                ("decode_gamma", config.gamma),
153            ],
154            ..wgpu::PipelineCompilationOptions::default()
155        };
156
157        let initial_immediates_buffer = device.create_buffer(&wgpu::BufferDescriptor {
158            label: "ntsc_immediates_buffer".into(),
159            size: size_of::<ImmediateParams>() as u64,
160            usage: wgpu::BufferUsages::UNIFORM,
161            mapped_at_creation: false,
162        });
163
164        let immediates_bind_group_layout =
165            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
166                label: "ntsc_immediates_bind_group_layout".into(),
167                entries: &[wgpu::BindGroupLayoutEntry {
168                    binding: 0,
169                    visibility: wgpu::ShaderStages::COMPUTE,
170                    ty: wgpu::BindingType::Buffer {
171                        ty: wgpu::BufferBindingType::Uniform,
172                        has_dynamic_offset: false,
173                        min_binding_size: None,
174                    },
175                    count: None,
176                }],
177            });
178
179        let initial_immediates_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
180            label: "ntsc_immediates_bind_group".into(),
181            layout: &immediates_bind_group_layout,
182            entries: &[wgpu::BindGroupEntry {
183                binding: 0,
184                resource: wgpu::BindingResource::Buffer(
185                    initial_immediates_buffer.as_entire_buffer_binding(),
186                ),
187            }],
188        });
189
190        let rgb_to_ntsc_bind_group_layout =
191            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
192                label: "rgb_to_ntsc_bind_group_layout".into(),
193                entries: &[
194                    wgpu::BindGroupLayoutEntry {
195                        binding: 0,
196                        visibility: wgpu::ShaderStages::COMPUTE,
197                        ty: wgpu::BindingType::Buffer {
198                            ty: wgpu::BufferBindingType::Uniform,
199                            has_dynamic_offset: false,
200                            min_binding_size: None,
201                        },
202                        count: None,
203                    },
204                    wgpu::BindGroupLayoutEntry {
205                        binding: 1,
206                        visibility: wgpu::ShaderStages::COMPUTE,
207                        ty: wgpu::BindingType::Buffer {
208                            ty: wgpu::BufferBindingType::Uniform,
209                            has_dynamic_offset: false,
210                            min_binding_size: None,
211                        },
212                        count: None,
213                    },
214                    wgpu::BindGroupLayoutEntry {
215                        binding: 2,
216                        visibility: wgpu::ShaderStages::COMPUTE,
217                        ty: wgpu::BindingType::Texture {
218                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
219                            view_dimension: wgpu::TextureViewDimension::D2,
220                            multisampled: false,
221                        },
222                        count: None,
223                    },
224                    wgpu::BindGroupLayoutEntry {
225                        binding: 3,
226                        visibility: wgpu::ShaderStages::COMPUTE,
227                        ty: wgpu::BindingType::StorageTexture {
228                            access: wgpu::StorageTextureAccess::WriteOnly,
229                            format: wgpu::TextureFormat::R32Float,
230                            view_dimension: wgpu::TextureViewDimension::D2,
231                        },
232                        count: None,
233                    },
234                ],
235            });
236
237        let rgb_to_ntsc_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
238            label: "rgb_to_ntsc_bind_group".into(),
239            layout: &rgb_to_ntsc_bind_group_layout,
240            entries: &[
241                wgpu::BindGroupEntry {
242                    binding: 0,
243                    resource: wgpu::BindingResource::Buffer(
244                        y_encode_lpf_fir_buffer.as_entire_buffer_binding(),
245                    ),
246                },
247                wgpu::BindGroupEntry {
248                    binding: 1,
249                    resource: wgpu::BindingResource::Buffer(
250                        uv_lpf_fir_buffer.as_entire_buffer_binding(),
251                    ),
252                },
253                wgpu::BindGroupEntry {
254                    binding: 2,
255                    resource: wgpu::BindingResource::TextureView(&input_view),
256                },
257                wgpu::BindGroupEntry {
258                    binding: 3,
259                    resource: wgpu::BindingResource::TextureView(&ntsc_view),
260                },
261            ],
262        });
263
264        let rgb_to_ntsc_pipeline_layout =
265            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
266                label: "rgb_to_ntsc_pipeline_layout".into(),
267                bind_group_layouts: &[
268                    Some(&rgb_to_ntsc_bind_group_layout),
269                    Some(&immediates_bind_group_layout),
270                ],
271                immediate_size: 0,
272            });
273
274        let rgb_to_ntsc_shader = match variant {
275            NtscShaderVariant::Rgb => "rgb_to_ntsc",
276            NtscShaderVariant::NesPpu => "nes_to_ntsc",
277        };
278        let rgb_to_ntsc_pipeline =
279            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
280                label: "rgb_to_ntsc_pipeline".into(),
281                layout: Some(&rgb_to_ntsc_pipeline_layout),
282                module: &shaders.ntsc,
283                entry_point: Some(rgb_to_ntsc_shader),
284                compilation_options: pipeline_compilation_options.clone(),
285                cache: None,
286            });
287
288        let separate_luma_chroma_pipeline =
289            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
290                label: "separate_luma_chroma_pipeline".into(),
291                layout: None,
292                module: &shaders.ntsc,
293                entry_point: Some("separate_luma_chroma"),
294                compilation_options: pipeline_compilation_options.clone(),
295                cache: None,
296            });
297
298        let separate_luma_chroma_bind_group =
299            device.create_bind_group(&wgpu::BindGroupDescriptor {
300                label: "separate_luma_chroma_bind_group".into(),
301                layout: &separate_luma_chroma_pipeline.get_bind_group_layout(0),
302                entries: &[
303                    wgpu::BindGroupEntry {
304                        binding: 4,
305                        resource: wgpu::BindingResource::Buffer(
306                            luma_bsf_fir_buffer.as_entire_buffer_binding(),
307                        ),
308                    },
309                    wgpu::BindGroupEntry {
310                        binding: 5,
311                        resource: wgpu::BindingResource::Buffer(
312                            chroma_bpf_fir_buffer.as_entire_buffer_binding(),
313                        ),
314                    },
315                    wgpu::BindGroupEntry {
316                        binding: 6,
317                        resource: wgpu::BindingResource::TextureView(&ntsc_view),
318                    },
319                    wgpu::BindGroupEntry {
320                        binding: 7,
321                        resource: wgpu::BindingResource::TextureView(&ntsc_pass_view),
322                    },
323                    wgpu::BindGroupEntry {
324                        binding: 8,
325                        resource: wgpu::BindingResource::TextureView(&ntsc_stop_view),
326                    },
327                ],
328            });
329
330        let luma_chroma_to_rgb_bind_group_layout =
331            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
332                label: "luma_chroma_to_rgb_bind_group_layout".into(),
333                entries: &[
334                    wgpu::BindGroupLayoutEntry {
335                        binding: 9,
336                        visibility: wgpu::ShaderStages::COMPUTE,
337                        ty: wgpu::BindingType::Buffer {
338                            ty: wgpu::BufferBindingType::Uniform,
339                            has_dynamic_offset: false,
340                            min_binding_size: None,
341                        },
342                        count: None,
343                    },
344                    wgpu::BindGroupLayoutEntry {
345                        binding: 10,
346                        visibility: wgpu::ShaderStages::COMPUTE,
347                        ty: wgpu::BindingType::Buffer {
348                            ty: wgpu::BufferBindingType::Uniform,
349                            has_dynamic_offset: false,
350                            min_binding_size: None,
351                        },
352                        count: None,
353                    },
354                    wgpu::BindGroupLayoutEntry {
355                        binding: 11,
356                        visibility: wgpu::ShaderStages::COMPUTE,
357                        ty: wgpu::BindingType::Texture {
358                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
359                            view_dimension: wgpu::TextureViewDimension::D2,
360                            multisampled: false,
361                        },
362                        count: None,
363                    },
364                    wgpu::BindGroupLayoutEntry {
365                        binding: 12,
366                        visibility: wgpu::ShaderStages::COMPUTE,
367                        ty: wgpu::BindingType::Texture {
368                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
369                            view_dimension: wgpu::TextureViewDimension::D2,
370                            multisampled: false,
371                        },
372                        count: None,
373                    },
374                    wgpu::BindGroupLayoutEntry {
375                        binding: 13,
376                        visibility: wgpu::ShaderStages::COMPUTE,
377                        ty: wgpu::BindingType::StorageTexture {
378                            access: wgpu::StorageTextureAccess::WriteOnly,
379                            format: wgpu::TextureFormat::Rgba8Unorm,
380                            view_dimension: wgpu::TextureViewDimension::D2,
381                        },
382                        count: None,
383                    },
384                ],
385            });
386
387        let luma_chroma_to_rgb_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
388            label: "luma_chroma_to_rgb_bind_group".into(),
389            layout: &luma_chroma_to_rgb_bind_group_layout,
390            entries: &[
391                wgpu::BindGroupEntry {
392                    binding: 9,
393                    resource: wgpu::BindingResource::Buffer(
394                        y_decode_lpf_fir_buffer.as_entire_buffer_binding(),
395                    ),
396                },
397                wgpu::BindGroupEntry {
398                    binding: 10,
399                    resource: wgpu::BindingResource::Buffer(
400                        uv_lpf_fir_buffer.as_entire_buffer_binding(),
401                    ),
402                },
403                wgpu::BindGroupEntry {
404                    binding: 11,
405                    resource: wgpu::BindingResource::TextureView(&ntsc_pass_view),
406                },
407                wgpu::BindGroupEntry {
408                    binding: 12,
409                    resource: wgpu::BindingResource::TextureView(&ntsc_stop_view),
410                },
411                wgpu::BindGroupEntry {
412                    binding: 13,
413                    resource: wgpu::BindingResource::TextureView(&output_view),
414                },
415            ],
416        });
417
418        let luma_chroma_to_rgb_pipeline_layout =
419            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
420                label: "luma_chroma_to_rgb_pipeline_layout".into(),
421                bind_group_layouts: &[
422                    Some(&luma_chroma_to_rgb_bind_group_layout),
423                    Some(&immediates_bind_group_layout),
424                ],
425                immediate_size: 0,
426            });
427
428        let luma_chroma_to_rgb_pipeline =
429            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
430                label: "luma_chroma_to_rgb_pipeline".into(),
431                layout: Some(&luma_chroma_to_rgb_pipeline_layout),
432                module: &shaders.ntsc,
433                entry_point: Some("luma_chroma_to_rgb"),
434                compilation_options: pipeline_compilation_options.clone(),
435                cache: None,
436            });
437
438        Self {
439            output: Arc::new(output_frame),
440            ntsc_frame_size: ntsc_texture_descriptor.size,
441            samples_per_color_cycle: params.samples_per_color_cycle,
442            immediates_bind_group_layout,
443            immediates_bind_group: initial_immediates_bind_group,
444            rgb_to_ntsc_bind_group,
445            rgb_to_ntsc_pipeline,
446            separate_luma_chroma_bind_group,
447            separate_luma_chroma_pipeline,
448            luma_chroma_to_rgb_bind_group,
449            luma_chroma_to_rgb_pipeline,
450        }
451    }
452}
453
454impl PipelineShader for NtscShader {
455    fn prepare(&mut self, device: &wgpu::Device, options: RenderFrameOptions) {
456        let immediate_params =
457            options.ntsc_per_frame_params.map_or(ImmediateParams::ZERO, |per_frame_params| {
458                let samples_per_color_cycle: u64 = self.samples_per_color_cycle.into();
459
460                ImmediateParams {
461                    frame_phase_offset: (per_frame_params.frame_phase_offset
462                        % samples_per_color_cycle) as i32,
463                    per_line_phase_offset: (per_frame_params.per_line_phase_offset
464                        % samples_per_color_cycle)
465                        as i32,
466                }
467            });
468
469        let immediates_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
470            label: "ntsc_immediates_buffer".into(),
471            contents: bytemuck::cast_slice(&[immediate_params]),
472            usage: wgpu::BufferUsages::UNIFORM,
473        });
474
475        self.immediates_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
476            label: "ntsc_immediates_bind_group".into(),
477            layout: &self.immediates_bind_group_layout,
478            entries: &[wgpu::BindGroupEntry {
479                binding: 0,
480                resource: wgpu::BindingResource::Buffer(
481                    immediates_buffer.as_entire_buffer_binding(),
482                ),
483            }],
484        });
485    }
486
487    fn draw(&mut self, encoder: &mut wgpu::CommandEncoder) {
488        let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
489            label: "ntsc_compute_pass".into(),
490            ..wgpu::ComputePassDescriptor::default()
491        });
492
493        let ntsc_x_workgroups = self.ntsc_frame_size.width / 16
494            + u32::from(!self.ntsc_frame_size.width.is_multiple_of(16));
495        let output_x_workgroups =
496            self.output.width() / 16 + u32::from(!self.output.width().is_multiple_of(16));
497        let y_workgroups =
498            self.output.height() / 16 + u32::from(!self.output.height().is_multiple_of(16));
499
500        compute_pass.set_bind_group(1, &self.immediates_bind_group, &[]);
501
502        compute_pass.set_bind_group(0, &self.rgb_to_ntsc_bind_group, &[]);
503        compute_pass.set_pipeline(&self.rgb_to_ntsc_pipeline);
504        compute_pass.dispatch_workgroups(ntsc_x_workgroups, y_workgroups, 1);
505
506        compute_pass.set_bind_group(0, &self.separate_luma_chroma_bind_group, &[]);
507        compute_pass.set_pipeline(&self.separate_luma_chroma_pipeline);
508        compute_pass.dispatch_workgroups(ntsc_x_workgroups, y_workgroups, 1);
509
510        compute_pass.set_bind_group(0, &self.luma_chroma_to_rgb_bind_group, &[]);
511        compute_pass.set_pipeline(&self.luma_chroma_to_rgb_pipeline);
512        compute_pass.dispatch_workgroups(output_x_workgroups, y_workgroups, 1);
513    }
514
515    fn output_texture(&self) -> &Arc<wgpu::Texture> {
516        &self.output
517    }
518}
519
520fn create_fir_buffer(device: &wgpu::Device, coefficients: &[f32]) -> wgpu::Buffer {
521    let mut fir: Vec<f32> = vec![0.0; FIR_BUFFER_LEN];
522
523    let slice_len = coefficients.len().min(FIR_BUFFER_LEN);
524    fir[..slice_len].copy_from_slice(&coefficients[..slice_len]);
525
526    device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
527        label: None,
528        contents: bytemuck::cast_slice(&fir),
529        usage: wgpu::BufferUsages::UNIFORM,
530    })
531}