Make Force Loop Bounding Optional

Co-authored-by: rudderbucky <anandkwork7@gmail.com>

Author: cwfitzgerald

CHANGELOG.md

@@ -53,6 +53,20 @@ By @cwfitzgerald in [#6619](https://github.com/gfx-rs/wgpu/pull/6619).
 A regression intoduced in 23.0.0 caused lifetimes of render and compute passes to be incorrectly enforced. While this is not
 a soundness issue, the intent is to move an error from runtime to compile time. This issue has been fixed and restored to the 22.0.0 behavior.
 
+### `Device::create_shader_module_unchecked` Now Has Configuration Options
+
+This allows you to customize which exact checks are omitted so that you can get the correct balance of performance and safety for your use case. Calling the function is still unsafe, but now can be used to skip certain checks only on certain builds.
+
+This also allows users to disable the workarounds in the `msl-out` backend to prevent the compiler from optimizing infinite loops. This can have a big impact on performance, but is not recommended for untrusted shaders.
+
+```diff
+let desc: ShaderModuleDescriptor = include_wgsl!(...)
+- let module = unsafe { device.create_shader_module_unchecked(desc) };
++ let module = unsafe { device.create_shader_module_unchecked(desc, wgpu::ShaderRuntimeChecks::unchecked()) };
+```
+
+By @cwfitzgerald and @rudderbucky in [#6662](https://github.com/gfx-rs/wgpu/pull/6662).
+
 ### The `diagnostic(…);` directive is now supported in WGSL
 
 Naga now parses `diagnostic(…);` directives according to the WGSL spec. This allows users to control certain lints, similar to Rust's `allow`, `warn`, and `deny` attributes. For example, in standard WGSL (but, notably, not Naga yet—see <https://github.com/gfx-rs/wgpu/issues/4369>) this snippet would emit a uniformity error:

deno_webgpu/shader.rs

@@ -42,7 +42,7 @@ pub fn op_webgpu_create_shader_module(
 
     let descriptor = wgpu_core::pipeline::ShaderModuleDescriptor {
         label: Some(label),
-        shader_bound_checks: wgpu_types::ShaderBoundChecks::default(),
+        runtime_checks: wgpu_types::ShaderRuntimeChecks::default(),
     };
 
     gfx_put!(instance.device_create_shader_module(

naga/src/back/msl/mod.rs

@@ -211,6 +211,9 @@ pub struct Options {
     pub bounds_check_policies: index::BoundsCheckPolicies,
     /// Should workgroup variables be zero initialized (by polyfilling)?
     pub zero_initialize_workgroup_memory: bool,
+    /// If set, loops will have code injected into them, forcing the compiler
+    /// to think the number of iterations is bounded.
+    pub force_loop_bounding: bool,
 }
 
 impl Default for Options {
@@ -223,6 +226,7 @@ impl Default for Options {
             fake_missing_bindings: true,
             bounds_check_policies: index::BoundsCheckPolicies::default(),
             zero_initialize_workgroup_memory: true,
+            force_loop_bounding: true,
         }
     }
 }

naga/src/back/msl/writer.rs

@@ -600,6 +600,8 @@ struct ExpressionContext<'a> {
     /// accesses. These may need to be cached in temporary variables. See
     /// `index::find_checked_indexes` for details.
     guarded_indices: HandleSet<crate::Expression>,
+    /// See [`emit_force_bounded_loop_macro`] for details.
+    force_loop_bounding: bool,
 }
 
 impl<'a> ExpressionContext<'a> {
@@ -3066,13 +3068,15 @@ impl<W: Write> Writer<W> {
                         writeln!(self.out, "{level}while(true) {{",)?;
                     }
                     self.put_block(level.next(), body, context)?;
-                    self.emit_force_bounded_loop_macro()?;
-                    writeln!(
-                        self.out,
-                        "{}{}",
-                        level.next(),
-                        self.force_bounded_loop_macro_name
-                    )?;
+                    if context.expression.force_loop_bounding {
+                        self.emit_force_bounded_loop_macro()?;
+                        writeln!(
+                            self.out,
+                            "{}{}",
+                            level.next(),
+                            self.force_bounded_loop_macro_name
+                        )?;
+                    }
                     writeln!(self.out, "{level}}}")?;
                 }
                 crate::Statement::Break => {
@@ -4880,6 +4884,7 @@ template <typename A>
                     module,
                     mod_info,
                     pipeline_options,
+                    force_loop_bounding: options.force_loop_bounding,
                 },
                 result_struct: None,
             };
@@ -5780,6 +5785,7 @@ template <typename A>
                     module,
                     mod_info,
                     pipeline_options,
+                    force_loop_bounding: options.force_loop_bounding,
                 },
                 result_struct: Some(&stage_out_name),
             };

naga/tests/in/interface.param.ron

@@ -28,7 +28,7 @@
 	),
 	msl_pipeline: (
 		allow_and_force_point_size: true,
-    vertex_pulling_transform: false,
-    vertex_buffer_mappings: [],
+		vertex_pulling_transform: false,
+		vertex_buffer_mappings: [],
 	),
 )

wgpu-core/src/device/resource.rs

@@ -1549,7 +1549,7 @@ impl Device {
         });
         let hal_desc = hal::ShaderModuleDescriptor {
             label: desc.label.to_hal(self.instance_flags),
-            runtime_checks: desc.shader_bound_checks.runtime_checks(),
+            runtime_checks: desc.runtime_checks,
         };
         let raw = match unsafe { self.raw().create_shader_module(&hal_desc, hal_shader) } {
             Ok(raw) => raw,
@@ -1589,7 +1589,7 @@ impl Device {
         self.require_features(wgt::Features::SPIRV_SHADER_PASSTHROUGH)?;
         let hal_desc = hal::ShaderModuleDescriptor {
             label: desc.label.to_hal(self.instance_flags),
-            runtime_checks: desc.shader_bound_checks.runtime_checks(),
+            runtime_checks: desc.runtime_checks,
         };
         let hal_shader = hal::ShaderInput::SpirV(source);
         let raw = match unsafe { self.raw().create_shader_module(&hal_desc, hal_shader) } {

wgpu-core/src/indirect_validation.rs

@@ -123,7 +123,7 @@ impl IndirectValidation {
         });
         let hal_desc = hal::ShaderModuleDescriptor {
             label: None,
-            runtime_checks: false,
+            runtime_checks: unsafe { wgt::ShaderRuntimeChecks::unchecked() },
         };
         let module =
             unsafe { device.create_shader_module(&hal_desc, hal_shader) }.map_err(|error| {

wgpu-core/src/pipeline.rs

@@ -41,7 +41,7 @@ pub enum ShaderModuleSource<'a> {
 pub struct ShaderModuleDescriptor<'a> {
     pub label: Label<'a>,
     #[cfg_attr(feature = "serde", serde(default))]
-    pub shader_bound_checks: wgt::ShaderBoundChecks,
+    pub runtime_checks: wgt::ShaderRuntimeChecks,
 }
 
 #[derive(Debug)]

wgpu-hal/examples/halmark/main.rs

@@ -177,7 +177,7 @@ impl<A: hal::Api> Example<A> {
         };
         let shader_desc = hal::ShaderModuleDescriptor {
             label: None,
-            runtime_checks: false,
+            runtime_checks: wgt::ShaderRuntimeChecks::checked(),
         };
         let shader = unsafe {
             device

wgpu-hal/examples/ray-traced-triangle/main.rs

@@ -371,7 +371,7 @@ impl<A: hal::Api> Example<A> {
         };
         let shader_desc = hal::ShaderModuleDescriptor {
             label: None,
-            runtime_checks: false,
+            runtime_checks: wgt::ShaderRuntimeChecks::checked(),
         };
         let shader_module = unsafe {
             device

wgpu-hal/src/dx12/device.rs

@@ -273,12 +273,12 @@ impl super::Device {
 
         let needs_temp_options = stage.zero_initialize_workgroup_memory
             != layout.naga_options.zero_initialize_workgroup_memory
-            || stage.module.runtime_checks != layout.naga_options.restrict_indexing;
+            || stage.module.runtime_checks.bounds_checks() != layout.naga_options.restrict_indexing;
         let mut temp_options;
         let naga_options = if needs_temp_options {
             temp_options = layout.naga_options.clone();
             temp_options.zero_initialize_workgroup_memory = stage.zero_initialize_workgroup_memory;
-            temp_options.restrict_indexing = stage.module.runtime_checks;
+            temp_options.restrict_indexing = stage.module.runtime_checks.bounds_checks();
             &temp_options
         } else {
             &layout.naga_options

wgpu-hal/src/dx12/mod.rs

@@ -967,7 +967,7 @@ impl crate::DynPipelineLayout for PipelineLayout {}
 pub struct ShaderModule {
     naga: crate::NagaShader,
     raw_name: Option<ffi::CString>,
-    runtime_checks: bool,
+    runtime_checks: wgt::ShaderRuntimeChecks,
 }
 
 impl crate::DynShaderModule for ShaderModule {}

wgpu-hal/src/lib.rs

@@ -2103,6 +2103,10 @@ pub enum ShaderInput<'a> {
 pub struct ShaderModuleDescriptor<'a> {
     pub label: Label<'a>,
 
+    /// # Safety
+    ///
+    /// ## `runtime_checks.bounds_checks()`:
+    ///
     /// Enforce bounds checks in shaders, even if the underlying driver doesn't
     /// support doing so natively.
     ///
@@ -2122,7 +2126,15 @@ pub struct ShaderModuleDescriptor<'a> {
     /// making such deferred initialization visible to the application.
     ///
     /// [ar]: struct.BufferBinding.html#accessible-region
-    pub runtime_checks: bool,
+    ///
+    /// ## `runtime_checks.force_loop_bounding()`:
+    ///
+    /// If false, the caller MUST ensure that all passed shaders do not contain any infinite loops.
+    ///
+    /// If it does, backend compilers MAY treat such a loop as unreachable code and draw
+    /// conclusions about other safety-critical code paths. This option SHOULD NOT be disabled
+    /// when running untrusted code.
+    pub runtime_checks: wgt::ShaderRuntimeChecks,
 }
 
 #[derive(Debug, Clone)]

wgpu-hal/src/metal/device.rs

@@ -117,7 +117,7 @@ impl super::Device {
 
         let ep_resources = &layout.per_stage_map[naga_stage];
 
-        let bounds_check_policy = if stage.module.runtime_checks {
+        let bounds_check_policy = if stage.module.bounds_checks.bounds_checks() {
             naga::proc::BoundsCheckPolicy::Restrict
         } else {
             naga::proc::BoundsCheckPolicy::Unchecked
@@ -151,6 +151,7 @@ impl super::Device {
                 binding_array: naga::proc::BoundsCheckPolicy::Unchecked,
             },
             zero_initialize_workgroup_memory: stage.zero_initialize_workgroup_memory,
+            force_loop_bounding: stage.module.bounds_checks.force_loop_bounding(),
         };
 
         let pipeline_options = naga::back::msl::PipelineOptions {
@@ -888,7 +889,7 @@ impl crate::Device for super::Device {
         match shader {
             crate::ShaderInput::Naga(naga) => Ok(super::ShaderModule {
                 naga,
-                runtime_checks: desc.runtime_checks,
+                bounds_checks: desc.runtime_checks,
             }),
             crate::ShaderInput::SpirV(_) => {
                 panic!("SPIRV_SHADER_PASSTHROUGH is not enabled for this backend")

wgpu-hal/src/metal/mod.rs

@@ -726,7 +726,7 @@ unsafe impl Sync for BindGroup {}
 #[derive(Debug)]
 pub struct ShaderModule {
     naga: crate::NagaShader,
-    runtime_checks: bool,
+    bounds_checks: wgt::ShaderRuntimeChecks,
 }
 
 impl crate::DynShaderModule for ShaderModule {}

wgpu-hal/src/vulkan/device.rs

@@ -897,14 +897,14 @@ impl super::Device {
                     entry_point: stage.entry_point.to_string(),
                     shader_stage: naga_stage,
                 };
-                let needs_temp_options = !runtime_checks
+                let needs_temp_options = !runtime_checks.bounds_checks()
                     || !binding_map.is_empty()
                     || naga_shader.debug_source.is_some()
                     || !stage.zero_initialize_workgroup_memory;
                 let mut temp_options;
                 let options = if needs_temp_options {
                     temp_options = self.naga_options.clone();
-                    if !runtime_checks {
+                    if !runtime_checks.bounds_checks() {
                         temp_options.bounds_check_policies = naga::proc::BoundsCheckPolicies {
                             index: naga::proc::BoundsCheckPolicy::Unchecked,
                             buffer: naga::proc::BoundsCheckPolicy::Unchecked,
@@ -1813,7 +1813,7 @@ impl crate::Device for super::Device {
                             file_name: d.file_name.as_ref().as_ref(),
                             language: naga::back::spv::SourceLanguage::WGSL,
                         });
-                if !desc.runtime_checks {
+                if !desc.runtime_checks.bounds_checks() {
                     naga_options.bounds_check_policies = naga::proc::BoundsCheckPolicies {
                         index: naga::proc::BoundsCheckPolicy::Unchecked,
                         buffer: naga::proc::BoundsCheckPolicy::Unchecked,

wgpu-hal/src/vulkan/mod.rs

@@ -974,7 +974,7 @@ pub enum ShaderModule {
     Raw(vk::ShaderModule),
     Intermediate {
         naga_shader: crate::NagaShader,
-        runtime_checks: bool,
+        runtime_checks: wgt::ShaderRuntimeChecks,
     },
 }
 

wgpu-types/src/lib.rs

@@ -7441,22 +7441,47 @@ impl DispatchIndirectArgs {
 }
 
 /// Describes how shader bound checks should be performed.
-#[derive(Clone, Debug)]
+#[derive(Copy, Clone, Debug)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub struct ShaderBoundChecks {
-    runtime_checks: bool,
+pub struct ShaderRuntimeChecks {
+    bounds_checks: bool,
+    force_loop_bounding: bool,
 }
 
-impl ShaderBoundChecks {
-    /// Creates a new configuration where the shader is bound checked.
+impl ShaderRuntimeChecks {
+    /// Creates a new configuration where the shader is fully checked.
+    #[must_use]
+    pub fn checked() -> Self {
+        unsafe { Self::all(true) }
+    }
+
+    /// Creates a new configuration where none of the checks are performed.
+    ///
+    /// # Safety
+    ///
+    /// See the documentation for the `set_*` methods for the safety requirements
+    /// of each sub-configuration.
+    #[must_use]
+    pub unsafe fn unchecked() -> Self {
+        unsafe { Self::all(false) }
+    }
+
+    /// Creates a new configuration where all checks are enabled or disabled. To safely
+    /// create a configuration with all checks enabled, use [`ShaderRuntimeChecks::checked`].
+    ///
+    /// # Safety
+    ///
+    /// See the documentation for the `set_*` methods for the safety requirements
+    /// of each sub-configuration.
     #[must_use]
-    pub fn new() -> Self {
-        ShaderBoundChecks {
-            runtime_checks: true,
+    pub unsafe fn all(all_checks: bool) -> Self {
+        Self {
+            bounds_checks: all_checks,
+            force_loop_bounding: all_checks,
         }
     }
 
-    /// Creates a new configuration where the shader isn't bound checked.
+    /// Sets the enablement of runtime bounds checks.
     ///
     /// # Safety
     ///
@@ -7469,22 +7494,46 @@ impl ShaderBoundChecks {
     /// configuration with ill-behaved shaders could expose uninitialized GPU
     /// memory contents to the application.
     #[must_use]
-    pub unsafe fn unchecked() -> Self {
-        ShaderBoundChecks {
-            runtime_checks: false,
+    pub unsafe fn set_bounds_checks(self, runtime_checks: bool) -> Self {
+        Self {
+            bounds_checks: runtime_checks,
+            ..self
+        }
+    }
+
+    /// Sets the enablement of loop bounding.
+    ///
+    /// # Safety
+    ///
+    /// If false, the caller MUST ensure that all passed shaders do not contain any infinite loops.
+    ///
+    /// If it does, backend compilers MAY treat such a loop as unreachable code and draw
+    /// conclusions about other safety-critical code paths. This option SHOULD NOT be disabled
+    /// when running untrusted code.
+    #[must_use]
+    pub unsafe fn set_force_loop_bounding(self, force_loop_bounding: bool) -> Self {
+        Self {
+            force_loop_bounding,
+            ..self
         }
     }
 
     /// Query whether runtime bound checks are enabled in this configuration
     #[must_use]
-    pub fn runtime_checks(&self) -> bool {
-        self.runtime_checks
+    pub fn bounds_checks(&self) -> bool {
+        self.bounds_checks
+    }
+
+    /// Query whether loop bounding is enabled in this configuration
+    #[must_use]
+    pub fn force_loop_bounding(&self) -> bool {
+        self.force_loop_bounding
     }
 }
 
-impl Default for ShaderBoundChecks {
+impl Default for ShaderRuntimeChecks {
     fn default() -> Self {
-        Self::new()
+        Self::checked()
     }
 }