[naga hlsl-out] Handle external texture color space conversion

This adds several fields to `ExternalTextureDescriptor`, specifying
how to handle color space conversion for an external texture. These
fields consist of transfer functions for the source and destination
color spaces, and a matrix for converting between gamuts. This allows
`ImageSample` and `ImageLoad` operations on external textures to
return values in a desired destination color space rather than the
source color space of the underlying planes.

These fields are plumbed through to the `ExternalTextureParams`
uniform buffer from which they are exposed to the shader. Following
conversion from YUV to RGB after sampling/loading from the external
texture planes, the shader uses them to gamma decode to linear RGB in
the source color space, convert from source to destination gamut, then
finally gamma encode to non-linear RGB in the destination color space.

Author: jamienicol

naga/src/back/hlsl/help.rs

@@ -264,6 +264,61 @@ impl<W: Write> super::Writer<'_, W> {
         Ok(())
     }
 
+    /// Helper function used by [`Self::write_wrapped_image_load_function`] and
+    /// [`Self::write_wrapped_image_sample_function`] to write the shared YUV
+    /// to RGB conversion code for external textures. Expects the preceding
+    /// code to declare the Y component as a `float` variable of name `y`, the
+    /// UV components as a `float2` variable of name `uv`, and the external
+    /// texture params as a variable of name `params`. The emitted code will
+    /// return the result.
+    fn write_convert_yuv_to_rgb_and_return(
+        &mut self,
+        level: crate::back::Level,
+        y: &str,
+        uv: &str,
+        params: &str,
+    ) -> BackendResult {
+        let l1 = level;
+        let l2 = l1.next();
+
+        // Convert from YUV to non-linear RGB in the source color space. We
+        // declare our matrices as row_major in HLSL, therefore we must reverse
+        // the order of this multiplication
+        writeln!(
+            self.out,
+            "{l1}float3 srcGammaRgb = mul(float4({y}, {uv}, 1.0), {params}.yuv_conversion_matrix).rgb;"
+        )?;
+
+        // Apply the inverse of the source transfer function to convert to
+        // linear RGB in the source color space.
+        writeln!(
+            self.out,
+            "{l1}float3 srcLinearRgb = srcGammaRgb < {params}.src_tf.k * {params}.src_tf.b ?"
+        )?;
+        writeln!(self.out, "{l2}srcGammaRgb / {params}.src_tf.k :")?;
+        writeln!(self.out, "{l2}pow((srcGammaRgb + {params}.src_tf.a - 1.0) / {params}.src_tf.a, {params}.src_tf.g);")?;
+
+        // Multiply by the gamut conversion matrix to convert to linear RGB in
+        // the destination color space. We declare our matrices as row_major in
+        // HLSL, therefore we must reverse the order of this multiplication.
+        writeln!(
+            self.out,
+            "{l1}float3 dstLinearRgb = mul(srcLinearRgb, {params}.gamut_conversion_matrix);"
+        )?;
+
+        // Finally, apply the dest transfer function to convert to non-linear
+        // RGB in the destination color space, and return the result.
+        writeln!(
+            self.out,
+            "{l1}float3 dstGammaRgb = dstLinearRgb < {params}.dst_tf.b ?"
+        )?;
+        writeln!(self.out, "{l2}{params}.dst_tf.k * dstLinearRgb :")?;
+        writeln!(self.out, "{l2}{params}.dst_tf.a * pow(dstLinearRgb, 1.0 / {params}.dst_tf.g) - ({params}.dst_tf.a - 1);")?;
+
+        writeln!(self.out, "{l1}return float4(dstGammaRgb, 1.0);")?;
+        Ok(())
+    }
+
     pub(super) fn write_wrapped_image_load_function(
         &mut self,
         module: &crate::Module,
@@ -346,12 +401,7 @@ impl<W: Write> super::Writer<'_, W> {
                 writeln!(self.out, "{l3}uv = float2(plane1.Load(uint3(plane1_coords, 0u)).x, plane2.Load(uint3(plane2_coords, 0u)).x);")?;
                 writeln!(self.out, "{l2}}}")?;
 
-                // Convert from YUV to RGB. We declare our matrices as row_major in HLSL,
-                // therefore we must reverse the order of this multiplication
-                writeln!(
-                    self.out,
-                    "{l2}return mul(float4(y, uv, 1.0), params.yuv_conversion_matrix);"
-                )?;
+                self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "params")?;
 
                 writeln!(self.out, "{l1}}}")?;
                 writeln!(self.out, "}}")?;
@@ -478,12 +528,8 @@ impl<W: Write> super::Writer<'_, W> {
                 writeln!(self.out, "{l3}uv = float2(plane1.SampleLevel(samp, plane1_coords, 0.0f).x, plane2.SampleLevel(samp, plane2_coords, 0.0f).x);")?;
                 writeln!(self.out, "{l2}}}")?;
 
-                // Convert from YUV to RGB. We declare our matrices as row_major in HLSL,
-                // therefore we must reverse the order of this multiplication
-                writeln!(
-                    self.out,
-                    "{l2}return mul(float4(y, uv, 1.0), params.yuv_conversion_matrix);"
-                )?;
+                self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "params")?;
+
                 writeln!(self.out, "{l1}}}")?;
                 writeln!(self.out, "}}")?;
                 writeln!(self.out)?;

naga/src/back/wgsl/writer.rs

@@ -131,6 +131,7 @@ impl<W: Write> Writer<W> {
             .values()
             .any(|t| *t == ty)
             || Some(ty) == module.special_types.external_texture_params
+            || Some(ty) == module.special_types.external_texture_transfer_function
     }
 
     pub fn write(&mut self, module: &Module, info: &valid::ModuleInfo) -> BackendResult {

naga/src/compact/mod.rs

@@ -381,6 +381,7 @@ impl<'module> ModuleTracer<'module> {
             ref ray_vertex_return,
             ref predeclared_types,
             ref external_texture_params,
+            ref external_texture_transfer_function,
         } = *special_types;
 
         if let Some(ray_desc) = *ray_desc {
@@ -398,6 +399,9 @@ impl<'module> ModuleTracer<'module> {
         if let Some(external_texture_params) = *external_texture_params {
             self.types_used.insert(external_texture_params);
         }
+        if let Some(external_texture_transfer_function) = *external_texture_transfer_function {
+            self.types_used.insert(external_texture_transfer_function);
+        }
         for (_, &handle) in predeclared_types {
             self.types_used.insert(handle);
         }
@@ -540,6 +544,7 @@ impl ModuleMap {
             ref mut ray_vertex_return,
             ref mut predeclared_types,
             ref mut external_texture_params,
+            ref mut external_texture_transfer_function,
         } = *special;
 
         if let Some(ref mut ray_desc) = *ray_desc {
@@ -557,6 +562,12 @@ impl ModuleMap {
             self.types.adjust(external_texture_params);
         }
 
+        if let Some(ref mut external_texture_transfer_function) =
+            *external_texture_transfer_function
+        {
+            self.types.adjust(external_texture_transfer_function);
+        }
+
         for handle in predeclared_types.values_mut() {
             self.types.adjust(handle);
         }

naga/src/front/type_gen.rs

@@ -276,14 +276,29 @@ impl crate::Module {
         handle
     }
 
-    /// Generate [`SpecialTypes::external_texture_params`].
+    /// Generate [`SpecialTypes::external_texture_params`] and
+    /// [`SpecialTypes::external_texture_transfer_function`].
+    ///
+    /// Other than the WGSL backend, every backend that supports external
+    /// textures does so by lowering them to a set of ordinary textures and
+    /// some parameters saying how to sample from them. These types are used
+    /// for said parameters. Note that they are not used by the IR, but
+    /// generated purely as a convenience for the backends.
     ///
     /// [`SpecialTypes::external_texture_params`]: crate::ir::SpecialTypes::external_texture_params
-    pub fn generate_external_texture_params_type(&mut self) -> Handle<crate::Type> {
-        if let Some(handle) = self.special_types.external_texture_params {
-            return handle;
+    /// [`SpecialTypes::external_texture_transfer_function`]: crate::ir::SpecialTypes::external_texture_transfer_function
+    pub fn generate_external_texture_types(&mut self) {
+        if self.special_types.external_texture_params.is_some() {
+            return;
         }
 
+        let ty_f32 = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Scalar(crate::Scalar::F32),
+            },
+            Span::UNDEFINED,
+        );
         let ty_u32 = self.types.insert(
             crate::Type {
                 name: None,
@@ -312,6 +327,17 @@ impl crate::Module {
             },
             Span::UNDEFINED,
         );
+        let ty_mat3x3f = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Matrix {
+                    columns: crate::VectorSize::Tri,
+                    rows: crate::VectorSize::Tri,
+                    scalar: crate::Scalar::F32,
+                },
+            },
+            Span::UNDEFINED,
+        );
         let ty_mat4x4f = self.types.insert(
             crate::Type {
                 name: None,
@@ -324,7 +350,44 @@ impl crate::Module {
             Span::UNDEFINED,
         );
 
-        let handle = self.types.insert(
+        let transfer_fn_handle = self.types.insert(
+            crate::Type {
+                name: Some("NagaExternalTextureTransferFn".to_string()),
+                inner: crate::TypeInner::Struct {
+                    members: vec![
+                        crate::StructMember {
+                            name: Some("a".to_string()),
+                            ty: ty_f32,
+                            binding: None,
+                            offset: 0,
+                        },
+                        crate::StructMember {
+                            name: Some("b".to_string()),
+                            ty: ty_f32,
+                            binding: None,
+                            offset: 4,
+                        },
+                        crate::StructMember {
+                            name: Some("g".to_string()),
+                            ty: ty_f32,
+                            binding: None,
+                            offset: 8,
+                        },
+                        crate::StructMember {
+                            name: Some("k".to_string()),
+                            ty: ty_f32,
+                            binding: None,
+                            offset: 12,
+                        },
+                    ],
+                    span: 16,
+                },
+            },
+            Span::UNDEFINED,
+        );
+        self.special_types.external_texture_transfer_function = Some(transfer_fn_handle);
+
+        let params_handle = self.types.insert(
             crate::Type {
                 name: Some("NagaExternalTextureParams".to_string()),
                 inner: crate::TypeInner::Struct {
@@ -335,39 +398,55 @@ impl crate::Module {
                             binding: None,
                             offset: 0,
                         },
+                        crate::StructMember {
+                            name: Some("gamut_conversion_matrix".to_string()),
+                            ty: ty_mat3x3f,
+                            binding: None,
+                            offset: 64,
+                        },
+                        crate::StructMember {
+                            name: Some("src_tf".to_string()),
+                            ty: transfer_fn_handle,
+                            binding: None,
+                            offset: 112,
+                        },
+                        crate::StructMember {
+                            name: Some("dst_tf".to_string()),
+                            ty: transfer_fn_handle,
+                            binding: None,
+                            offset: 128,
+                        },
                         crate::StructMember {
                             name: Some("sample_transform".to_string()),
                             ty: ty_mat3x2f,
                             binding: None,
-                            offset: 64,
+                            offset: 144,
                         },
                         crate::StructMember {
                             name: Some("load_transform".to_string()),
                             ty: ty_mat3x2f,
                             binding: None,
-                            offset: 88,
+                            offset: 168,
                         },
                         crate::StructMember {
                             name: Some("size".to_string()),
                             ty: ty_vec2u,
                             binding: None,
-                            offset: 112,
+                            offset: 192,
                         },
                         crate::StructMember {
                             name: Some("num_planes".to_string()),
                             ty: ty_u32,
                             binding: None,
-                            offset: 120,
+                            offset: 200,
                         },
                     ],
-                    span: 128,
+                    span: 208,
                 },
             },
             Span::UNDEFINED,
         );
-
-        self.special_types.external_texture_params = Some(handle);
-        handle
+        self.special_types.external_texture_params = Some(params_handle);
     }
 
     /// Populate this module's [`SpecialTypes::predeclared_types`] type and return the handle.

naga/src/front/wgsl/lower/mod.rs

@@ -3991,14 +3991,15 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     // sample from them. We don't know which backend will
                     // consume the `Module` we're building, but in case it's not
                     // WGSL, populate `SpecialTypes::external_texture_params`
-                    // with the type the backend will use for the parameter
+                    // and `SpecialTypes::external_texture_transfer_function`
+                    // with the types the backend will use for the parameter
                     // buffer.
                     //
-                    // This is *not* the type we are lowering here: that's an
-                    // ordinary `TypeInner::Image`. But the fact we are
-                    // lowering a `texture_external` implies the backends may
-                    // need `SpecialTypes::external_texture_params` too.
-                    ctx.module.generate_external_texture_params_type();
+                    // Neither of these are the type we are lowering here:
+                    // that's an ordinary `TypeInner::Image`. But the fact we
+                    // are lowering a `texture_external` implies the backends
+                    // may need these additional types too.
+                    ctx.module.generate_external_texture_types();
                 }
                 ir::TypeInner::Image {
                     dim,

naga/src/ir/mod.rs

@@ -2369,19 +2369,42 @@ pub struct SpecialTypes {
     /// In WGSL, this type would be:
     ///
     /// ```ignore
-    /// struct NagaExternalTextureParams {      // align size offset
-    ///     yuv_conversion_matrix: mat4x4<f32>, //    16   64      0
-    ///     sample_transform: mat3x2<f32>,      //     8   24     64
-    ///     load_transform: mat3x2<f32>,        //     8   24     88
-    ///     size: vec2<u32>,                    //     8    8    112
-    ///     num_planes: u32,                    //     4    4    120
-    /// }                         // whole struct:    16  128
+    /// struct NagaExternalTextureParams {         // align size offset
+    ///     yuv_conversion_matrix: mat4x4<f32>,    //    16   64      0
+    ///     gamut_conversion_matrix: mat3x3<f32>,  //    16   48     64
+    ///     src_tf: NagaExternalTextureTransferFn, //     4   16    112
+    ///     dst_tf: NagaExternalTextureTransferFn, //     4   16    128
+    ///     sample_transform: mat3x2<f32>,         //     8   24    144
+    ///     load_transform: mat3x2<f32>,           //     8   24    168
+    ///     size: vec2<u32>,                       //     8    8    192
+    ///     num_planes: u32,                       //     4    4    200
+    /// }                            // whole struct:    16  208
     /// ```
     ///
-    /// Call [`Module::generate_external_texture_params_type`] to populate this
-    /// if needed and return the handle.
+    /// Call [`Module::generate_external_texture_types`] to populate this if
+    /// needed.
     pub external_texture_params: Option<Handle<Type>>,
 
+    /// Struct describing a gamma encoding transfer function. Member of
+    /// `NagaExternalTextureParams`, describing how the backend should perform
+    /// color space conversion when sampling from [`ImageClass::External`]
+    /// textures.
+    ///
+    /// In WGSL, this type would be:
+    ///
+    /// ```ignore
+    /// struct NagaExternalTextureTransferFn { // align size offset
+    ///     a: f32,                            //     4    4      0
+    ///     b: f32,                            //     4    4      4
+    ///     g: f32,                            //     4    4      8
+    ///     k: f32,                            //     4    4     12
+    /// }                         // whole struct:    4   16
+    /// ```
+    ///
+    /// Call [`Module::generate_external_texture_types`] to populate this if
+    /// needed.
+    pub external_texture_transfer_function: Option<Handle<Type>>,
+
     /// Types for predeclared wgsl types instantiated on demand.
     ///
     /// Call [`Module::generate_predeclared_type`] to populate this if