[naga hlsl-out] Implement external texture support

This adds HLSL backend support for `ImageClass::External` (ie WGSL's
`external_texture` texture type).

For each external texture global variable in the IR, we declare 3
`Texture2D` globals as well as a `cbuffer` for the params. The
additional bindings required by these are found in the newly added
`external_texture_binding_map`. Unique names for each can be obtained
using `NameKey::ExternalTextureGlobalVariable`.

For functions that contain ImageQuery::Size, ImageLoad, or ImageSample
expressions for external textures, ensure we have generated wrapper
functions for those expressions. When emitting code for the
expressions themselves, simply insert a call to the wrapper function.

For size queries, we return the value provided in the params
struct. If that value is [0, 0] then we query the size of the plane 0
texture and return that.

For load and sample, we sample the textures based on the number of
planes specified in the params struct. If there is more than one plane
we additionally perform YUV to RGB conversion using the provided
matrix.

Unfortunately HLSL does not allow structs to contain textures, meaning
we are unable to wrap the 3 textures and params struct variables in a
single variable that can be passed around.

For our wrapper functions we therefore ensure they take the three
textures and the params as consecutive arguments. Likewise, when
declaring user-defined functions with external texture arguments, we
expand the single external texture argument into 4 consecutive
arguments. (Using NameKey::ExternalTextureFunctionArgument to ensure
unique names for each.)

Thankfully external textures can only be used as either global
variables or function arguments. This means we only have to handle the
`Expression::GlobalVariable` and `Expression::FunctionArgument` cases
of `write_expr()`. Since in both cases we know the external texture
can only be an argument to either a user-defined function or one of
our wrapper functions, we can simply emit the names of the variables
for each three textures and the params struct in a comma-separated
list.

Author: jamienicol

naga-cli/src/bin/naga.rs

@@ -539,6 +539,7 @@ fn run() -> anyhow::Result<()> {
                 let missing = match Path::new(path).extension().and_then(|ex| ex.to_str()) {
                     Some("wgsl") => C::CLIP_DISTANCE | C::CULL_DISTANCE,
                     Some("metal") => C::CULL_DISTANCE | C::TEXTURE_EXTERNAL,
+                    Some("hlsl") => C::empty(),
                     _ => C::TEXTURE_EXTERNAL,
                 };
                 caps & !missing

naga/src/back/hlsl/help.rs

@@ -826,6 +826,7 @@ pub const RESERVED: &[&str] = &[
     super::writer::INSERT_BITS_FUNCTION,
     super::writer::SAMPLER_HEAP_VAR,
     super::writer::COMPARISON_SAMPLER_HEAP_VAR,
+    super::writer::SAMPLE_EXTERNAL_TEXTURE_FUNCTION,
     super::writer::ABS_FUNCTION,
     super::writer::DIV_FUNCTION,
     super::writer::MOD_FUNCTION,
@@ -834,6 +835,7 @@ pub const RESERVED: &[&str] = &[
     super::writer::F2U32_FUNCTION,
     super::writer::F2I64_FUNCTION,
     super::writer::F2U64_FUNCTION,
+    super::writer::IMAGE_LOAD_EXTERNAL_FUNCTION,
     super::writer::IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION,
 ];
 

naga/src/back/hlsl/keywords.rs

@@ -106,6 +106,37 @@ index buffer for each bind group. This buffer is accessed in the shader to get t
 sampler index within the heap. See the wgpu_hal dx12 backend documentation for more
 information.
 
+# External textures
+
+Support for [`crate::ImageClass::External`] textures is implemented by lowering
+each external texture global variable to 3 `Texture2D<float4>`s, and a `cbuffer`
+of type `NagaExternalTextureParams`. This provides up to 3 planes of texture
+data (for example single planar RGBA, or separate Y, Cb, and Cr planes), and the
+parameters buffer containing information describing how to handle these
+correctly. The bind target to use for each of these globals is specified via
+[`Options::external_texture_binding_map`].
+
+External textures are supported by WGSL's `textureDimensions()`,
+`textureLoad()`, and `textureSampleBaseClampToEdge()` built-in functions. These
+are implemented using helper functions. See the following functions for how
+these are generated:
+ * `Writer::write_wrapped_image_query_function`
+ * `Writer::write_wrapped_image_load_function`
+ * `Writer::write_wrapped_image_sample_function`
+
+Ideally the set of global variables could be wrapped in a single struct that
+could conveniently be passed around. But, alas, HLSL does not allow structs to
+have `Texture2D` members. Fortunately, however, external textures can only be
+used as arguments to either built-in or user-defined functions. We therefore
+expand any external texture function argument to four consecutive arguments (3
+textures and the params struct) when declaring user-defined functions, and
+ensure our built-in function implementations take the same arguments. Then,
+whenever we need to emit an external texture in `Writer::write_expr`, which
+fortunately can only ever be for a global variable or function argument, we
+simply emit the variable name of each of the three textures and the parameters
+struct in a comma-separated list. This won't win any awards for elegance, but
+it works for our purposes.
+
 [hlsl]: https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl
 [ilov]: https://gpuweb.github.io/gpuweb/wgsl/#internal-value-layout
 [16bb]: https://github.com/microsoft/DirectXShaderCompiler/wiki/Buffer-Packing#constant-buffer-packing
@@ -126,6 +157,35 @@ use thiserror::Error;
 
 use crate::{back, ir, proc};
 
+/// Direct3D 12 binding information for a global variable.
+///
+/// This type provides the HLSL-specific information Naga needs to declare and
+/// access an HLSL global variable that cannot be derived from the `Module`
+/// itself.
+///
+/// An HLSL global variable declaration includes details that the Direct3D API
+/// will use to refer to it. For example:
+///
+///    RWByteAddressBuffer s_sasm : register(u0, space2);
+///
+/// This defines a global `s_sasm` that a Direct3D root signature would refer to
+/// as register `0` in register space `2` in a `UAV` descriptor range. Naga can
+/// infer the register's descriptor range type from the variable's address class
+/// (writable [`Storage`] variables are implemented by Direct3D Unordered Access
+/// Views, the `u` register type), but the register number and register space
+/// must be supplied by the user.
+///
+/// The [`back::hlsl::Options`] structure provides `BindTarget`s for various
+/// situations in which Naga may need to generate an HLSL global variable, like
+/// [`binding_map`] for Naga global variables, or [`push_constants_target`] for
+/// a module's sole [`PushConstant`] variable. See those fields' documentation
+/// for details.
+///
+/// [`Storage`]: crate::ir::AddressSpace::Storage
+/// [`back::hlsl::Options`]: Options
+/// [`binding_map`]: Options::binding_map
+/// [`push_constants_target`]: Options::push_constants_target
+/// [`PushConstant`]: crate::ir::AddressSpace::PushConstant
 #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
@@ -335,6 +395,62 @@ where
 
 pub type DynamicStorageBufferOffsetsTargets = alloc::collections::BTreeMap<u32, OffsetsBindTarget>;
 
+/// HLSL binding information for a Naga [`External`] image global variable.
+///
+/// See the module documentation's section on [External textures][mod] for details.
+///
+/// [`External`]: crate::ir::ImageClass::External
+/// [mod]: #external-textures
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+pub struct ExternalTextureBindTarget {
+    /// HLSL binding information for the individual plane textures.
+    ///
+    /// Each of these should refer to an HLSL `Texture2D<float4>` holding one
+    /// plane of data for the external texture. The exact meaning of each plane
+    /// varies at runtime depending on where the external texture's data
+    /// originated.
+    pub planes: [BindTarget; 3],
+
+    /// HLSL binding information for a buffer holding the sampling parameters.
+    ///
+    /// This should refer to a cbuffer of type `NagaExternalTextureParams`, that
+    /// the code Naga generates for `textureSampleBaseClampToEdge` consults to
+    /// decide how to combine the data in [`planes`] to get the result required
+    /// by the spec.
+    ///
+    /// [`planes`]: Self::planes
+    pub params: BindTarget,
+}
+
+#[cfg(any(feature = "serialize", feature = "deserialize"))]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+struct ExternalTextureBindingMapSerialization {
+    resource_binding: crate::ResourceBinding,
+    bind_target: ExternalTextureBindTarget,
+}
+
+#[cfg(feature = "deserialize")]
+fn deserialize_external_texture_binding_map<'de, D>(
+    deserializer: D,
+) -> Result<ExternalTextureBindingMap, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    use serde::Deserialize;
+
+    let vec = Vec::<ExternalTextureBindingMapSerialization>::deserialize(deserializer)?;
+    let mut map = ExternalTextureBindingMap::default();
+    for item in vec {
+        map.insert(item.resource_binding, item.bind_target);
+    }
+    Ok(map)
+}
+pub type ExternalTextureBindingMap =
+    alloc::collections::BTreeMap<crate::ResourceBinding, ExternalTextureBindTarget>;
+
 /// Shorthand result used internally by the backend
 type BackendResult = Result<(), Error>;
 
@@ -354,21 +470,47 @@ pub enum EntryPointError {
 pub struct Options {
     /// The hlsl shader model to be used
     pub shader_model: ShaderModel,
-    /// Map of resources association to binding locations.
+
+    /// HLSL binding information for each Naga global variable.
+    ///
+    /// This maps Naga [`GlobalVariable`]'s [`ResourceBinding`]s to a
+    /// [`BindTarget`] specifying its register number and space, along with
+    /// other details necessary to generate a full HLSL declaration for it,
+    /// or to access its value.
+    ///
+    /// This must provide a [`BindTarget`] for every [`GlobalVariable`] in the
+    /// [`Module`] that has a [`binding`].
+    ///
+    /// [`GlobalVariable`]: crate::ir::GlobalVariable
+    /// [`ResourceBinding`]: crate::ir::ResourceBinding
+    /// [`Module`]: crate::ir::Module
+    /// [`binding`]: crate::ir::GlobalVariable::binding
     #[cfg_attr(
         feature = "deserialize",
         serde(deserialize_with = "deserialize_binding_map")
     )]
     pub binding_map: BindingMap,
+
     /// Don't panic on missing bindings, instead generate any HLSL.
     pub fake_missing_bindings: bool,
     /// Add special constants to `SV_VertexIndex` and `SV_InstanceIndex`,
     /// to make them work like in Vulkan/Metal, with help of the host.
     pub special_constants_binding: Option<BindTarget>,
-    /// Bind target of the push constant buffer
+
+    /// HLSL binding information for the [`PushConstant`] global, if present.
+    ///
+    /// If a module contains a global in the [`PushConstant`] address space, the
+    /// `dx12` backend stores its value directly in the root signature as a
+    /// series of [`D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS`], whose binding
+    /// information is given here.
+    ///
+    /// [`PushConstant`]: crate::ir::AddressSpace::PushConstant
+    /// [`D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS`]: https://learn.microsoft.com/en-us/windows/win32/api/d3d12/ne-d3d12-d3d12_root_parameter_type
     pub push_constants_target: Option<BindTarget>,
-    /// Bind target of the sampler heap and comparison sampler heap.
+
+    /// HLSL binding information for the sampler heap and comparison sampler heap.
     pub sampler_heap_target: SamplerHeapBindTargets,
+
     /// Mapping of each bind group's sampler index buffer to a bind target.
     #[cfg_attr(
         feature = "deserialize",
@@ -381,6 +523,18 @@ pub struct Options {
         serde(deserialize_with = "deserialize_storage_buffer_offsets")
     )]
     pub dynamic_storage_buffer_offsets_targets: DynamicStorageBufferOffsetsTargets,
+    #[cfg_attr(
+        feature = "deserialize",
+        serde(deserialize_with = "deserialize_external_texture_binding_map")
+    )]
+
+    /// HLSL binding information for [`External`] image global variables.
+    ///
+    /// See [`ExternalTextureBindTarget`] for details.
+    ///
+    /// [`External`]: crate::ir::ImageClass::External
+    pub external_texture_binding_map: ExternalTextureBindingMap,
+
     /// Should workgroup variables be zero initialized (by polyfilling)?
     pub zero_initialize_workgroup_memory: bool,
     /// Should we restrict indexing of vectors, matrices and arrays?
@@ -401,6 +555,7 @@ impl Default for Options {
             sampler_buffer_binding_map: alloc::collections::BTreeMap::default(),
             push_constants_target: None,
             dynamic_storage_buffer_offsets_targets: alloc::collections::BTreeMap::new(),
+            external_texture_binding_map: ExternalTextureBindingMap::default(),
             zero_initialize_workgroup_memory: true,
             restrict_indexing: true,
             force_loop_bounding: true,
@@ -425,6 +580,29 @@ impl Options {
             None => Err(EntryPointError::MissingBinding(*res_binding)),
         }
     }
+
+    fn resolve_external_texture_resource_binding(
+        &self,
+        res_binding: &crate::ResourceBinding,
+    ) -> Result<ExternalTextureBindTarget, EntryPointError> {
+        match self.external_texture_binding_map.get(res_binding) {
+            Some(target) => Ok(*target),
+            None if self.fake_missing_bindings => {
+                let fake = BindTarget {
+                    space: res_binding.group as u8,
+                    register: res_binding.binding,
+                    binding_array_size: None,
+                    dynamic_storage_buffer_offsets_index: None,
+                    restrict_indexing: false,
+                };
+                Ok(ExternalTextureBindTarget {
+                    planes: [fake, fake, fake],
+                    params: fake,
+                })
+            }
+            None => Err(EntryPointError::MissingBinding(*res_binding)),
+        }
+    }
 }
 
 /// Reflection info for entry point names.
@@ -479,6 +657,7 @@ enum WrappedType {
     ArrayLength(help::WrappedArrayLength),
     ImageSample(help::WrappedImageSample),
     ImageQuery(help::WrappedImageQuery),
+    ImageLoad(help::WrappedImageLoad),
     ImageLoadScalar(crate::Scalar),
     Constructor(help::WrappedConstructor),
     StructMatrixAccess(help::WrappedStructMatrixAccess),