add support for zero-initializing workgroup memory

benches/criterion.rs

@@ -242,6 +242,7 @@ fn backends(c: &mut Criterion) {
                 version: naga::back::glsl::Version::new_gles(320),
                 writer_flags: naga::back::glsl::WriterFlags::empty(),
                 binding_map: Default::default(),
+                zero_initialize_workgroup_memory: true,
             };
             for &(ref module, ref info) in inputs.iter() {
                 for ep in module.entry_points.iter() {

src/back/glsl/mod.rs

@@ -228,6 +228,8 @@ pub struct Options {
     pub writer_flags: WriterFlags,
     /// Map of resources association to binding locations.
     pub binding_map: BindingMap,
+    /// Should workgroup variables be zero initialized (by polyfilling)?
+    pub zero_initialize_workgroup_memory: bool,
 }
 
 impl Default for Options {
@@ -236,6 +238,7 @@ impl Default for Options {
             version: Version::new_gles(310),
             writer_flags: WriterFlags::ADJUST_COORDINATE_SPACE,
             binding_map: BindingMap::default(),
+            zero_initialize_workgroup_memory: true,
         }
     }
 }
@@ -1399,6 +1402,8 @@ impl<'a, W: Write> Writer<'a, W> {
         // Close the parentheses and open braces to start the function body
         writeln!(self.out, ") {{")?;
 
+        self.write_workgroup_variables_initialization(&ctx)?;
+
         // Compose the function arguments from globals, in case of an entry point.
         if let back::FunctionType::EntryPoint(ep_index) = ctx.ty {
             let stage = self.module.entry_points[ep_index as usize].stage;
@@ -1487,6 +1492,41 @@ impl<'a, W: Write> Writer<'a, W> {
         Ok(())
     }
 
+    fn write_workgroup_variables_initialization(
+        &mut self,
+        ctx: &back::FunctionCtx,
+    ) -> BackendResult {
+        if !self.options.zero_initialize_workgroup_memory
+            || !ctx.ty.is_compute_entry_point(self.module)
+        {
+            return Ok(());
+        }
+
+        let vars = self
+            .module
+            .global_variables
+            .iter()
+            .filter(|&(handle, var)| {
+                !ctx.info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup
+            });
+
+        let mut write_barrier = false;
+
+        for (handle, var) in vars {
+            let name = &self.names[&NameKey::GlobalVariable(handle)];
+            write!(self.out, "{}{} = ", back::INDENT, name)?;
+            self.write_zero_init_value(var.ty)?;
+            writeln!(self.out, ";")?;
+            write_barrier = true;
+        }
+
+        if write_barrier {
+            self.write_barrier(crate::Barrier::WORK_GROUP, back::Level(1))?;
+        }
+
+        Ok(())
+    }
+
     /// Helper method that writes a list of comma separated `T` with a writer function `F`
     ///
     /// The writer function `F` receives a mutable reference to `self` that if needed won't cause
@@ -1996,18 +2036,8 @@ impl<'a, W: Write> Writer<'a, W> {
             // keyword which ceases all further processing in a fragment shader, it's called OpKill
             // in spir-v that's why it's called `Statement::Kill`
             Statement::Kill => writeln!(self.out, "{}discard;", level)?,
-            // Issue a memory barrier. Please note that to ensure visibility,
-            // OpenGL always requires a call to the `barrier()` function after a `memoryBarrier*()`
             Statement::Barrier(flags) => {
-                if flags.contains(crate::Barrier::STORAGE) {
-                    writeln!(self.out, "{}memoryBarrierBuffer();", level)?;
-                }
-
-                if flags.contains(crate::Barrier::WORK_GROUP) {
-                    writeln!(self.out, "{}memoryBarrierShared();", level)?;
-                }
-
-                writeln!(self.out, "{}barrier();", level)?;
+                self.write_barrier(flags, level)?;
             }
             // Stores in glsl are just variable assignments written as `pointer = value;`
             Statement::Store { pointer, value } => {
@@ -3499,7 +3529,7 @@ impl<'a, W: Write> Writer<'a, W> {
     fn write_zero_init_value(&mut self, ty: Handle<crate::Type>) -> BackendResult {
         let inner = &self.module.types[ty].inner;
         match *inner {
-            TypeInner::Scalar { kind, .. } => {
+            TypeInner::Scalar { kind, .. } | TypeInner::Atomic { kind, .. } => {
                 self.write_zero_init_scalar(kind)?;
             }
             TypeInner::Vector { kind, .. } => {
@@ -3544,7 +3574,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
                 write!(self.out, ")")?;
             }
-            _ => {} // TODO:
+            _ => unreachable!(),
         }
 
         Ok(())
@@ -3562,6 +3592,19 @@ impl<'a, W: Write> Writer<'a, W> {
         Ok(())
     }
 
+    /// Issue a memory barrier. Please note that to ensure visibility,
+    /// OpenGL always requires a call to the `barrier()` function after a `memoryBarrier*()`
+    fn write_barrier(&mut self, flags: crate::Barrier, level: back::Level) -> BackendResult {
+        if flags.contains(crate::Barrier::STORAGE) {
+            writeln!(self.out, "{}memoryBarrierBuffer();", level)?;
+        }
+        if flags.contains(crate::Barrier::WORK_GROUP) {
+            writeln!(self.out, "{}memoryBarrierShared();", level)?;
+        }
+        writeln!(self.out, "{}barrier();", level)?;
+        Ok(())
+    }
+
     /// Helper function that return the glsl storage access string of [`StorageAccess`](crate::StorageAccess)
     ///
     /// glsl allows adding both `readonly` and `writeonly` but this means that

src/back/hlsl/mod.rs

@@ -191,6 +191,8 @@ pub struct Options {
     pub special_constants_binding: Option<BindTarget>,
     /// Bind target of the push constant buffer
     pub push_constants_target: Option<BindTarget>,
+    /// Should workgroup variables be zero initialized (by polyfilling)?
+    pub zero_initialize_workgroup_memory: bool,
 }
 
 impl Default for Options {
@@ -201,6 +203,7 @@ impl Default for Options {
             fake_missing_bindings: true,
             special_constants_binding: None,
             push_constants_target: None,
+            zero_initialize_workgroup_memory: true,
         }
     }
 }

src/back/hlsl/writer.rs

@@ -1154,6 +1154,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         writeln!(self.out)?;
         writeln!(self.out, "{{")?;
 
+        self.write_workgroup_variables_initialization(func_ctx, module)?;
+
         if let back::FunctionType::EntryPoint(index) = func_ctx.ty {
             self.write_ep_arguments_initialization(module, func, index)?;
         }
@@ -1207,6 +1209,38 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         Ok(())
     }
 
+    fn write_workgroup_variables_initialization(
+        &mut self,
+        func_ctx: &back::FunctionCtx,
+        module: &Module,
+    ) -> BackendResult {
+        if !self.options.zero_initialize_workgroup_memory
+            || !func_ctx.ty.is_compute_entry_point(module)
+        {
+            return Ok(());
+        }
+
+        let vars = module.global_variables.iter().filter(|&(handle, var)| {
+            !func_ctx.info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup
+        });
+
+        let mut write_barrier = false;
+
+        for (handle, var) in vars {
+            let name = &self.names[&NameKey::GlobalVariable(handle)];
+            write!(self.out, "{}{} = ", back::INDENT, name)?;
+            self.write_default_init(module, var.ty)?;
+            writeln!(self.out, ";")?;
+            write_barrier = true;
+        }
+
+        if write_barrier {
+            self.write_barrier(crate::Barrier::WORK_GROUP, back::Level(1))?;
+        }
+
+        Ok(())
+    }
+
     /// Helper method used to write statements
     ///
     /// # Notes
@@ -1693,13 +1727,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             Statement::Break => writeln!(self.out, "{}break;", level)?,
             Statement::Continue => writeln!(self.out, "{}continue;", level)?,
             Statement::Barrier(barrier) => {
-                if barrier.contains(crate::Barrier::STORAGE) {
-                    writeln!(self.out, "{}DeviceMemoryBarrierWithGroupSync();", level)?;
-                }
-
-                if barrier.contains(crate::Barrier::WORK_GROUP) {
-                    writeln!(self.out, "{}GroupMemoryBarrierWithGroupSync();", level)?;
-                }
+                self.write_barrier(barrier, level)?;
             }
             Statement::ImageStore {
                 image,
@@ -2817,27 +2845,21 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
     /// Helper function that write default zero initialization
     fn write_default_init(&mut self, module: &Module, ty: Handle<crate::Type>) -> BackendResult {
-        match module.types[ty].inner {
-            TypeInner::Array {
-                size: crate::ArraySize::Constant(const_handle),
-                base,
-                ..
-            } => {
-                write!(self.out, "{{")?;
-                let count = module.constants[const_handle].to_array_length().unwrap();
-                for i in 0..count {
-                    if i != 0 {
-                        write!(self.out, ",")?;
-                    }
-                    self.write_default_init(module, base)?;
-                }
-                write!(self.out, "}}")?;
-            }
-            _ => {
-                write!(self.out, "(")?;
-                self.write_type(module, ty)?;
-                write!(self.out, ")0")?;
-            }
+        write!(self.out, "(")?;
+        self.write_type(module, ty)?;
+        if let TypeInner::Array { base, size, .. } = module.types[ty].inner {
+            self.write_array_size(module, base, size)?;
+        }
+        write!(self.out, ")0")?;
+        Ok(())
+    }
+
+    fn write_barrier(&mut self, barrier: crate::Barrier, level: back::Level) -> BackendResult {
+        if barrier.contains(crate::Barrier::STORAGE) {
+            writeln!(self.out, "{}DeviceMemoryBarrierWithGroupSync();", level)?;
+        }
+        if barrier.contains(crate::Barrier::WORK_GROUP) {
+            writeln!(self.out, "{}GroupMemoryBarrierWithGroupSync();", level)?;
         }
         Ok(())
     }

src/back/mod.rs

@@ -47,6 +47,17 @@ enum FunctionType {
     EntryPoint(crate::proc::EntryPointIndex),
 }
 
+impl FunctionType {
+    fn is_compute_entry_point(&self, module: &crate::Module) -> bool {
+        match *self {
+            FunctionType::EntryPoint(index) => {
+                module.entry_points[index as usize].stage == crate::ShaderStage::Compute
+            }
+            _ => false,
+        }
+    }
+}
+
 /// Helper structure that stores data needed when writing the function
 struct FunctionCtx<'a> {
     /// The current function being written

src/back/msl/mod.rs

@@ -207,6 +207,8 @@ pub struct Options {
     /// Bounds checking policies.
     #[cfg_attr(feature = "deserialize", serde(default))]
     pub bounds_check_policies: index::BoundsCheckPolicies,
+    /// Should workgroup variables be zero initialized (by polyfilling)?
+    pub zero_initialize_workgroup_memory: bool,
 }
 
 impl Default for Options {
@@ -218,6 +220,7 @@ impl Default for Options {
             spirv_cross_compatibility: false,
             fake_missing_bindings: true,
             bounds_check_policies: index::BoundsCheckPolicies::default(),
+            zero_initialize_workgroup_memory: true,
         }
     }
 }