fix(encoder): only convert RGB input where the driver sites chroma as the stream implies

shader/color_convert.comp

@@ -169,6 +169,23 @@ uint q10_c(float c) {
     return (val << 6u) & 0xFFC0u;
 }
 
+// 4:2:0 chroma for the block whose top-left sample is (x, y), sited where a
+// player puts it when the stream says nothing else: on the left luma column,
+// midway between the two rows (chroma_sample_loc_type 0). That is a [1 2 1]
+// filter centred on column x, averaged over both rows. A sample past an edge
+// repeats the edge. `yuv` is the sample at (x, y), already converted.
+vec2 chroma_420(uint x, uint y, vec3 yuv) {
+    uint xl = x > 0u ? x - 1u : x;
+    uint xr = min(x + 1u, params.width - 1u);
+    uint yb = min(y + 1u, params.height - 1u);
+    vec2 top = rgb_to_yuv(read_rgb(ivec2(xl, y))).yz + 2.0 * yuv.yz
+        + rgb_to_yuv(read_rgb(ivec2(xr, y))).yz;
+    vec2 bottom = rgb_to_yuv(read_rgb(ivec2(xl, yb))).yz
+        + 2.0 * rgb_to_yuv(read_rgb(ivec2(x, yb))).yz
+        + rgb_to_yuv(read_rgb(ivec2(xr, yb))).yz;
+    return (top + bottom) / 8.0;
+}
+
 void main() {
     uint x = gl_GlobalInvocationID.x;
     uint y = gl_GlobalInvocationID.y;
@@ -222,16 +239,9 @@ void main() {
             uint uv_width = params.width / 2u;
             uint uv_idx = uv_y * uv_width + uv_x;
 
-            vec3 yuv00 = yuv;
-            vec3 yuv10 = (x + 1u < params.width) ?
-                rgb_to_yuv(read_rgb(ivec2(x + 1u, y))) : yuv00;
-            vec3 yuv01 = (y + 1u < params.height) ?
-                rgb_to_yuv(read_rgb(ivec2(x, y + 1u))) : yuv00;
-            vec3 yuv11 = (x + 1u < params.width && y + 1u < params.height) ?
-                rgb_to_yuv(read_rgb(ivec2(x + 1u, y + 1u))) : yuv00;
-
-            float avg_u = (yuv00.y + yuv10.y + yuv01.y + yuv11.y) / 4.0;
-            float avg_v = (yuv00.z + yuv10.z + yuv01.z + yuv11.z) / 4.0;
+            vec2 uv = chroma_420(x, y, yuv);
+            float avg_u = uv.x;
+            float avg_v = uv.y;
 
             uint uv_base_words = pixel_count / 2u;
             uint uv_word_idx = uv_base_words + uv_idx;
@@ -251,16 +261,9 @@ void main() {
             uint uv_width = params.width / 2u;
             uint uv_idx = uv_y * uv_width + uv_x;
 
-            vec3 yuv00 = yuv;
-            vec3 yuv10 = (x + 1u < params.width) ?
-                rgb_to_yuv(read_rgb(ivec2(x + 1u, y))) : yuv00;
-            vec3 yuv01 = (y + 1u < params.height) ?
-                rgb_to_yuv(read_rgb(ivec2(x, y + 1u))) : yuv00;
-            vec3 yuv11 = (x + 1u < params.width && y + 1u < params.height) ?
-                rgb_to_yuv(read_rgb(ivec2(x + 1u, y + 1u))) : yuv00;
-
-            float avg_u = (yuv00.y + yuv10.y + yuv01.y + yuv11.y) / 4.0;
-            float avg_v = (yuv00.z + yuv10.z + yuv01.z + yuv11.z) / 4.0;
+            vec2 uv = chroma_420(x, y, yuv);
+            float avg_u = uv.x;
+            float avg_v = uv.y;
 
             if (params.output_format == 0u) {
                 // NV12: Interleaved UV after Y plane.

src/encoder/av1/session_params.rs

@@ -69,7 +69,13 @@ impl Av1 {
             color_primaries: av1_color_primaries,
             transfer_characteristics: av1_transfer,
             matrix_coefficients: av1_matrix,
-            chroma_sample_position: ash::vk::native::StdVideoAV1ChromaSamplePosition_STD_VIDEO_AV1_CHROMA_SAMPLE_POSITION_UNKNOWN,
+            // 4:2:0 chroma sits on the left luma sample, midway between the
+            // two rows. Only 4:2:0 writes a position at all.
+            chroma_sample_position: if subsampling_x == 1 && subsampling_y == 1 {
+                ash::vk::native::StdVideoAV1ChromaSamplePosition_STD_VIDEO_AV1_CHROMA_SAMPLE_POSITION_VERTICAL
+            } else {
+                ash::vk::native::StdVideoAV1ChromaSamplePosition_STD_VIDEO_AV1_CHROMA_SAMPLE_POSITION_UNKNOWN
+            },
         };
 
         let profile = match config.pixel_format {

src/encoder/h264/session_params.rs

@@ -75,6 +75,12 @@ impl H264 {
         vui_flags.set_video_signal_type_present_flag(1);
         vui_flags.set_video_full_range_flag(if desc.full_range { 1 } else { 0 });
         vui_flags.set_color_description_present_flag(1);
+        // 4:2:0 chroma sits on the left luma sample, midway between the two
+        // rows (chroma_sample_loc_type 0), and the stream says so.
+        vui_flags
+            .set_chroma_loc_info_present_flag(
+                matches!(config.pixel_format, PixelFormat::Yuv420) as u32
+            );
         // Avoid HRD with zeroed bitrates (device loss on some drivers).
         vui_flags.set_nal_hrd_parameters_present_flag(0);
         vui_flags.set_bitstream_restriction_flag(1);

src/encoder/h265/session_params.rs

@@ -128,7 +128,9 @@ impl H265 {
                 1,
                 if desc.full_range { 1 } else { 0 },
                 1,
-                0,
+                // chroma_loc_info_present_flag: 4:2:0 chroma sits on the left
+                // luma sample, midway between the two rows (type 0).
+                matches!(config.pixel_format, PixelFormat::Yuv420) as u32,
                 0,
                 0,
                 0,

src/encoder/rgb.rs

@@ -129,23 +129,22 @@ impl RgbConversionCaps {
         // and a player that is not told assumes chroma sits on the left luma
         // sample of each pair, and midway between the two rows (H.264 and
         // H.265 chroma location type 0, which most AV1 players follow as well).
-        // Asking the hardware for exactly that keeps colour edges where the
-        // player will put them. A driver that cannot gets the other option.
-        let pick = |supported: Offset, preferred: Offset| {
-            if supported.contains(preferred) {
-                preferred
-            } else if supported.contains(Offset::COSITED_EVEN) {
-                Offset::COSITED_EVEN
-            } else {
-                Offset::MIDPOINT
-            }
-        };
+        // The hardware has to produce exactly that: a driver that can only
+        // site chroma elsewhere would shift color edges without the stream
+        // saying so, so it gets no RGB input and the caller converts instead.
+        let (x, y) = (Offset::COSITED_EVEN, Offset::MIDPOINT);
+        if !self.x_chroma_offsets.contains(x) || !self.y_chroma_offsets.contains(y) {
+            return Err(PixelForgeError::NoSuitableDevice(format!(
+                "RGB input: the driver sites chroma at {:?} x {:?}, not {x:?} x {y:?}, which the stream implies",
+                self.x_chroma_offsets, self.y_chroma_offsets
+            )));
+        }
         Ok(
             vk::VideoEncodeSessionRgbConversionCreateInfoVALVE::default()
                 .rgb_model(model)
                 .rgb_range(range)
-                .x_chroma_offset(pick(self.x_chroma_offsets, Offset::COSITED_EVEN))
-                .y_chroma_offset(pick(self.y_chroma_offsets, Offset::MIDPOINT)),
+                .x_chroma_offset(x)
+                .y_chroma_offset(y),
         )
     }
 }
@@ -221,3 +220,63 @@ impl VideoContext {
         true
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use vk::VideoEncodeRgbChromaOffsetFlagsVALVE as Offset;
+    use vk::VideoEncodeRgbModelConversionFlagsVALVE as Model;
+    use vk::VideoEncodeRgbRangeCompressionFlagsVALVE as Range;
+
+    /// What RADV reports (Mesa radv_video.c): chroma cosited horizontally,
+    /// either way vertically.
+    fn radv() -> RgbConversionCaps {
+        RgbConversionCaps {
+            models: Model::YCBCR_709 | Model::YCBCR_2020,
+            ranges: Range::FULL_RANGE | Range::NARROW_RANGE,
+            x_chroma_offsets: Offset::COSITED_EVEN,
+            y_chroma_offsets: Offset::MIDPOINT | Offset::COSITED_EVEN,
+        }
+    }
+
+    #[test]
+    fn sites_chroma_where_the_stream_implies() {
+        let info = radv().session_info(&ColorDescription::bt709()).unwrap();
+        assert_eq!(info.x_chroma_offset, Offset::COSITED_EVEN);
+        assert_eq!(info.y_chroma_offset, Offset::MIDPOINT);
+        assert_eq!(info.rgb_model, Model::YCBCR_709);
+        assert_eq!(info.rgb_range, Range::NARROW_RANGE);
+    }
+
+    #[test]
+    fn refuses_other_chroma_siting() {
+        let midpoint_x = RgbConversionCaps {
+            x_chroma_offsets: Offset::MIDPOINT,
+            ..radv()
+        };
+        assert!(midpoint_x.session_info(&ColorDescription::bt709()).is_err());
+        let cosited_y = RgbConversionCaps {
+            y_chroma_offsets: Offset::COSITED_EVEN,
+            ..radv()
+        };
+        assert!(cosited_y.session_info(&ColorDescription::bt709()).is_err());
+    }
+
+    #[test]
+    fn refuses_an_unsupported_matrix_or_range() {
+        let no_2020 = RgbConversionCaps {
+            models: Model::YCBCR_709,
+            ..radv()
+        };
+        assert!(
+            no_2020
+                .session_info(&ColorDescription::bt2020_pq())
+                .is_err()
+        );
+        let full_only = RgbConversionCaps {
+            ranges: Range::FULL_RANGE,
+            ..radv()
+        };
+        assert!(full_only.session_info(&ColorDescription::bt709()).is_err());
+    }
+}

tests/color_conversion.rs

@@ -14,7 +14,7 @@
 //! target rather than from the shader's branches, so a wrong branch shows up as
 //! a large disagreement.
 //!
-//! Luma only, since chroma is averaged over 2x2 blocks and that averaging would
+//! Luma only, since chroma is filtered over neighbouring pixels and that filter would
 //! muddy the comparison. The test frame is deliberately not grey: grey converts
 //! to exact values and would hide the very errors being looked for.
 //!