Render surfaces correctly when encoded in `420YpCbCr8BiPlanarFullRange`

Antonio Scandurra created

Change summary

crates/gpui/src/platform/mac/renderer.rs           |  75 +++--
crates/gpui/src/platform/mac/shaders/shaders.h     | 185 ++++++++-------
crates/gpui/src/platform/mac/shaders/shaders.metal |  48 ++++
crates/media/src/media.rs                          |  15 +
4 files changed, 202 insertions(+), 121 deletions(-)

Detailed changes

crates/gpui/src/platform/mac/renderer.rs 🔗

@@ -28,6 +28,7 @@ pub struct Renderer {
     shadow_pipeline_state: metal::RenderPipelineState,
     sprite_pipeline_state: metal::RenderPipelineState,
     image_pipeline_state: metal::RenderPipelineState,
+    surface_pipeline_state: metal::RenderPipelineState,
     path_atlas_pipeline_state: metal::RenderPipelineState,
     underline_pipeline_state: metal::RenderPipelineState,
     unit_vertices: metal::Buffer,
@@ -116,6 +117,14 @@ impl Renderer {
             "image_fragment",
             pixel_format,
         );
+        let surface_pipeline_state = build_pipeline_state(
+            &device,
+            &library,
+            "surface",
+            "surface_vertex",
+            "surface_fragment",
+            pixel_format,
+        );
         let path_atlas_pipeline_state = build_path_atlas_pipeline_state(
             &device,
             &library,
@@ -141,6 +150,7 @@ impl Renderer {
             shadow_pipeline_state,
             sprite_pipeline_state,
             image_pipeline_state,
+            surface_pipeline_state,
             path_atlas_pipeline_state,
             underline_pipeline_state,
             unit_vertices,
@@ -798,14 +808,14 @@ impl Renderer {
             return;
         }
 
-        command_encoder.set_render_pipeline_state(&self.image_pipeline_state);
+        command_encoder.set_render_pipeline_state(&self.surface_pipeline_state);
         command_encoder.set_vertex_buffer(
-            shaders::GPUIImageVertexInputIndex_GPUIImageVertexInputIndexVertices as u64,
+            shaders::GPUISurfaceVertexInputIndex_GPUISurfaceVertexInputIndexVertices as u64,
             Some(&self.unit_vertices),
             0,
         );
         command_encoder.set_vertex_bytes(
-            shaders::GPUIImageVertexInputIndex_GPUIImageVertexInputIndexViewportSize as u64,
+            shaders::GPUISurfaceVertexInputIndex_GPUISurfaceVertexInputIndexViewportSize as u64,
             mem::size_of::<shaders::vector_float2>() as u64,
             [drawable_size.to_float2()].as_ptr() as *const c_void,
         );
@@ -817,64 +827,71 @@ impl Renderer {
                 surface.image_buffer.height() as i32,
             );
             let target_size = surface.bounds.size() * scale_factor;
-            let pixel_format = if surface.image_buffer.pixel_format_type()
-                == core_video::kCVPixelFormatType_32BGRA
-            {
-                MTLPixelFormat::BGRA8Unorm
-            } else {
-                MTLPixelFormat::R8Unorm
-            };
 
-            let texture = self
+            assert_eq!(
+                surface.image_buffer.pixel_format_type(),
+                core_video::kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
+            );
+
+            let y_texture = self
                 .cv_texture_cache
                 .create_texture_from_image(
                     surface.image_buffer.as_concrete_TypeRef(),
                     ptr::null(),
-                    pixel_format,
-                    source_size.x() as usize,
-                    source_size.y() as usize,
+                    MTLPixelFormat::R8Unorm,
+                    surface.image_buffer.plane_width(0),
+                    surface.image_buffer.plane_height(0),
                     0,
                 )
                 .unwrap();
+            let cb_cr_texture = self
+                .cv_texture_cache
+                .create_texture_from_image(
+                    surface.image_buffer.as_concrete_TypeRef(),
+                    ptr::null(),
+                    MTLPixelFormat::RG8Unorm,
+                    surface.image_buffer.plane_width(1),
+                    surface.image_buffer.plane_height(1),
+                    1,
+                )
+                .unwrap();
 
             align_offset(offset);
-            let next_offset = *offset + mem::size_of::<shaders::GPUIImage>();
+            let next_offset = *offset + mem::size_of::<shaders::GPUISurface>();
             assert!(
                 next_offset <= INSTANCE_BUFFER_SIZE,
                 "instance buffer exhausted"
             );
 
             command_encoder.set_vertex_buffer(
-                shaders::GPUIImageVertexInputIndex_GPUIImageVertexInputIndexImages as u64,
+                shaders::GPUISurfaceVertexInputIndex_GPUISurfaceVertexInputIndexSurfaces as u64,
                 Some(&self.instances),
                 *offset as u64,
             );
             command_encoder.set_vertex_bytes(
-                shaders::GPUIImageVertexInputIndex_GPUIImageVertexInputIndexAtlasSize as u64,
+                shaders::GPUISurfaceVertexInputIndex_GPUISurfaceVertexInputIndexAtlasSize as u64,
                 mem::size_of::<shaders::vector_float2>() as u64,
                 [source_size.to_float2()].as_ptr() as *const c_void,
             );
             command_encoder.set_fragment_texture(
-                shaders::GPUIImageFragmentInputIndex_GPUIImageFragmentInputIndexAtlas as u64,
-                Some(texture.as_texture_ref()),
+                shaders::GPUISurfaceFragmentInputIndex_GPUISurfaceFragmentInputIndexYAtlas as u64,
+                Some(y_texture.as_texture_ref()),
+            );
+            command_encoder.set_fragment_texture(
+                shaders::GPUISurfaceFragmentInputIndex_GPUISurfaceFragmentInputIndexCbCrAtlas
+                    as u64,
+                Some(cb_cr_texture.as_texture_ref()),
             );
 
             unsafe {
-                let buffer_contents =
-                    (self.instances.contents() as *mut u8).add(*offset) as *mut shaders::GPUIImage;
+                let buffer_contents = (self.instances.contents() as *mut u8).add(*offset)
+                    as *mut shaders::GPUISurface;
                 std::ptr::write(
                     buffer_contents,
-                    shaders::GPUIImage {
+                    shaders::GPUISurface {
                         origin: origin.to_float2(),
                         target_size: target_size.to_float2(),
                         source_size: source_size.to_float2(),
-                        atlas_origin: Default::default(),
-                        border_top: Default::default(),
-                        border_right: Default::default(),
-                        border_bottom: Default::default(),
-                        border_left: Default::default(),
-                        border_color: Default::default(),
-                        corner_radius: Default::default(),
                     },
                 );
             }

crates/gpui/src/platform/mac/shaders/shaders.h 🔗

@@ -1,122 +1,125 @@
 #include <simd/simd.h>
 
-typedef struct
-{
-    vector_float2 viewport_size;
+typedef struct {
+  vector_float2 viewport_size;
 } GPUIUniforms;
 
-typedef enum
-{
-    GPUIQuadInputIndexVertices = 0,
-    GPUIQuadInputIndexQuads = 1,
-    GPUIQuadInputIndexUniforms = 2,
+typedef enum {
+  GPUIQuadInputIndexVertices = 0,
+  GPUIQuadInputIndexQuads = 1,
+  GPUIQuadInputIndexUniforms = 2,
 } GPUIQuadInputIndex;
 
-typedef struct
-{
-    vector_float2 origin;
-    vector_float2 size;
-    vector_uchar4 background_color;
-    float border_top;
-    float border_right;
-    float border_bottom;
-    float border_left;
-    vector_uchar4 border_color;
-    float corner_radius;
+typedef struct {
+  vector_float2 origin;
+  vector_float2 size;
+  vector_uchar4 background_color;
+  float border_top;
+  float border_right;
+  float border_bottom;
+  float border_left;
+  vector_uchar4 border_color;
+  float corner_radius;
 } GPUIQuad;
 
-typedef enum
-{
-    GPUIShadowInputIndexVertices = 0,
-    GPUIShadowInputIndexShadows = 1,
-    GPUIShadowInputIndexUniforms = 2,
+typedef enum {
+  GPUIShadowInputIndexVertices = 0,
+  GPUIShadowInputIndexShadows = 1,
+  GPUIShadowInputIndexUniforms = 2,
 } GPUIShadowInputIndex;
 
-typedef struct
-{
-    vector_float2 origin;
-    vector_float2 size;
-    float corner_radius;
-    float sigma;
-    vector_uchar4 color;
+typedef struct {
+  vector_float2 origin;
+  vector_float2 size;
+  float corner_radius;
+  float sigma;
+  vector_uchar4 color;
 } GPUIShadow;
 
-typedef enum
-{
-    GPUISpriteVertexInputIndexVertices = 0,
-    GPUISpriteVertexInputIndexSprites = 1,
-    GPUISpriteVertexInputIndexViewportSize = 2,
-    GPUISpriteVertexInputIndexAtlasSize = 3,
+typedef enum {
+  GPUISpriteVertexInputIndexVertices = 0,
+  GPUISpriteVertexInputIndexSprites = 1,
+  GPUISpriteVertexInputIndexViewportSize = 2,
+  GPUISpriteVertexInputIndexAtlasSize = 3,
 } GPUISpriteVertexInputIndex;
 
-typedef enum
-{
-    GPUISpriteFragmentInputIndexAtlas = 0,
+typedef enum {
+  GPUISpriteFragmentInputIndexAtlas = 0,
 } GPUISpriteFragmentInputIndex;
 
-typedef struct
-{
-    vector_float2 origin;
-    vector_float2 target_size;
-    vector_float2 source_size;
-    vector_float2 atlas_origin;
-    vector_uchar4 color;
-    uint8_t compute_winding;
+typedef struct {
+  vector_float2 origin;
+  vector_float2 target_size;
+  vector_float2 source_size;
+  vector_float2 atlas_origin;
+  vector_uchar4 color;
+  uint8_t compute_winding;
 } GPUISprite;
 
-typedef enum
-{
-    GPUIPathAtlasVertexInputIndexVertices = 0,
-    GPUIPathAtlasVertexInputIndexAtlasSize = 1,
+typedef enum {
+  GPUIPathAtlasVertexInputIndexVertices = 0,
+  GPUIPathAtlasVertexInputIndexAtlasSize = 1,
 } GPUIPathAtlasVertexInputIndex;
 
-typedef struct
-{
-    vector_float2 xy_position;
-    vector_float2 st_position;
-    vector_float2 clip_rect_origin;
-    vector_float2 clip_rect_size;
+typedef struct {
+  vector_float2 xy_position;
+  vector_float2 st_position;
+  vector_float2 clip_rect_origin;
+  vector_float2 clip_rect_size;
 } GPUIPathVertex;
 
-typedef enum
-{
-    GPUIImageVertexInputIndexVertices = 0,
-    GPUIImageVertexInputIndexImages = 1,
-    GPUIImageVertexInputIndexViewportSize = 2,
-    GPUIImageVertexInputIndexAtlasSize = 3,
+typedef enum {
+  GPUIImageVertexInputIndexVertices = 0,
+  GPUIImageVertexInputIndexImages = 1,
+  GPUIImageVertexInputIndexViewportSize = 2,
+  GPUIImageVertexInputIndexAtlasSize = 3,
 } GPUIImageVertexInputIndex;
 
-typedef enum
-{
-    GPUIImageFragmentInputIndexAtlas = 0,
+typedef enum {
+  GPUIImageFragmentInputIndexAtlas = 0,
 } GPUIImageFragmentInputIndex;
 
-typedef struct
-{
-    vector_float2 origin;
-    vector_float2 target_size;
-    vector_float2 source_size;
-    vector_float2 atlas_origin;
-    float border_top;
-    float border_right;
-    float border_bottom;
-    float border_left;
-    vector_uchar4 border_color;
-    float corner_radius;
+typedef struct {
+  vector_float2 origin;
+  vector_float2 target_size;
+  vector_float2 source_size;
+  vector_float2 atlas_origin;
+  float border_top;
+  float border_right;
+  float border_bottom;
+  float border_left;
+  vector_uchar4 border_color;
+  float corner_radius;
 } GPUIImage;
 
-typedef enum
-{
-    GPUIUnderlineInputIndexVertices = 0,
-    GPUIUnderlineInputIndexUnderlines = 1,
-    GPUIUnderlineInputIndexUniforms = 2,
+typedef enum {
+  GPUISurfaceVertexInputIndexVertices = 0,
+  GPUISurfaceVertexInputIndexSurfaces = 1,
+  GPUISurfaceVertexInputIndexViewportSize = 2,
+  GPUISurfaceVertexInputIndexAtlasSize = 3,
+} GPUISurfaceVertexInputIndex;
+
+typedef enum {
+  GPUISurfaceFragmentInputIndexYAtlas = 0,
+  GPUISurfaceFragmentInputIndexCbCrAtlas = 1,
+} GPUISurfaceFragmentInputIndex;
+
+typedef struct {
+  vector_float2 origin;
+  vector_float2 target_size;
+  vector_float2 source_size;
+} GPUISurface;
+
+typedef enum {
+  GPUIUnderlineInputIndexVertices = 0,
+  GPUIUnderlineInputIndexUnderlines = 1,
+  GPUIUnderlineInputIndexUniforms = 2,
 } GPUIUnderlineInputIndex;
 
-typedef struct
-{
-    vector_float2 origin;
-    vector_float2 size;
-    float thickness;
-    vector_uchar4 color;
-    uint8_t squiggly;
+typedef struct {
+  vector_float2 origin;
+  vector_float2 size;
+  float thickness;
+  vector_uchar4 color;
+  uint8_t squiggly;
 } GPUIUnderline;

crates/gpui/src/platform/mac/shaders/shaders.metal 🔗

@@ -263,6 +263,54 @@ fragment float4 image_fragment(
     return quad_sdf(input);
 }
 
+vertex QuadFragmentInput surface_vertex(
+    uint unit_vertex_id [[vertex_id]],
+    uint image_id [[instance_id]],
+    constant float2 *unit_vertices [[buffer(GPUISurfaceVertexInputIndexVertices)]],
+    constant GPUISurface *images [[buffer(GPUISurfaceVertexInputIndexSurfaces)]],
+    constant float2 *viewport_size [[buffer(GPUISurfaceVertexInputIndexViewportSize)]],
+    constant float2 *atlas_size [[buffer(GPUISurfaceVertexInputIndexAtlasSize)]]
+) {
+    float2 unit_vertex = unit_vertices[unit_vertex_id];
+    GPUISurface image = images[image_id];
+    float2 position = unit_vertex * image.target_size + image.origin;
+    float4 device_position = to_device_position(position, *viewport_size);
+    float2 atlas_position = (unit_vertex * image.source_size) / *atlas_size;
+
+    return QuadFragmentInput {
+        device_position,
+        atlas_position,
+        image.origin,
+        image.target_size,
+        float4(0.),
+        0.,
+        0.,
+        0.,
+        0.,
+        float4(0.),
+        0.,
+    };
+}
+
+fragment float4 surface_fragment(
+    QuadFragmentInput input [[stage_in]],
+    texture2d<float> y_atlas [[ texture(GPUISurfaceFragmentInputIndexYAtlas) ]],
+    texture2d<float> cb_cr_atlas [[ texture(GPUISurfaceFragmentInputIndexCbCrAtlas) ]]
+) {
+    constexpr sampler atlas_sampler(mag_filter::linear, min_filter::linear);
+    const float4x4 ycbcrToRGBTransform = float4x4(
+        float4(+1.0000f, +1.0000f, +1.0000f, +0.0000f),
+        float4(+0.0000f, -0.3441f, +1.7720f, +0.0000f),
+        float4(+1.4020f, -0.7141f, +0.0000f, +0.0000f),
+        float4(-0.7010f, +0.5291f, -0.8860f, +1.0000f)
+    );
+    float4 ycbcr = float4(y_atlas.sample(atlas_sampler, input.atlas_position).r,
+                          cb_cr_atlas.sample(atlas_sampler, input.atlas_position).rg, 1.0);
+
+    input.background_color = ycbcrToRGBTransform * ycbcr;
+    return quad_sdf(input);
+}
+
 struct PathAtlasVertexOutput {
     float4 position [[position]];
     float2 st_position;

crates/media/src/media.rs 🔗

@@ -31,7 +31,10 @@ pub mod core_video {
     #![allow(non_snake_case)]
 
     use super::*;
-    pub use crate::bindings::kCVPixelFormatType_32BGRA;
+    pub use crate::bindings::{
+        kCVPixelFormatType_32BGRA, kCVPixelFormatType_420YpCbCr8BiPlanarFullRange,
+        kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange, kCVPixelFormatType_420YpCbCr8Planar,
+    };
     use crate::bindings::{kCVReturnSuccess, CVReturn, OSType};
     use anyhow::{anyhow, Result};
     use core_foundation::{
@@ -68,6 +71,14 @@ pub mod core_video {
             unsafe { CVPixelBufferGetHeight(self.as_concrete_TypeRef()) }
         }
 
+        pub fn plane_width(&self, plane: usize) -> usize {
+            unsafe { CVPixelBufferGetWidthOfPlane(self.as_concrete_TypeRef(), plane) }
+        }
+
+        pub fn plane_height(&self, plane: usize) -> usize {
+            unsafe { CVPixelBufferGetHeightOfPlane(self.as_concrete_TypeRef(), plane) }
+        }
+
         pub fn pixel_format_type(&self) -> OSType {
             unsafe { CVPixelBufferGetPixelFormatType(self.as_concrete_TypeRef()) }
         }
@@ -79,6 +90,8 @@ pub mod core_video {
         fn CVPixelBufferGetIOSurface(buffer: CVImageBufferRef) -> IOSurfaceRef;
         fn CVPixelBufferGetWidth(buffer: CVImageBufferRef) -> usize;
         fn CVPixelBufferGetHeight(buffer: CVImageBufferRef) -> usize;
+        fn CVPixelBufferGetWidthOfPlane(buffer: CVImageBufferRef, plane: usize) -> usize;
+        fn CVPixelBufferGetHeightOfPlane(buffer: CVImageBufferRef, plane: usize) -> usize;
         fn CVPixelBufferGetPixelFormatType(buffer: CVImageBufferRef) -> OSType;
     }