[naga hlsl-out] HLSL implemention of external textures

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/keywords.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/mod.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),