Make Force Loop Bounding Optional (#6662)

* Make Force Loop Bounding Optional

Co-authored-by: rudderbucky <[email protected]>

* Deprecate and Rename

---------

Co-authored-by: rudderbucky <[email protected]>

CHANGELOG.md

@@ -53,6 +53,22 @@ By @cwfitzgerald in [#6619](https://github.com/gfx-rs/wgpu/pull/6619).
 A regression introduced 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` Renamed and Now Has Configuration Options
+
+`create_shader_module_unchecked` became `create_shader_module_trusted`.
+
+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_trusted(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 [`Writer::emit_force_bounded_loop_macro`] for details.
+    force_loop_bounding: bool,
 }
 
 impl<'a> ExpressionContext<'a> {
@@ -3068,13 +3070,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 => {
@@ -4885,6 +4889,7 @@ template <typename A>
                     module,
                     mod_info,
                     pipeline_options,
+                    force_loop_bounding: options.force_loop_bounding,
                 },
                 result_struct: None,
             };
@@ -5785,6 +5790,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

@@ -1539,7 +1539,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,
@@ -1579,7 +1579,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: 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

@@ -2101,26 +2101,12 @@ pub enum ShaderInput<'a> {
 pub struct ShaderModuleDescriptor<'a> {
     pub label: Label<'a>,
 
-    /// Enforce bounds checks in shaders, even if the underlying driver doesn't
-    /// support doing so natively.
-    ///
-    /// When this is `true`, `wgpu_hal` promises that shaders can only read or
-    /// write the [accessible region][ar] of a bindgroup's buffer bindings. If
-    /// the underlying graphics platform cannot implement these bounds checks
-    /// itself, `wgpu_hal` will inject bounds checks before presenting the
-    /// shader to the platform.
-    ///
-    /// When this is `false`, `wgpu_hal` only enforces such bounds checks if the
-    /// underlying platform provides a way to do so itself. `wgpu_hal` does not
-    /// itself add any bounds checks to generated shader code.
+    /// # Safety
     ///
-    /// Note that `wgpu_hal` users may try to initialize only those portions of
-    /// buffers that they anticipate might be read from. Passing `false` here
-    /// may allow shaders to see wider regions of the buffers than expected,
-    /// making such deferred initialization visible to the application.
+    /// See the documentation for each flag in [`ShaderRuntimeChecks`][src].
     ///
-    /// [ar]: struct.BufferBinding.html#accessible-region
-    pub runtime_checks: bool,
+    /// [src]: wgt::ShaderRuntimeChecks
+    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

@@ -7447,50 +7447,73 @@ 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 {
+    /// Enforce bounds checks in shaders, even if the underlying driver doesn't
+    /// support doing so natively.
+    ///
+    /// When this is `true`, `wgpu` promises that shaders can only read or
+    /// write the accessible region of a bindgroup's buffer bindings. If
+    /// the underlying graphics platform cannot implement these bounds checks
+    /// itself, `wgpu` will inject bounds checks before presenting the
+    /// shader to the platform.
+    ///
+    /// When this is `false`, `wgpu` only enforces such bounds checks if the
+    /// underlying platform provides a way to do so itself. `wgpu` does not
+    /// itself add any bounds checks to generated shader code.
+    ///
+    /// Note that `wgpu` users may try to initialize only those portions of
+    /// buffers that they anticipate might be read from. Passing `false` here
+    /// may allow shaders to see wider regions of the buffers than expected,
+    /// making such deferred initialization visible to the application.
+    pub bounds_checks: bool,
+    ///
+    /// 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 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 new() -> Self {
-        ShaderBoundChecks {
-            runtime_checks: true,
-        }
+    pub fn checked() -> Self {
+        unsafe { Self::all(true) }
     }
 
-    /// Creates a new configuration where the shader isn't bound checked.
+    /// Creates a new configuration where none of the checks are performed.
     ///
     /// # Safety
     ///
-    /// The caller MUST ensure that all shaders built with this configuration
-    /// don't perform any out of bounds reads or writes.
-    ///
-    /// Note that `wgpu_core`, in particular, initializes only those portions of
-    /// buffers that it expects might be read, and it does not expect contents
-    /// outside the ranges bound in bindgroups to be accessible, so using this
-    /// configuration with ill-behaved shaders could expose uninitialized GPU
-    /// memory contents to the application.
+    /// See the documentation for the `set_*` methods for the safety requirements
+    /// of each sub-configuration.
     #[must_use]
-    pub unsafe fn unchecked() -> Self {
-        ShaderBoundChecks {
-            runtime_checks: false,
-        }
+    pub fn unchecked() -> Self {
+        unsafe { Self::all(false) }
     }
 
-    /// Query whether runtime bound checks are enabled in this configuration
+    /// 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 runtime_checks(&self) -> bool {
-        self.runtime_checks
+    pub unsafe fn all(all_checks: bool) -> Self {
+        Self {
+            bounds_checks: all_checks,
+            force_loop_bounding: all_checks,
+        }
     }
 }
 
-impl Default for ShaderBoundChecks {
+impl Default for ShaderRuntimeChecks {
     fn default() -> Self {
-        Self::new()
+        Self::checked()
     }
 }