From 5a27bfdcb1e0155f4d469e7a8c48c924b71b968e Mon Sep 17 00:00:00 2001
From: Erich Gubler <erichdongubler@gmail.com>
Date: Tue, 18 Jul 2023 21:55:34 -0400
Subject: [PATCH] fix(limits): properly calculate `max_bindings_per_bind_group`

---
 wgpu-hal/src/auxil/mod.rs      | 77 ++++++++++++++++++++++++++++++++++
 wgpu-hal/src/dx11/adapter.rs   | 27 +++++++++---
 wgpu-hal/src/dx12/adapter.rs   | 49 +++++++++++++++-------
 wgpu-hal/src/gles/adapter.rs   | 20 +++++++--
 wgpu-hal/src/metal/adapter.rs  | 28 ++++++++++---
 wgpu-hal/src/vulkan/adapter.rs | 29 ++++++++++---
 6 files changed, 194 insertions(+), 36 deletions(-)

diff --git a/wgpu-hal/src/auxil/mod.rs b/wgpu-hal/src/auxil/mod.rs
index f0aa6a4a892..2df9c61f75f 100644
--- a/wgpu-hal/src/auxil/mod.rs
+++ b/wgpu-hal/src/auxil/mod.rs
@@ -1,3 +1,5 @@
+use wgt::Limits;
+
 #[cfg(all(any(feature = "dx11", feature = "dx12"), windows))]
 pub(super) mod dxgi;
 
@@ -46,6 +48,81 @@ pub mod db {
 /// offset at some intermediate point, internally, as i32.
 pub const MAX_I32_BINDING_SIZE: u32 = 1 << 31;
 
+/// Per the [WebGPU spec.]:
+///
+/// > **_max shader stages per pipeline_** is `2`, because a `GPURenderPipeline` supports both
+/// > a vertex and fragment shader.
+///
+/// [WebGPU spec.]: https://gpuweb.github.io/gpuweb/#max-shader-stages-per-pipeline
+//#[cfg(not(target_arch = "wasm32"))]
+const MAX_SHADER_STAGES_PER_PIPELINE: u32 = 2;
+
+/// Input for [`max_bindings_per_bind_group`].
+pub(crate) struct MaxBindingsPerBindGroupInput {
+    pub max_sampled_textures_per_shader_stage: u32,
+    pub max_samplers_per_shader_stage: u32,
+    pub max_storage_buffers_per_shader_stage: u32,
+    pub max_storage_textures_per_shader_stage: u32,
+    pub max_uniform_buffers_per_shader_stage: u32,
+}
+
+/// Calculates the maximum bindings per bind group, according to [this formula from the adapter
+/// capabilities guarantees list in the WebGPU spec.]:
+///
+/// > `maxBindingsPerBindGroup` must be must be ≥ (max bindings per shader stage × max shader
+/// > stages per pipeline), where:
+/// >
+/// > - max bindings per shader stage is (`maxSampledTexturesPerShaderStage`
+/// +   `maxSamplersPerShaderStage` + `maxStorageBuffersPerShaderStage`
+/// +   `maxStorageTexturesPerShaderStage` + `maxUniformBuffersPerShaderStage`).
+/// > - max shader stages per pipeline is `2`, because
+/// >   a `[GPURenderPipeline](https://gpuweb.github.io/gpuweb/#gpurenderpipeline)` supports both
+/// >   a vertex and fragment shader.
+///
+/// We choose to interpret the above additions as saturating operations. If, for some reason, the
+/// output of this formula is <= default, it is clamped to the default.
+///
+/// See also from the spec.:
+///
+/// * Documentation for
+///   [`maxBindingsPerBindGroup`](https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxbindingsperbindgroup)
+/// * [4.2.1 Adapter Capability Guarantees](adapter-cap-guarantees)
+///
+/// [adapter-cap-guarantees]: https://gpuweb.github.io/gpuweb/#adapter-capability-guarantees
+pub(crate) fn max_bindings_per_bind_group(input: MaxBindingsPerBindGroupInput) -> u32 {
+    let minimum = Limits::default().max_bindings_per_bind_group;
+
+    let MaxBindingsPerBindGroupInput {
+        max_sampled_textures_per_shader_stage,
+        max_samplers_per_shader_stage,
+        max_storage_buffers_per_shader_stage,
+        max_storage_textures_per_shader_stage,
+        max_uniform_buffers_per_shader_stage,
+    } = input;
+
+    let mut max_bindings_per_bind_group = (max_sampled_textures_per_shader_stage
+        .saturating_add(max_samplers_per_shader_stage)
+        .saturating_add(max_storage_buffers_per_shader_stage)
+        .saturating_add(max_storage_textures_per_shader_stage)
+        .saturating_add(max_uniform_buffers_per_shader_stage))
+    .saturating_mul(MAX_SHADER_STAGES_PER_PIPELINE);
+
+    if max_bindings_per_bind_group < minimum {
+        log::warn!(
+            "`max_bindings_per_bind_group` was < 1000, clamping to 1000 to adhere to WebGPU spec."
+        );
+        max_bindings_per_bind_group = minimum;
+    }
+
+    if max_bindings_per_bind_group > minimum {
+        // Yes, we're throwing away the calculated value! We're clamping to this value right now
+        // because we want to limit exposure to driver bugs, like Vulkan is known to have.
+        max_bindings_per_bind_group = minimum;
+    }
+
+    max_bindings_per_bind_group
+}
+
 pub fn map_naga_stage(stage: naga::ShaderStage) -> wgt::ShaderStages {
     match stage {
         naga::ShaderStage::Vertex => wgt::ShaderStages::VERTEX,
diff --git a/wgpu-hal/src/dx11/adapter.rs b/wgpu-hal/src/dx11/adapter.rs
index a28106a9bbf..6f4ca1f2b33 100644
--- a/wgpu-hal/src/dx11/adapter.rs
+++ b/wgpu-hal/src/dx11/adapter.rs
@@ -2,6 +2,8 @@ use std::num::NonZeroU64;
 
 use winapi::um::{d3d11, d3dcommon};
 
+use crate::auxil::{max_bindings_per_bind_group, MaxBindingsPerBindGroupInput};
+
 impl crate::Adapter<super::Api> for super::Adapter {
     unsafe fn open(
         &self,
@@ -203,6 +205,21 @@ impl super::Adapter {
         let max_compute_workgroups_per_dimension =
             d3d11::D3D11_CS_DISPATCH_MAX_THREAD_GROUPS_PER_DIMENSION;
 
+        let max_sampled_textures_per_shader_stage = max_sampled_textures;
+        let max_samplers_per_shader_stage = max_samplers;
+        let max_storage_buffers_per_shader_stage = max_uavs;
+        let max_storage_textures_per_shader_stage = max_uavs;
+        let max_uniform_buffers_per_shader_stage = max_constant_buffers;
+
+        let max_bindings_per_bind_group =
+            max_bindings_per_bind_group(MaxBindingsPerBindGroupInput {
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
+            });
+
         let limits = wgt::Limits {
             max_texture_dimension_1d: max_texture_dimension_2d,
             max_texture_dimension_2d,
@@ -212,11 +229,11 @@ impl super::Adapter {
             max_bindings_per_bind_group: 65535,
             max_dynamic_uniform_buffers_per_pipeline_layout: max_constant_buffers,
             max_dynamic_storage_buffers_per_pipeline_layout: 0,
-            max_sampled_textures_per_shader_stage: max_sampled_textures,
-            max_samplers_per_shader_stage: max_samplers,
-            max_storage_buffers_per_shader_stage: max_uavs,
-            max_storage_textures_per_shader_stage: max_uavs,
-            max_uniform_buffers_per_shader_stage: max_constant_buffers,
+            max_sampled_textures_per_shader_stage,
+            max_samplers_per_shader_stage,
+            max_storage_buffers_per_shader_stage,
+            max_storage_textures_per_shader_stage,
+            max_uniform_buffers_per_shader_stage,
             max_uniform_buffer_binding_size: 1 << 16,
             max_storage_buffer_binding_size: u32::MAX,
             max_vertex_buffers,
diff --git a/wgpu-hal/src/dx12/adapter.rs b/wgpu-hal/src/dx12/adapter.rs
index 27e8e8e05f8..6e5bfa6982f 100644
--- a/wgpu-hal/src/dx12/adapter.rs
+++ b/wgpu-hal/src/dx12/adapter.rs
@@ -1,5 +1,7 @@
 use crate::{
-    auxil::{self, dxgi::result::HResult as _},
+    auxil::{
+        self, dxgi::result::HResult as _, max_bindings_per_bind_group, MaxBindingsPerBindGroupInput,
+    },
     dx12::SurfaceTarget,
 };
 use std::{mem, ptr, sync::Arc, thread};
@@ -277,6 +279,30 @@ impl super::Adapter {
 
         let base = wgt::Limits::default();
 
+        let max_sampled_textures_per_shader_stage = match options.ResourceBindingTier {
+            d3d12_ty::D3D12_RESOURCE_BINDING_TIER_1 => 128,
+            _ => full_heap_count,
+        };
+        let max_samplers_per_shader_stage = match options.ResourceBindingTier {
+            d3d12_ty::D3D12_RESOURCE_BINDING_TIER_1 => 16,
+            _ => d3d12_ty::D3D12_MAX_SHADER_VISIBLE_SAMPLER_HEAP_SIZE,
+        };
+        // these both account towards `uav_count`, but we can't express the limit as as sum
+        // of the two, so we divide it by 4 to account for the worst case scenario
+        // (2 shader stages, with both using 16 storage textures and 16 storage buffers)
+        let max_storage_buffers_per_shader_stage = uav_count / 4;
+        let max_storage_textures_per_shader_stage = uav_count / 4;
+        let max_uniform_buffers_per_shader_stage = full_heap_count;
+
+        let max_bindings_per_bind_group =
+            max_bindings_per_bind_group(MaxBindingsPerBindGroupInput {
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
+            });
+
         Some(crate::ExposedAdapter {
             adapter: super::Adapter {
                 raw: adapter,
@@ -297,26 +323,17 @@ impl super::Adapter {
                     max_texture_dimension_3d: d3d12_ty::D3D12_REQ_TEXTURE3D_U_V_OR_W_DIMENSION,
                     max_texture_array_layers: d3d12_ty::D3D12_REQ_TEXTURE2D_ARRAY_AXIS_DIMENSION,
                     max_bind_groups: crate::MAX_BIND_GROUPS as u32,
-                    max_bindings_per_bind_group: 65535,
+                    max_bindings_per_bind_group,
                     // dynamic offsets take a root constant, so we expose the minimum here
                     max_dynamic_uniform_buffers_per_pipeline_layout: base
                         .max_dynamic_uniform_buffers_per_pipeline_layout,
                     max_dynamic_storage_buffers_per_pipeline_layout: base
                         .max_dynamic_storage_buffers_per_pipeline_layout,
-                    max_sampled_textures_per_shader_stage: match options.ResourceBindingTier {
-                        d3d12_ty::D3D12_RESOURCE_BINDING_TIER_1 => 128,
-                        _ => full_heap_count,
-                    },
-                    max_samplers_per_shader_stage: match options.ResourceBindingTier {
-                        d3d12_ty::D3D12_RESOURCE_BINDING_TIER_1 => 16,
-                        _ => d3d12_ty::D3D12_MAX_SHADER_VISIBLE_SAMPLER_HEAP_SIZE,
-                    },
-                    // these both account towards `uav_count`, but we can't express the limit as as sum
-                    // of the two, so we divide it by 4 to account for the worst case scenario
-                    // (2 shader stages, with both using 16 storage textures and 16 storage buffers)
-                    max_storage_buffers_per_shader_stage: uav_count / 4,
-                    max_storage_textures_per_shader_stage: uav_count / 4,
-                    max_uniform_buffers_per_shader_stage: full_heap_count,
+                    max_sampled_textures_per_shader_stage,
+                    max_samplers_per_shader_stage,
+                    max_storage_buffers_per_shader_stage,
+                    max_storage_textures_per_shader_stage,
+                    max_uniform_buffers_per_shader_stage,
                     max_uniform_buffer_binding_size:
                         d3d12_ty::D3D12_REQ_CONSTANT_BUFFER_ELEMENT_COUNT * 16,
                     max_storage_buffer_binding_size: crate::auxil::MAX_I32_BINDING_SIZE,
diff --git a/wgpu-hal/src/gles/adapter.rs b/wgpu-hal/src/gles/adapter.rs
index 5594dfa237b..8effca4be86 100644
--- a/wgpu-hal/src/gles/adapter.rs
+++ b/wgpu-hal/src/gles/adapter.rs
@@ -182,6 +182,8 @@ impl super::Adapter {
     pub(super) unsafe fn expose(
         context: super::AdapterContext,
     ) -> Option<crate::ExposedAdapter<super::Api>> {
+        use crate::auxil::{max_bindings_per_bind_group, MaxBindingsPerBindGroupInput};
+
         let gl = context.lock();
         let extensions = gl.supported_extensions();
 
@@ -499,6 +501,18 @@ impl super::Adapter {
             0
         };
 
+        let max_sampled_textures_per_shader_stage = super::MAX_TEXTURE_SLOTS as u32;
+        let max_samplers_per_shader_stage = super::MAX_SAMPLERS as u32;
+
+        let max_bindings_per_bind_group =
+            max_bindings_per_bind_group(MaxBindingsPerBindGroupInput {
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
+            });
+
         let limits = wgt::Limits {
             max_texture_dimension_1d: max_texture_size,
             max_texture_dimension_2d: max_texture_size,
@@ -507,11 +521,11 @@ impl super::Adapter {
                 gl.get_parameter_i32(glow::MAX_ARRAY_TEXTURE_LAYERS)
             } as u32,
             max_bind_groups: crate::MAX_BIND_GROUPS as u32,
-            max_bindings_per_bind_group: 65535,
+            max_bindings_per_bind_group,
             max_dynamic_uniform_buffers_per_pipeline_layout: max_uniform_buffers_per_shader_stage,
             max_dynamic_storage_buffers_per_pipeline_layout: max_storage_buffers_per_shader_stage,
-            max_sampled_textures_per_shader_stage: super::MAX_TEXTURE_SLOTS as u32,
-            max_samplers_per_shader_stage: super::MAX_SAMPLERS as u32,
+            max_sampled_textures_per_shader_stage,
+            max_samplers_per_shader_stage,
             max_storage_buffers_per_shader_stage,
             max_storage_textures_per_shader_stage,
             max_uniform_buffers_per_shader_stage,
diff --git a/wgpu-hal/src/metal/adapter.rs b/wgpu-hal/src/metal/adapter.rs
index c5e6316c195..5059932ce83 100644
--- a/wgpu-hal/src/metal/adapter.rs
+++ b/wgpu-hal/src/metal/adapter.rs
@@ -1,3 +1,4 @@
+use crate::auxil::{max_bindings_per_bind_group, MaxBindingsPerBindGroupInput};
 use metal::{MTLFeatureSet, MTLGPUFamily, MTLLanguageVersion, MTLReadWriteTextureTier};
 use objc::{class, msg_send, sel, sel_impl};
 use parking_lot::Mutex;
@@ -829,6 +830,21 @@ impl super::PrivateCapabilities {
             .flags
             .set(wgt::DownlevelFlags::ANISOTROPIC_FILTERING, true);
 
+        let max_sampled_textures_per_shader_stage = self.max_textures_per_stage;
+        let max_samplers_per_shader_stage = self.max_samplers_per_stage;
+        let max_storage_buffers_per_shader_stage = self.max_buffers_per_stage;
+        let max_storage_textures_per_shader_stage = self.max_textures_per_stage;
+        let max_uniform_buffers_per_shader_stage = self.max_buffers_per_stage;
+
+        let max_bindings_per_bind_group =
+            max_bindings_per_bind_group(MaxBindingsPerBindGroupInput {
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
+            });
+
         let base = wgt::Limits::default();
         crate::Capabilities {
             limits: wgt::Limits {
@@ -837,16 +853,16 @@ impl super::PrivateCapabilities {
                 max_texture_dimension_3d: self.max_texture_3d_size as u32,
                 max_texture_array_layers: self.max_texture_layers as u32,
                 max_bind_groups: 8,
-                max_bindings_per_bind_group: 65535,
+                max_bindings_per_bind_group,
                 max_dynamic_uniform_buffers_per_pipeline_layout: base
                     .max_dynamic_uniform_buffers_per_pipeline_layout,
                 max_dynamic_storage_buffers_per_pipeline_layout: base
                     .max_dynamic_storage_buffers_per_pipeline_layout,
-                max_sampled_textures_per_shader_stage: self.max_textures_per_stage,
-                max_samplers_per_shader_stage: self.max_samplers_per_stage,
-                max_storage_buffers_per_shader_stage: self.max_buffers_per_stage,
-                max_storage_textures_per_shader_stage: self.max_textures_per_stage,
-                max_uniform_buffers_per_shader_stage: self.max_buffers_per_stage,
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
                 max_uniform_buffer_binding_size: self.max_buffer_size.min(!0u32 as u64) as u32,
                 max_storage_buffer_binding_size: self.max_buffer_size.min(!0u32 as u64) as u32,
                 max_vertex_buffers: self.max_vertex_buffers,
diff --git a/wgpu-hal/src/vulkan/adapter.rs b/wgpu-hal/src/vulkan/adapter.rs
index d23aca76a35..1253d29f811 100644
--- a/wgpu-hal/src/vulkan/adapter.rs
+++ b/wgpu-hal/src/vulkan/adapter.rs
@@ -1,3 +1,5 @@
+use crate::auxil::{max_bindings_per_bind_group, MaxBindingsPerBindGroupInput};
+
 use super::conv;
 
 use ash::{extensions::khr, vk};
@@ -712,6 +714,21 @@ impl PhysicalDeviceCapabilities {
                 u64::MAX
             };
 
+        let max_sampled_textures_per_shader_stage = limits.max_per_stage_descriptor_sampled_images;
+        let max_samplers_per_shader_stage = limits.max_per_stage_descriptor_samplers;
+        let max_storage_buffers_per_shader_stage = limits.max_per_stage_descriptor_storage_buffers;
+        let max_storage_textures_per_shader_stage = limits.max_per_stage_descriptor_storage_images;
+        let max_uniform_buffers_per_shader_stage = limits.max_per_stage_descriptor_uniform_buffers;
+
+        let max_bindings_per_bind_group =
+            max_bindings_per_bind_group(MaxBindingsPerBindGroupInput {
+                max_sampled_textures_per_shader_stage,
+                max_samplers_per_shader_stage,
+                max_storage_buffers_per_shader_stage,
+                max_storage_textures_per_shader_stage,
+                max_uniform_buffers_per_shader_stage,
+            });
+
         wgt::Limits {
             max_texture_dimension_1d: limits.max_image_dimension1_d,
             max_texture_dimension_2d: limits.max_image_dimension2_d,
@@ -720,16 +737,16 @@ impl PhysicalDeviceCapabilities {
             max_bind_groups: limits
                 .max_bound_descriptor_sets
                 .min(crate::MAX_BIND_GROUPS as u32),
-            max_bindings_per_bind_group: wgt::Limits::default().max_bindings_per_bind_group,
+            max_bindings_per_bind_group,
             max_dynamic_uniform_buffers_per_pipeline_layout: limits
                 .max_descriptor_set_uniform_buffers_dynamic,
             max_dynamic_storage_buffers_per_pipeline_layout: limits
                 .max_descriptor_set_storage_buffers_dynamic,
-            max_sampled_textures_per_shader_stage: limits.max_per_stage_descriptor_sampled_images,
-            max_samplers_per_shader_stage: limits.max_per_stage_descriptor_samplers,
-            max_storage_buffers_per_shader_stage: limits.max_per_stage_descriptor_storage_buffers,
-            max_storage_textures_per_shader_stage: limits.max_per_stage_descriptor_storage_images,
-            max_uniform_buffers_per_shader_stage: limits.max_per_stage_descriptor_uniform_buffers,
+            max_sampled_textures_per_shader_stage,
+            max_samplers_per_shader_stage,
+            max_storage_buffers_per_shader_stage,
+            max_storage_textures_per_shader_stage,
+            max_uniform_buffers_per_shader_stage,
             max_uniform_buffer_binding_size: limits
                 .max_uniform_buffer_range
                 .min(crate::auxil::MAX_I32_BINDING_SIZE),