diff --git a/crates/rendering/src/shaders/yuv420p_to_rgba.wgsl b/crates/rendering/src/shaders/yuv420p_to_rgba.wgsl index c95d3ae38d9..93ad874ed37 100644 --- a/crates/rendering/src/shaders/yuv420p_to_rgba.wgsl +++ b/crates/rendering/src/shaders/yuv420p_to_rgba.wgsl @@ -15,20 +15,41 @@ override BU: f32 = 2.017232; fn main(@builtin(global_invocation_id) global_id: vec3) { let coords = global_id.xy; let dims = textureDimensions(output); + let uv = load_uv(coords / 2u); + convert_pixel(coords, dims, uv.x, uv.y); +} - if (coords.x >= dims.x || coords.y >= dims.y) { - return; - } +// One invocation per 2x2 block: the block shares its chroma sample, so this loads U and V once instead of four times. +@compute @workgroup_size(8, 8) +fn main_block(@builtin(global_invocation_id) global_id: vec3) { + let uv_coords = global_id.xy; + let dims = textureDimensions(output); + let uv = load_uv(uv_coords); - let y_raw = textureLoad(y_plane, coords, 0).r; + let origin = uv_coords * 2u; + convert_pixel(origin, dims, uv.x, uv.y); + convert_pixel(origin + vec2(1u, 0u), dims, uv.x, uv.y); + convert_pixel(origin + vec2(0u, 1u), dims, uv.x, uv.y); + convert_pixel(origin + vec2(1u, 1u), dims, uv.x, uv.y); +} - let uv_coords = coords / 2; +fn load_uv(uv_coords: vec2) -> vec2 { let u_dims = textureDimensions(u_plane); let v_dims = textureDimensions(v_plane); let u_clamped = min(uv_coords, u_dims - vec2(1, 1)); let v_clamped = min(uv_coords, v_dims - vec2(1, 1)); - let u_raw = textureLoad(u_plane, u_clamped, 0).r; - let v_raw = textureLoad(v_plane, v_clamped, 0).r; + return vec2( + textureLoad(u_plane, u_clamped, 0).r, + textureLoad(v_plane, v_clamped, 0).r + ); +} + +fn convert_pixel(coords: vec2, dims: vec2, u_raw: f32, v_raw: f32) { + if (coords.x >= dims.x || coords.y >= dims.y) { + return; + } + + let y_raw = textureLoad(y_plane, coords, 0).r; let y = (y_raw - 16.0 / 255.0) * (255.0 / 219.0); let u = u_raw - 128.0 / 255.0; diff --git a/crates/rendering/src/yuv_converter.rs b/crates/rendering/src/yuv_converter.rs index 2dc25523e6a..7f31d0ab72d 100644 --- a/crates/rendering/src/yuv_converter.rs +++ b/crates/rendering/src/yuv_converter.rs @@ -107,6 +107,16 @@ const INITIAL_TEXTURE_HEIGHT: u32 = 1080; const TEXTURE_SIZE_PADDING: u32 = 64; +// The 2x2 block kernel was only measured as a win on Linux; Apple GPUs were slower and Windows is untested. +#[cfg(target_os = "linux")] +const YUV420P_ENTRY_POINT: &str = "main_block"; +#[cfg(target_os = "linux")] +const YUV420P_PIXELS_PER_WORKGROUP: u32 = 16; +#[cfg(not(target_os = "linux"))] +const YUV420P_ENTRY_POINT: &str = "main"; +#[cfg(not(target_os = "linux"))] +const YUV420P_PIXELS_PER_WORKGROUP: u32 = 8; + fn align_dimension(dim: u32) -> u32 { dim.div_ceil(TEXTURE_SIZE_PADDING) * TEXTURE_SIZE_PADDING } @@ -266,6 +276,7 @@ struct MatrixPipelines { label: &'static str, layout: wgpu::PipelineLayout, module: wgpu::ShaderModule, + entry_point: &'static str, bt601: std::sync::OnceLock, bt709: std::sync::OnceLock, } @@ -276,12 +287,14 @@ impl MatrixPipelines { label: &'static str, layout: wgpu::PipelineLayout, module: wgpu::ShaderModule, + entry_point: &'static str, ) -> Self { Self { device: device.clone(), label, layout, module, + entry_point, bt601: std::sync::OnceLock::new(), bt709: std::sync::OnceLock::new(), } @@ -299,7 +312,7 @@ impl MatrixPipelines { label: Some(self.label), layout: Some(&self.layout), module: &self.module, - entry_point: Some("main"), + entry_point: Some(self.entry_point), compilation_options: wgpu::PipelineCompilationOptions { constants: &constants, ..Default::default() @@ -437,12 +450,14 @@ impl YuvConverterPipelines { "NV12 Converter Pipeline", nv12_pipeline_layout, nv12_shader, + "main", ); let yuv420p = MatrixPipelines::new( device, "YUV420P Converter Pipeline", yuv420p_pipeline_layout, yuv420p_shader, + YUV420P_ENTRY_POINT, ); tracing::info!("Shared YUV converter pipelines created successfully"); @@ -1258,7 +1273,11 @@ impl YuvToRgbaConverter { }); compute_pass.set_pipeline(self.pipelines.yuv420p.get(self.matrix)); compute_pass.set_bind_group(0, bind_group, &[]); - compute_pass.dispatch_workgroups(width.div_ceil(8), height.div_ceil(8), 1); + compute_pass.dispatch_workgroups( + width.div_ceil(YUV420P_PIXELS_PER_WORKGROUP), + height.div_ceil(YUV420P_PIXELS_PER_WORKGROUP), + 1, + ); } queue.submit(std::iter::once(encoder.finish())); @@ -1329,7 +1348,11 @@ impl YuvToRgbaConverter { }); compute_pass.set_pipeline(self.pipelines.yuv420p.get(self.matrix)); compute_pass.set_bind_group(0, bind_group, &[]); - compute_pass.dispatch_workgroups(width.div_ceil(8), height.div_ceil(8), 1); + compute_pass.dispatch_workgroups( + width.div_ceil(YUV420P_PIXELS_PER_WORKGROUP), + height.div_ceil(YUV420P_PIXELS_PER_WORKGROUP), + 1, + ); } Ok(self.current_output_view())