Unconditionally Generate Bindless Samplers in DX12

Cargo.toml

@@ -77,6 +77,7 @@ version = "24.0.0"
 anyhow = "1.0.95"
 argh = "0.1.13"
 arrayvec = "0.7"
+approx = "0.5"
 bincode = "1"
 bit-vec = "0.8"
 bitflags = "2.7"

naga/src/back/hlsl/help.rs

@@ -1190,6 +1190,83 @@ impl<W: Write> super::Writer<'_, W> {
         Ok(())
     }
 
+    /// Writes out the sampler heap declarations if they haven't been written yet.
+    pub(super) fn write_sampler_heaps(&mut self) -> BackendResult {
+        if self.wrapped.sampler_heaps {
+            return Ok(());
+        }
+
+        writeln!(
+            self.out,
+            "SamplerState {}[2048]: register(s{}, space{});",
+            super::writer::SAMPLER_HEAP_VAR,
+            self.options.sampler_heap_target.standard_samplers.register,
+            self.options.sampler_heap_target.standard_samplers.space
+        )?;
+        writeln!(
+            self.out,
+            "SamplerComparisonState {}[2048]: register(s{}, space{});",
+            super::writer::COMPARISON_SAMPLER_HEAP_VAR,
+            self.options
+                .sampler_heap_target
+                .comparison_samplers
+                .register,
+            self.options.sampler_heap_target.comparison_samplers.space
+        )?;
+
+        self.wrapped.sampler_heaps = true;
+
+        Ok(())
+    }
+
+    /// Writes out the sampler index buffer declaration if it hasn't been written yet.
+    pub(super) fn write_wrapped_sampler_buffer(
+        &mut self,
+        key: super::SamplerIndexBufferKey,
+    ) -> BackendResult {
+        // The astute will notice that we do a double hash lookup, but we do this to avoid
+        // holding a mutable reference to `self` while trying to call `write_sampler_heaps`.
+        //
+        // We only pay this double lookup cost when we actually need to write out the sampler
+        // buffer, which should be not be common.
+
+        if self.wrapped.sampler_index_buffers.contains_key(&key) {
+            return Ok(());
+        };
+
+        self.write_sampler_heaps()?;
+
+        // Because the group number can be arbitrary, we use the namer to generate a unique name
+        // instead of adding it to the reserved name list.
+        let sampler_array_name = self
+            .namer
+            .call(&format!("nagaGroup{}SamplerIndexArray", key.group));
+
+        let bind_target = match self.options.sampler_buffer_binding_map.get(&key) {
+            Some(&bind_target) => bind_target,
+            None if self.options.fake_missing_bindings => super::BindTarget {
+                space: u8::MAX,
+                register: key.group,
+                binding_array_size: None,
+            },
+            None => {
+                unreachable!("Sampler buffer of group {key:?} not bound to a register");
+            }
+        };
+
+        writeln!(
+            self.out,
+            "StructuredBuffer<uint> {sampler_array_name} : register(t{}, space{});",
+            bind_target.register, bind_target.space
+        )?;
+
+        self.wrapped
+            .sampler_index_buffers
+            .insert(key, sampler_array_name);
+
+        Ok(())
+    }
+
     pub(super) fn write_texture_coordinates(
         &mut self,
         kind: &str,

naga/src/back/hlsl/keywords.rs

@@ -820,6 +820,8 @@ pub const RESERVED: &[&str] = &[
     super::writer::FREXP_FUNCTION,
     super::writer::EXTRACT_BITS_FUNCTION,
     super::writer::INSERT_BITS_FUNCTION,
+    super::writer::SAMPLER_HEAP_VAR,
+    super::writer::COMPARISON_SAMPLER_HEAP_VAR,
 ];
 
 // DXC scalar types, from https://github.com/microsoft/DirectXShaderCompiler/blob/18c9e114f9c314f93e68fbc72ce207d4ed2e65ae/tools/clang/lib/AST/ASTContextHLSL.cpp#L48-L254

naga/src/back/hlsl/mod.rs

@@ -92,6 +92,15 @@ float3x2 GetMatmOnBaz(Baz obj) {
 We also emit an analogous `Set` function, as well as functions for
 accessing individual columns by dynamic index.
 
+## Sampler Handling
+
+Due to limitations in how sampler heaps work in D3D12, we need to access samplers
+through a layer of indirection. Instead of directly binding samplers, we bind the entire
+sampler heap as both a standard and a comparison sampler heap. We then use a sampler
+index buffer for each bind group. This buffer is accessed in the shader to get the actual
+sampler index within the heap. See the wgpu_hal dx12 backend documentation for more
+information.
+
 [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
@@ -110,7 +119,7 @@ use thiserror::Error;
 
 use crate::{back, proc};
 
-#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
 pub struct BindTarget {
@@ -179,6 +188,43 @@ impl crate::ImageDimension {
     }
 }
 
+#[derive(Clone, Copy, Debug, Hash, Eq, Ord, PartialEq, PartialOrd)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+pub struct SamplerIndexBufferKey {
+    pub group: u32,
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+#[cfg_attr(feature = "deserialize", serde(default))]
+pub struct SamplerHeapBindTargets {
+    pub standard_samplers: BindTarget,
+    pub comparison_samplers: BindTarget,
+}
+
+impl Default for SamplerHeapBindTargets {
+    fn default() -> Self {
+        Self {
+            standard_samplers: BindTarget {
+                space: 0,
+                register: 0,
+                binding_array_size: None,
+            },
+            comparison_samplers: BindTarget {
+                space: 1,
+                register: 0,
+                binding_array_size: None,
+            },
+        }
+    }
+}
+
+// We use a BTreeMap here so that we can hash it.
+pub type SamplerIndexBufferBindingMap =
+    std::collections::BTreeMap<SamplerIndexBufferKey, BindTarget>;
+
 /// Shorthand result used internally by the backend
 type BackendResult = Result<(), Error>;
 
@@ -207,6 +253,10 @@ pub struct Options {
     pub special_constants_binding: Option<BindTarget>,
     /// Bind target of the push constant buffer
     pub push_constants_target: Option<BindTarget>,
+    /// Bind target of 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.
+    pub sampler_buffer_binding_map: SamplerIndexBufferBindingMap,
     /// Should workgroup variables be zero initialized (by polyfilling)?
     pub zero_initialize_workgroup_memory: bool,
     /// Should we restrict indexing of vectors, matrices and arrays?
@@ -220,6 +270,8 @@ impl Default for Options {
             binding_map: BindingMap::default(),
             fake_missing_bindings: true,
             special_constants_binding: None,
+            sampler_heap_target: SamplerHeapBindTargets::default(),
+            sampler_buffer_binding_map: std::collections::BTreeMap::default(),
             push_constants_target: None,
             zero_initialize_workgroup_memory: true,
             restrict_indexing: true,
@@ -233,13 +285,13 @@ impl Options {
         res_binding: &crate::ResourceBinding,
     ) -> Result<BindTarget, EntryPointError> {
         match self.binding_map.get(res_binding) {
-            Some(target) => Ok(target.clone()),
+            Some(target) => Ok(*target),
             None if self.fake_missing_bindings => Ok(BindTarget {
                 space: res_binding.group as u8,
                 register: res_binding.binding,
                 binding_array_size: None,
             }),
-            None => Err(EntryPointError::MissingBinding(res_binding.clone())),
+            None => Err(EntryPointError::MissingBinding(*res_binding)),
         }
     }
 }
@@ -279,6 +331,10 @@ struct Wrapped {
     struct_matrix_access: crate::FastHashSet<help::WrappedStructMatrixAccess>,
     mat_cx2s: crate::FastHashSet<help::WrappedMatCx2>,
     math: crate::FastHashSet<help::WrappedMath>,
+    /// If true, the sampler heaps have been written out.
+    sampler_heaps: bool,
+    // Mapping from SamplerIndexBufferKey to the name the namer returned.
+    sampler_index_buffers: crate::FastHashMap<SamplerIndexBufferKey, String>,
 }
 
 impl Wrapped {