[vulkan, naga spv-out] Remap resource bindings

In order to support external textures, we must be able to map a single
external texture resource binding to multiple Vulkan descriptors. This
means we must be able to override the `Binding` and `DescriptorSet`
values for global variables when generating SPIR-V, rather than simply
passing through the group and binding values from Naga IR.

This patch extends the existing SPIR-V Naga backend's `BindingMap` to
contain a descriptor set and binding value in addition to the existing
array size. When creating BindGroupLayouts/BindGroups we use a
sequentially incrementing value for each entry's binding value,
continuing to just use the bind group index as the descriptor set
value.

The Naga backend looks up each resource in the map when emitting its
`Binding` and `DescriptorSet` decorations. If the entry cannot be
found in the map, it will either error or emit fake bindings based on
its configuration.

naga-test/src/lib.rs

@@ -155,6 +155,7 @@ impl SpirvOutParameters {
                 Some(self.capabilities.clone())
             },
             bounds_check_policies,
+            fake_missing_bindings: true,
             binding_map: self.binding_map.clone(),
             zero_initialize_workgroup_memory: spv::ZeroInitializeWorkgroupMemoryMode::Polyfill,
             force_loop_bounding: true,

naga/src/back/spv/mod.rs

@@ -79,6 +79,8 @@ pub enum Error {
     Override,
     #[error(transparent)]
     ResolveArraySizeError(#[from] crate::proc::ResolveArraySizeError),
+    #[error("mapping of {0:?} is missing")]
+    MissingBinding(crate::ResourceBinding),
 }
 
 #[derive(Default)]
@@ -760,6 +762,7 @@ pub struct Writer {
     constant_ids: HandleVec<crate::Expression, Word>,
     cached_constants: crate::FastHashMap<CachedConstant, Word>,
     global_variables: HandleVec<crate::GlobalVariable, GlobalVariable>,
+    fake_missing_bindings: bool,
     binding_map: BindingMap,
 
     // Cached expressions are only meaningful within a BlockContext, but we
@@ -811,10 +814,12 @@ bitflags::bitflags! {
     }
 }
 
-#[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 BindingInfo {
+    pub descriptor_set: u32,
+    pub binding: u32,
     /// If the binding is an unsized binding array, this overrides the size.
     pub binding_array_size: Option<u32>,
 }
@@ -839,6 +844,10 @@ pub struct Options<'a> {
     /// Configuration flags for the writer.
     pub flags: WriterFlags,
 
+    /// Don't panic on missing bindings. Instead use fake values for `Binding`
+    /// and `DescriptorSet` decorations. This may result in invalid SPIR-V.
+    pub fake_missing_bindings: bool,
+
     /// Map of resources to information about the binding.
     pub binding_map: BindingMap,
 
@@ -877,6 +886,7 @@ impl Default for Options<'_> {
         Options {
             lang_version: (1, 0),
             flags,
+            fake_missing_bindings: true,
             binding_map: BindingMap::default(),
             capabilities: None,
             bounds_check_policies: BoundsCheckPolicies::default(),

naga/src/back/spv/writer.rs

@@ -87,6 +87,7 @@ impl Writer {
             constant_ids: HandleVec::new(),
             cached_constants: crate::FastHashMap::default(),
             global_variables: HandleVec::new(),
+            fake_missing_bindings: options.fake_missing_bindings,
             binding_map: options.binding_map.clone(),
             saved_cached: CachedExpressions::default(),
             gl450_ext_inst_id,
@@ -149,6 +150,7 @@ impl Writer {
             force_loop_bounding: self.force_loop_bounding,
             use_storage_input_output_16: self.use_storage_input_output_16,
             capabilities_available: take(&mut self.capabilities_available),
+            fake_missing_bindings: self.fake_missing_bindings,
             binding_map: take(&mut self.binding_map),
 
             // Initialized afresh:
@@ -469,6 +471,26 @@ impl Writer {
         })
     }
 
+    /// Resolve the [`BindingInfo`] for a [`crate::ResourceBinding`] from the
+    /// provided [`Writer::binding_map`].
+    ///
+    /// If the specified resource is not present in the binding map this will
+    /// return an error, unless [`Writer::fake_missing_bindings`] is set.
+    fn resolve_resource_binding(
+        &self,
+        res_binding: &crate::ResourceBinding,
+    ) -> Result<BindingInfo, Error> {
+        match self.binding_map.get(res_binding) {
+            Some(target) => Ok(*target),
+            None if self.fake_missing_bindings => Ok(BindingInfo {
+                descriptor_set: res_binding.group,
+                binding: res_binding.binding,
+                binding_array_size: None,
+            }),
+            None => Err(Error::MissingBinding(*res_binding)),
+        }
+    }
+
     /// Emits code for any wrapper functions required by the expressions in ir_function.
     /// The IDs of any emitted functions will be stored in [`Self::wrapped_functions`].
     fn write_wrapped_functions(
@@ -2241,13 +2263,11 @@ impl Writer {
         // and it is failing on 0.
         let mut substitute_inner_type_lookup = None;
         if let Some(ref res_binding) = global_variable.binding {
-            self.decorate(id, Decoration::DescriptorSet, &[res_binding.group]);
-            self.decorate(id, Decoration::Binding, &[res_binding.binding]);
+            let bind_target = self.resolve_resource_binding(res_binding)?;
+            self.decorate(id, Decoration::DescriptorSet, &[bind_target.descriptor_set]);
+            self.decorate(id, Decoration::Binding, &[bind_target.binding]);
 
-            if let Some(&BindingInfo {
-                binding_array_size: Some(remapped_binding_array_size),
-            }) = self.binding_map.get(res_binding)
-            {
+            if let Some(remapped_binding_array_size) = bind_target.binding_array_size {
                 if let crate::TypeInner::BindingArray { base, .. } =
                     ir_module.types[global_variable.ty].inner
                 {

naga/tests/in/wgsl/binding-arrays.toml

@@ -62,5 +62,5 @@ resource_binding = { group = 0, binding = 8 }
 version = [1, 1]
 
 [[spv.binding_map]]
-bind_target = { binding_array_size = 10 }
+bind_target = { descriptor_set = 0, binding = 0, binding_array_size = 10 }
 resource_binding = { group = 0, binding = 0 }

naga/tests/in/wgsl/binding-buffer-arrays.toml

@@ -9,5 +9,5 @@ image = "ReadZeroSkipWrite"
 version = [1, 1]
 
 [[spv.binding_map]]
-bind_target = { binding_array_size = 10 }
+bind_target = { descriptor_set = 0, binding = 0, binding_array_size = 10 }
 resource_binding = { group = 0, binding = 0 }

wgpu-hal/src/vulkan/adapter.rs

@@ -2166,6 +2166,7 @@ impl super::Adapter {
                 force_loop_bounding: true,
                 use_storage_input_output_16: features.contains(wgt::Features::SHADER_F16)
                     && self.phd_features.supports_storage_input_output_16(),
+                fake_missing_bindings: false,
                 // We need to build this separately for each invocation, so just default it out here
                 binding_map: BTreeMap::default(),
                 debug_info: None,

wgpu-hal/src/vulkan/device.rs

@@ -1482,16 +1482,17 @@ impl crate::Device for super::Device {
     ) -> Result<super::BindGroupLayout, crate::DeviceError> {
         // Iterate through the entries and accumulate our Vulkan
         // DescriptorSetLayoutBindings and DescriptorBindingFlags, as well as
-        // the list of which bindings are binding arrays, and our descriptor
-        // counts.
+        // our binding map and our descriptor counts.
         // Note: not bothering with on stack arrays here as it's low frequency
         let mut vk_bindings = Vec::new();
         let mut binding_flags = Vec::new();
-        let mut binding_arrays = Vec::new();
+        let mut binding_map = Vec::new();
+        let mut next_binding = 0;
+        let mut contains_binding_arrays = false;
         let mut desc_count = gpu_descriptor::DescriptorTotalCount::default();
-        for (i, entry) in desc.entries.iter().enumerate() {
-            if let Some(count) = entry.count {
-                binding_arrays.push((i as u32, count))
+        for entry in desc.entries {
+            if entry.count.is_some() {
+                contains_binding_arrays = true;
             }
 
             let partially_bound = desc
@@ -1510,7 +1511,7 @@ impl crate::Device for super::Device {
                 wgt::BindingType::ExternalTexture => unimplemented!(),
                 _ => {
                     vk_bindings.push(vk::DescriptorSetLayoutBinding {
-                        binding: entry.binding,
+                        binding: next_binding,
                         descriptor_type: conv::map_binding_type(entry.ty),
                         descriptor_count: count,
                         stage_flags: conv::map_shader_stage(entry.visibility),
@@ -1518,6 +1519,14 @@ impl crate::Device for super::Device {
                         _marker: Default::default(),
                     });
                     binding_flags.push(flags);
+                    binding_map.push((
+                        entry.binding,
+                        super::BindingInfo {
+                            binding: next_binding,
+                            binding_array_size: entry.count,
+                        },
+                    ));
+                    next_binding += 1;
                 }
             }
 
@@ -1560,7 +1569,7 @@ impl crate::Device for super::Device {
 
         let vk_info = vk::DescriptorSetLayoutCreateInfo::default()
             .bindings(&vk_bindings)
-            .flags(if !binding_arrays.is_empty() {
+            .flags(if contains_binding_arrays {
                 vk::DescriptorSetLayoutCreateFlags::UPDATE_AFTER_BIND_POOL
             } else {
                 vk::DescriptorSetLayoutCreateFlags::empty()
@@ -1588,7 +1597,8 @@ impl crate::Device for super::Device {
             raw,
             desc_count,
             entries: desc.entries.into(),
-            binding_arrays,
+            binding_map,
+            contains_binding_arrays,
         })
     }
     unsafe fn destroy_bind_group_layout(&self, bg_layout: super::BindGroupLayout) {
@@ -1640,27 +1650,25 @@ impl crate::Device for super::Device {
             unsafe { self.shared.set_object_name(raw, label) };
         }
 
-        let mut binding_arrays = BTreeMap::new();
+        let mut binding_map = BTreeMap::new();
         for (group, &layout) in desc.bind_group_layouts.iter().enumerate() {
-            for &(binding, binding_array_size) in &layout.binding_arrays {
-                binding_arrays.insert(
+            for &(binding, binding_info) in &layout.binding_map {
+                binding_map.insert(
                     naga::ResourceBinding {
                         group: group as u32,
                         binding,
                     },
                     naga::back::spv::BindingInfo {
-                        binding_array_size: Some(binding_array_size.get()),
+                        descriptor_set: group as u32,
+                        binding: binding_info.binding,
+                        binding_array_size: binding_info.binding_array_size.map(NonZeroU32::get),
                     },
                 );
             }
         }
 
         self.counters.pipeline_layouts.add(1);
-
-        Ok(super::PipelineLayout {
-            raw,
-            binding_arrays,
-        })
+        Ok(super::PipelineLayout { raw, binding_map })
     }
     unsafe fn destroy_pipeline_layout(&self, pipeline_layout: super::PipelineLayout) {
         unsafe {
@@ -1682,9 +1690,7 @@ impl crate::Device for super::Device {
             super::AccelerationStructure,
         >,
     ) -> Result<super::BindGroup, crate::DeviceError> {
-        let contains_binding_arrays = !desc.layout.binding_arrays.is_empty();
-
-        let desc_set_layout_flags = if contains_binding_arrays {
+        let desc_set_layout_flags = if desc.layout.contains_binding_arrays {
             gpu_descriptor::DescriptorSetLayoutCreateFlags::UPDATE_AFTER_BIND
         } else {
             gpu_descriptor::DescriptorSetLayoutCreateFlags::empty()
@@ -1780,6 +1786,7 @@ impl crate::Device for super::Device {
                 .expect("internal error: no layout entry found with binding slot");
             (layout, entry)
         });
+        let mut next_binding = 0;
         for (layout, entry) in layout_and_entry_iter {
             let write = vk::WriteDescriptorSet::default().dst_set(*set.raw());
 
@@ -1794,10 +1801,11 @@ impl crate::Device for super::Device {
                         ));
                     writes.push(
                         write
-                            .dst_binding(entry.binding)
+                            .dst_binding(next_binding)
                             .descriptor_type(conv::map_binding_type(layout.ty))
                             .image_info(local_image_infos),
                     );
+                    next_binding += 1;
                 }
                 wgt::BindingType::Texture { .. } | wgt::BindingType::StorageTexture { .. } => {
                     let start = entry.resource_index;
@@ -1815,10 +1823,11 @@ impl crate::Device for super::Device {
                         ));
                     writes.push(
                         write
-                            .dst_binding(entry.binding)
+                            .dst_binding(next_binding)
                             .descriptor_type(conv::map_binding_type(layout.ty))
                             .image_info(local_image_infos),
                     );
+                    next_binding += 1;
                 }
                 wgt::BindingType::Buffer { .. } => {
                     let start = entry.resource_index;
@@ -1837,10 +1846,11 @@ impl crate::Device for super::Device {
                         ));
                     writes.push(
                         write
-                            .dst_binding(entry.binding)
+                            .dst_binding(next_binding)
                             .descriptor_type(conv::map_binding_type(layout.ty))
                             .buffer_info(local_buffer_infos),
                     );
+                    next_binding += 1;
                 }
                 wgt::BindingType::AccelerationStructure { .. } => {
                     let start = entry.resource_index;
@@ -1867,11 +1877,12 @@ impl crate::Device for super::Device {
 
                     writes.push(
                         write
-                            .dst_binding(entry.binding)
+                            .dst_binding(next_binding)
                             .descriptor_type(conv::map_binding_type(layout.ty))
                             .descriptor_count(entry.count)
                             .push_next(local_acceleration_structure_infos),
                     );
+                    next_binding += 1;
                 }
                 wgt::BindingType::ExternalTexture => unimplemented!(),
             }
@@ -2033,7 +2044,7 @@ impl crate::Device for super::Device {
                 compiled_vs = Some(self.compile_stage(
                     vertex_stage,
                     naga::ShaderStage::Vertex,
-                    &desc.layout.binding_arrays,
+                    &desc.layout.binding_map,
                 )?);
                 stages.push(compiled_vs.as_ref().unwrap().create_info);
             }
@@ -2045,14 +2056,14 @@ impl crate::Device for super::Device {
                     compiled_ts = Some(self.compile_stage(
                         t,
                         naga::ShaderStage::Task,
-                        &desc.layout.binding_arrays,
+                        &desc.layout.binding_map,
                     )?);
                     stages.push(compiled_ts.as_ref().unwrap().create_info);
                 }
                 compiled_ms = Some(self.compile_stage(
                     mesh_stage,
                     naga::ShaderStage::Mesh,
-                    &desc.layout.binding_arrays,
+                    &desc.layout.binding_map,
                 )?);
                 stages.push(compiled_ms.as_ref().unwrap().create_info);
             }
@@ -2062,7 +2073,7 @@ impl crate::Device for super::Device {
                 let compiled = self.compile_stage(
                     stage,
                     naga::ShaderStage::Fragment,
-                    &desc.layout.binding_arrays,
+                    &desc.layout.binding_map,
                 )?;
                 stages.push(compiled.create_info);
                 Some(compiled)
@@ -2270,7 +2281,7 @@ impl crate::Device for super::Device {
         let compiled = self.compile_stage(
             &desc.stage,
             naga::ShaderStage::Compute,
-            &desc.layout.binding_arrays,
+            &desc.layout.binding_map,
         )?;
 
         let vk_infos = [{

wgpu-hal/src/vulkan/mod.rs

@@ -1001,14 +1001,25 @@ pub struct Sampler {
 
 impl crate::DynSampler for Sampler {}
 
+/// Information about a binding within a specific BindGroupLayout / BindGroup.
+/// This will be used to construct a [`naga::back::spv::BindingInfo`], where
+/// the descriptor set value will be taken from the index of the group.
+#[derive(Copy, Clone, Debug)]
+struct BindingInfo {
+    binding: u32,
+    binding_array_size: Option<NonZeroU32>,
+}
+
 #[derive(Debug)]
 pub struct BindGroupLayout {
     raw: vk::DescriptorSetLayout,
     desc_count: gpu_descriptor::DescriptorTotalCount,
     /// Sorted list of entries.
     entries: Box<[wgt::BindGroupLayoutEntry]>,
-    /// Map of binding index to size,
-    binding_arrays: Vec<(u32, NonZeroU32)>,
+    /// Map of original binding index to remapped binding index and optional
+    /// array size.
+    binding_map: Vec<(u32, BindingInfo)>,
+    contains_binding_arrays: bool,
 }
 
 impl crate::DynBindGroupLayout for BindGroupLayout {}
@@ -1016,7 +1027,7 @@ impl crate::DynBindGroupLayout for BindGroupLayout {}
 #[derive(Debug)]
 pub struct PipelineLayout {
     raw: vk::PipelineLayout,
-    binding_arrays: naga::back::spv::BindingMap,
+    binding_map: naga::back::spv::BindingMap,
 }
 
 impl crate::DynPipelineLayout for PipelineLayout {}