Allow configuring whether workgroup memory is zero initialised

CHANGELOG.md

@@ -118,6 +118,16 @@ Bottom level categories:
   - `wgpu::Texture::as_hal` now returns a user-defined type to match the other as_hal functions
 
 - Added support for pipeline-overridable constants. By @teoxoy & @jimblandy in [#5500](https://github.com/gfx-rs/wgpu/pull/5500)
+- Support disabling zero-initialisation of workgroup local memory in compute shaders.
+
+  ```diff
+   ComputePipelineDescriptor {
+  +    zero_initialise_workgroup_memory: Default::default()
+       // ...
+   }
+  ```
+
+  By @DJMcNab in [#5508](https://github.com/gfx-rs/wgpu/pull/5508)
 
 #### GLES
 

deno_webgpu/pipeline.rs

@@ -10,6 +10,7 @@ use serde::Serialize;
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::rc::Rc;
+use wgpu_types::ZeroInitializeWorkgroupMemory;
 
 use super::error::WebGpuError;
 use super::error::WebGpuResult;
@@ -114,6 +115,7 @@ pub fn op_webgpu_create_compute_pipeline(
             entry_point: compute.entry_point.map(Cow::from),
             constants: Cow::Owned(compute.constants),
         },
+        zero_initialise_workgroup_memory: ZeroInitializeWorkgroupMemory::always(),
     };
     let implicit_pipelines = match layout {
         GPUPipelineLayoutOrGPUAutoLayoutMode::Layout(_) => None,

examples/src/boids/mod.rs

@@ -165,6 +165,7 @@ impl crate::framework::Example for Example {
             layout: Some(&compute_pipeline_layout),
             module: &compute_shader,
             entry_point: "main",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 

examples/src/hello_compute/mod.rs

@@ -109,6 +109,7 @@ async fn execute_gpu_inner(
         layout: None,
         module: &cs_module,
         entry_point: "main",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
 

examples/src/hello_synchronization/mod.rs

@@ -103,13 +103,15 @@ async fn execute(
         layout: Some(&pipeline_layout),
         module: &shaders_module,
         entry_point: "patient_main",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
     let hasty_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
         label: None,
         layout: Some(&pipeline_layout),
         module: &shaders_module,
         entry_point: "hasty_main",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
 

examples/src/hello_workgroups/mod.rs

@@ -110,6 +110,7 @@ async fn run() {
         layout: Some(&pipeline_layout),
         module: &shader,
         entry_point: "main",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
 

examples/src/repeated_compute/mod.rs

@@ -245,6 +245,7 @@ impl WgpuContext {
             layout: Some(&pipeline_layout),
             module: &shader,
             entry_point: "main",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 

examples/src/storage_texture/mod.rs

@@ -100,6 +100,7 @@ async fn run(_path: Option<String>) {
         layout: Some(&pipeline_layout),
         module: &shader,
         entry_point: "main",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
 

examples/src/timestamp_queries/mod.rs

@@ -298,6 +298,7 @@ fn compute_pass(
         layout: None,
         module,
         entry_point: "main_cs",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
     let bind_group_layout = compute_pipeline.get_bind_group_layout(0);

player/tests/data/bind-group.ron

@@ -59,6 +59,7 @@
                     entry_point: None,
                     constants: {},
                 ),
+                zero_initialise_workgroup_memory: ZeroInitializeWorkgroupMemory(Always)
             ),
         ),
         Submit(1, [

player/tests/data/pipeline-statistics-query.ron

@@ -32,6 +32,7 @@
                     entry_point: None,
                     constants: {},
                 ),
+                zero_initialise_workgroup_memory: ZeroInitializeWorkgroupMemory(Always)
             ),
         ),
         CreateQuerySet(

player/tests/data/zero-init-buffer.ron

@@ -136,6 +136,7 @@
                     entry_point: None,
                     constants: {},
                 ),
+                zero_initialise_workgroup_memory: ZeroInitializeWorkgroupMemory(Always)
             ),
         ),
         Submit(1, [

player/tests/data/zero-init-texture-binding.ron

@@ -137,6 +137,7 @@
                     entry_point: None,
                     constants: {},
                 ),
+                zero_initialise_workgroup_memory: ZeroInitializeWorkgroupMemory(Always)
             ),
         ),
 

tests/src/image.rs

@@ -369,6 +369,7 @@ fn copy_via_compute(
         layout: Some(&pll),
         module: &sm,
         entry_point: "copy_texture_to_buffer",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     });
 

tests/tests/bgra8unorm_storage.rs

@@ -97,6 +97,7 @@ static BGRA8_UNORM_STORAGE: GpuTestConfiguration = GpuTestConfiguration::new()
             layout: Some(&pl),
             entry_point: "main",
             constants: &Default::default(),
+            zero_initialise_workgroup_memory: Default::default(),
             module: &module,
         });
 

tests/tests/bind_group_layout_dedup.rs

@@ -90,6 +90,7 @@ async fn bgl_dedupe(ctx: TestingContext) {
             layout: Some(&pipeline_layout),
             module: &module,
             entry_point: "no_resources",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         };
 
@@ -219,6 +220,7 @@ fn bgl_dedupe_with_dropped_user_handle(ctx: TestingContext) {
             layout: Some(&pipeline_layout),
             module: &module,
             entry_point: "no_resources",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 
@@ -265,6 +267,7 @@ fn bgl_dedupe_derived(ctx: TestingContext) {
             layout: None,
             module: &module,
             entry_point: "resources",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 
@@ -336,6 +339,7 @@ fn separate_programs_have_incompatible_derived_bgls(ctx: TestingContext) {
         layout: None,
         module: &module,
         entry_point: "resources",
+        zero_initialise_workgroup_memory: Default::default(),
         constants: &Default::default(),
     };
     // Create two pipelines, creating a BG from the second.
@@ -398,6 +402,7 @@ fn derived_bgls_incompatible_with_regular_bgls(ctx: TestingContext) {
             layout: None,
             module: &module,
             entry_point: "resources",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 

tests/tests/buffer.rs

@@ -224,6 +224,7 @@ static MINIMUM_BUFFER_BINDING_SIZE_LAYOUT: GpuTestConfiguration = GpuTestConfigu
                     layout: Some(&pipeline_layout),
                     module: &shader_module,
                     entry_point: "main",
+                    zero_initialise_workgroup_memory: Default::default(),
                     constants: &Default::default(),
                 });
         });
@@ -293,6 +294,7 @@ static MINIMUM_BUFFER_BINDING_SIZE_DISPATCH: GpuTestConfiguration = GpuTestConfi
                 layout: Some(&pipeline_layout),
                 module: &shader_module,
                 entry_point: "main",
+                zero_initialise_workgroup_memory: Default::default(),
                 constants: &Default::default(),
             });
 

tests/tests/device.rs

@@ -499,6 +499,7 @@ static DEVICE_DESTROY_THEN_MORE: GpuTestConfiguration = GpuTestConfiguration::ne
                     layout: None,
                     module: &shader_module,
                     entry_point: "",
+                    zero_initialise_workgroup_memory: Default::default(),
                     constants: &Default::default(),
                 });
         });

tests/tests/partially_bounded_arrays/mod.rs

@@ -69,6 +69,7 @@ static PARTIALLY_BOUNDED_ARRAY: GpuTestConfiguration = GpuTestConfiguration::new
             layout: Some(&pipeline_layout),
             module: &cs_module,
             entry_point: "main",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 

tests/tests/pipeline.rs

@@ -28,6 +28,7 @@ static PIPELINE_DEFAULT_LAYOUT_BAD_MODULE: GpuTestConfiguration = GpuTestConfigu
                     layout: None,
                     module: &module,
                     entry_point: "doesn't exist",
+                    zero_initialise_workgroup_memory: Default::default(),
                     constants: &Default::default(),
                 });
 

tests/tests/push_constants.rs

@@ -103,6 +103,7 @@ async fn partial_update_test(ctx: TestingContext) {
             layout: Some(&pipeline_layout),
             module: &sm,
             entry_point: "main",
+            zero_initialise_workgroup_memory: Default::default(),
             constants: &Default::default(),
         });
 

tests/tests/shader/mod.rs

@@ -307,6 +307,7 @@ async fn shader_input_output_test(
                 layout: Some(&pll),
                 module: &sm,
                 entry_point: "cs_main",
+                zero_initialise_workgroup_memory: Default::default(),
                 constants: &Default::default(),
             });
 

tests/tests/shader/zero_init_workgroup_mem.rs

@@ -87,6 +87,7 @@ static ZERO_INIT_WORKGROUP_MEMORY: GpuTestConfiguration = GpuTestConfiguration::
                 layout: Some(&pll),
                 module: &sm,
                 entry_point: "read",
+                zero_initialise_workgroup_memory: Default::default(),
                 constants: &Default::default(),
             });
 
@@ -97,6 +98,7 @@ static ZERO_INIT_WORKGROUP_MEMORY: GpuTestConfiguration = GpuTestConfiguration::
                 layout: None,
                 module: &sm,
                 entry_point: "write",
+                zero_initialise_workgroup_memory: Default::default(),
                 constants: &Default::default(),
             });
 

wgpu-core/src/device/resource.rs

@@ -2766,6 +2766,7 @@ impl<A: HalApi> Device<A> {
                 entry_point: final_entry_point_name.as_ref(),
                 constants: desc.stage.constants.as_ref(),
             },
+            zero_initialise_workgroup_memory: desc.zero_initialise_workgroup_memory,
         };
 
         let raw = unsafe {

wgpu-core/src/pipeline.rs

@@ -268,6 +268,10 @@ pub struct ComputePipelineDescriptor<'a> {
     pub layout: Option<PipelineLayoutId>,
     /// The compiled compute stage and its entry point.
     pub stage: ProgrammableStageDescriptor<'a>,
+    /// Whether to initialise workgroup scoped memory to have a value of zero.
+    /// In most cases, you should set this to [`wgt::ZeroInitializeWorkgroupMemory::always()`],
+    /// which is the default value.
+    pub zero_initialise_workgroup_memory: wgt::ZeroInitializeWorkgroupMemory,
 }
 
 #[derive(Clone, Debug, Error)]

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

@@ -373,6 +373,7 @@ impl<A: hal::Api> Example<A> {
                     entry_point: "main",
                     constants: &Default::default(),
                 },
+                zero_initialise_workgroup_memory: Default::default(),
             })
         }
         .unwrap();

wgpu-hal/src/dx12/device.rs

@@ -214,6 +214,7 @@ impl super::Device {
         stage: &crate::ProgrammableStage<super::Api>,
         layout: &super::PipelineLayout,
         naga_stage: naga::ShaderStage,
+        zero_initialise_workgroup_memory: bool,
     ) -> Result<super::CompiledShader, crate::PipelineError> {
         use naga::back::hlsl;
 
@@ -226,9 +227,18 @@ impl super::Device {
         )
         .map_err(|e| crate::PipelineError::Linkage(stage_bit, format!("HLSL: {e:?}")))?;
 
+        let mut cloned_naga_options;
+        let naga_options = if !zero_initialise_workgroup_memory {
+            cloned_naga_options = layout.naga_options.clone();
+            cloned_naga_options.zero_initialize_workgroup_memory = zero_initialise_workgroup_memory;
+            &cloned_naga_options
+        } else {
+            &layout.naga_options
+        };
+
         //TODO: reuse the writer
         let mut source = String::new();
-        let mut writer = hlsl::Writer::new(&mut source, &layout.naga_options);
+        let mut writer = hlsl::Writer::new(&mut source, naga_options);
         let reflection_info = {
             profiling::scope!("naga::back::hlsl::write");
             writer
@@ -239,7 +249,7 @@ impl super::Device {
         let full_stage = format!(
             "{}_{}\0",
             naga_stage.to_hlsl_str(),
-            layout.naga_options.shader_model.to_str()
+            naga_options.shader_model.to_str()
         );
 
         let ep_index = module
@@ -1289,12 +1299,16 @@ impl crate::Device for super::Device {
         let (topology_class, topology) = conv::map_topology(desc.primitive.topology);
         let mut shader_stages = wgt::ShaderStages::VERTEX;
 
-        let blob_vs =
-            self.load_shader(&desc.vertex_stage, desc.layout, naga::ShaderStage::Vertex)?;
+        let blob_vs = self.load_shader(
+            &desc.vertex_stage,
+            desc.layout,
+            naga::ShaderStage::Vertex,
+            true,
+        )?;
         let blob_fs = match desc.fragment_stage {
             Some(ref stage) => {
                 shader_stages |= wgt::ShaderStages::FRAGMENT;
-                Some(self.load_shader(stage, desc.layout, naga::ShaderStage::Fragment)?)
+                Some(self.load_shader(stage, desc.layout, naga::ShaderStage::Fragment, true)?)
             }
             None => None,
         };
@@ -1473,7 +1487,12 @@ impl crate::Device for super::Device {
         &self,
         desc: &crate::ComputePipelineDescriptor<super::Api>,
     ) -> Result<super::ComputePipeline, crate::PipelineError> {
-        let blob_cs = self.load_shader(&desc.stage, desc.layout, naga::ShaderStage::Compute)?;
+        let blob_cs = self.load_shader(
+            &desc.stage,
+            desc.layout,
+            naga::ShaderStage::Compute,
+            desc.zero_initialise_workgroup_memory == wgt::ZeroInitializeWorkgroupMemory::always(),
+        )?;
 
         let pair = {
             profiling::scope!("ID3D12Device::CreateComputePipelineState");