[naga spv-in] Add support for memory barriers

naga/src/back/dot/mod.rs

@@ -112,7 +112,8 @@ impl StatementGraph {
                     }
                     "Continue"
                 }
-                S::Barrier(_flags) => "Barrier",
+                S::ControlBarrier(_flags) => "ControlBarrier",
+                S::MemoryBarrier(_flags) => "MemoryBarrier",
                 S::Block(ref b) => {
                     let (other, last) = self.add(b, targets);
                     self.flow.push((id, other, ""));

naga/src/back/glsl/mod.rs

@@ -1968,7 +1968,7 @@ impl<'a, W: Write> Writer<'a, W> {
             }
 
             writeln!(self.out, "{level}}}")?;
-            self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
+            self.write_control_barrier(crate::Barrier::WORK_GROUP, level)?;
         }
 
         Ok(())
@@ -2512,8 +2512,11 @@ 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, "{level}discard;")?,
-            Statement::Barrier(flags) => {
-                self.write_barrier(flags, level)?;
+            Statement::ControlBarrier(flags) => {
+                self.write_control_barrier(flags, level)?;
+            }
+            Statement::MemoryBarrier(flags) => {
+                self.write_memory_barrier(flags, level)?;
             }
             // Stores in glsl are just variable assignments written as `pointer = value;`
             Statement::Store { pointer, value } => {
@@ -2527,14 +2530,14 @@ impl<'a, W: Write> Writer<'a, W> {
                 // GLSL doesn't have pointers, which means that this backend needs to ensure that
                 // the actual "loading" is happening between the two barriers.
                 // This is done in `Emit` by never emitting a variable name for pointer variables
-                self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
+                self.write_control_barrier(crate::Barrier::WORK_GROUP, level)?;
 
                 let result_name = Baked(result).to_string();
                 write!(self.out, "{level}")?;
                 // Expressions cannot have side effects, so just writing the expression here is fine.
                 self.write_named_expr(pointer, result_name, result, ctx)?;
 
-                self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
+                self.write_control_barrier(crate::Barrier::WORK_GROUP, level)?;
             }
             // Stores a value into an image.
             Statement::ImageStore {
@@ -4912,9 +4915,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 {
+    /// Issue a control barrier.
+    fn write_control_barrier(
+        &mut self,
+        flags: crate::Barrier,
+        level: back::Level,
+    ) -> BackendResult {
+        self.write_memory_barrier(flags, level)?;
+        writeln!(self.out, "{level}barrier();")?;
+        Ok(())
+    }
+
+    /// Issue a memory barrier.
+    fn write_memory_barrier(&mut self, flags: crate::Barrier, level: back::Level) -> BackendResult {
         if flags.contains(crate::Barrier::STORAGE) {
             writeln!(self.out, "{level}memoryBarrierBuffer();")?;
         }
@@ -4927,7 +4940,6 @@ impl<'a, W: Write> Writer<'a, W> {
         if flags.contains(crate::Barrier::TEXTURE) {
             writeln!(self.out, "{level}memoryBarrierImage();")?;
         }
-        writeln!(self.out, "{level}barrier();")?;
         Ok(())
     }
 

naga/src/back/hlsl/writer.rs

@@ -1647,7 +1647,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         }
 
         writeln!(self.out, "{level}}}")?;
-        self.write_barrier(crate::Barrier::WORK_GROUP, level)
+        self.write_control_barrier(crate::Barrier::WORK_GROUP, level)
     }
 
     /// Helper method used to write switches
@@ -2291,8 +2291,11 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     writeln!(self.out, "{level}continue;")?
                 }
             }
-            Statement::Barrier(barrier) => {
-                self.write_barrier(barrier, level)?;
+            Statement::ControlBarrier(barrier) => {
+                self.write_control_barrier(barrier, level)?;
+            }
+            Statement::MemoryBarrier(barrier) => {
+                self.write_memory_barrier(barrier, level)?;
             }
             Statement::ImageStore {
                 image,
@@ -2464,12 +2467,12 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 writeln!(self.out, ");")?;
             }
             Statement::WorkGroupUniformLoad { pointer, result } => {
-                self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
+                self.write_control_barrier(crate::Barrier::WORK_GROUP, level)?;
                 write!(self.out, "{level}")?;
                 let name = Baked(result).to_string();
                 self.write_named_expr(module, pointer, name, result, func_ctx)?;
 
-                self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
+                self.write_control_barrier(crate::Barrier::WORK_GROUP, level)?;
             }
             Statement::Switch {
                 selector,
@@ -4287,7 +4290,11 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         Ok(())
     }
 
-    fn write_barrier(&mut self, barrier: crate::Barrier, level: back::Level) -> BackendResult {
+    fn write_control_barrier(
+        &mut self,
+        barrier: crate::Barrier,
+        level: back::Level,
+    ) -> BackendResult {
         if barrier.contains(crate::Barrier::STORAGE) {
             writeln!(self.out, "{level}DeviceMemoryBarrierWithGroupSync();")?;
         }
@@ -4303,6 +4310,26 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         Ok(())
     }
 
+    fn write_memory_barrier(
+        &mut self,
+        barrier: crate::Barrier,
+        level: back::Level,
+    ) -> BackendResult {
+        if barrier.contains(crate::Barrier::STORAGE) {
+            writeln!(self.out, "{level}DeviceMemoryBarrier();")?;
+        }
+        if barrier.contains(crate::Barrier::WORK_GROUP) {
+            writeln!(self.out, "{level}GroupMemoryBarrier();")?;
+        }
+        if barrier.contains(crate::Barrier::SUB_GROUP) {
+            // Does not exist in DirectX
+        }
+        if barrier.contains(crate::Barrier::TEXTURE) {
+            writeln!(self.out, "{level}DeviceMemoryBarrier();")?;
+        }
+        Ok(())
+    }
+
     /// Helper to emit the shared tail of an HLSL atomic call (arguments, value, result)
     fn emit_hlsl_atomic_tail(
         &mut self,

naga/src/back/msl/writer.rs

@@ -3700,7 +3700,8 @@ impl<W: Write> Writer<W> {
                 crate::Statement::Kill => {
                     writeln!(self.out, "{level}{NAMESPACE}::discard_fragment();")?;
                 }
-                crate::Statement::Barrier(flags) => {
+                crate::Statement::ControlBarrier(flags)
+                | crate::Statement::MemoryBarrier(flags) => {
                     self.write_barrier(flags, level)?;
                 }
                 crate::Statement::Store { pointer, value } => {

naga/src/back/pipeline_constants.rs

@@ -803,7 +803,11 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 crate::RayQueryFunction::Terminate => {}
             }
         }
-        Statement::Break | Statement::Continue | Statement::Kill | Statement::Barrier(_) => {}
+        Statement::Break
+        | Statement::Continue
+        | Statement::Kill
+        | Statement::ControlBarrier(_)
+        | Statement::MemoryBarrier(_) => {}
     }
 }
 

naga/src/back/spv/block.rs

@@ -3240,8 +3240,11 @@ impl BlockContext<'_> {
                     self.function.consume(block, Instruction::kill());
                     return Ok(BlockExitDisposition::Discarded);
                 }
-                Statement::Barrier(flags) => {
-                    self.writer.write_barrier(flags, &mut block);
+                Statement::ControlBarrier(flags) => {
+                    self.writer.write_control_barrier(flags, &mut block);
+                }
+                Statement::MemoryBarrier(flags) => {
+                    self.writer.write_memory_barrier(flags, &mut block);
                 }
                 Statement::Store { pointer, value } => {
                     let value_id = self.cached[value];
@@ -3576,7 +3579,7 @@ impl BlockContext<'_> {
                 }
                 Statement::WorkGroupUniformLoad { pointer, result } => {
                     self.writer
-                        .write_barrier(crate::Barrier::WORK_GROUP, &mut block);
+                        .write_control_barrier(crate::Barrier::WORK_GROUP, &mut block);
                     let result_type_id = self.get_expression_type_id(&self.fun_info[result].ty);
                     // Embed the body of
                     match self.write_access_chain(
@@ -3616,7 +3619,7 @@ impl BlockContext<'_> {
                         }
                     }
                     self.writer
-                        .write_barrier(crate::Barrier::WORK_GROUP, &mut block);
+                        .write_control_barrier(crate::Barrier::WORK_GROUP, &mut block);
                 }
                 Statement::RayQuery { query, ref fun } => {
                     self.write_ray_query_function(query, fun, &mut block);

naga/src/back/spv/instructions.rs

@@ -1138,6 +1138,12 @@ impl super::Instruction {
         instruction.add_operand(semantics_id);
         instruction
     }
+    pub(super) fn memory_barrier(mem_scope_id: Word, semantics_id: Word) -> Self {
+        let mut instruction = Self::new(Op::MemoryBarrier);
+        instruction.add_operand(mem_scope_id);
+        instruction.add_operand(semantics_id);
+        instruction
+    }
 
     // Group Instructions
 

naga/src/back/spv/writer.rs

@@ -1711,9 +1711,11 @@ impl Writer {
         Ok(id)
     }
 
-    pub(super) fn write_barrier(&mut self, flags: crate::Barrier, block: &mut Block) {
+    pub(super) fn write_control_barrier(&mut self, flags: crate::Barrier, block: &mut Block) {
         let memory_scope = if flags.contains(crate::Barrier::STORAGE) {
             spirv::Scope::Device
+        } else if flags.contains(crate::Barrier::SUB_GROUP) {
+            spirv::Scope::Subgroup
         } else {
             spirv::Scope::Workgroup
         };
@@ -1726,6 +1728,10 @@ impl Writer {
             spirv::MemorySemantics::WORKGROUP_MEMORY,
             flags.contains(crate::Barrier::WORK_GROUP),
         );
+        semantics.set(
+            spirv::MemorySemantics::SUBGROUP_MEMORY,
+            flags.contains(crate::Barrier::SUB_GROUP),
+        );
         semantics.set(
             spirv::MemorySemantics::IMAGE_MEMORY,
             flags.contains(crate::Barrier::TEXTURE),
@@ -1744,6 +1750,37 @@ impl Writer {
         ));
     }
 
+    pub(super) fn write_memory_barrier(&mut self, flags: crate::Barrier, block: &mut Block) {
+        let mut semantics = spirv::MemorySemantics::ACQUIRE_RELEASE;
+        semantics.set(
+            spirv::MemorySemantics::UNIFORM_MEMORY,
+            flags.contains(crate::Barrier::STORAGE),
+        );
+        semantics.set(
+            spirv::MemorySemantics::WORKGROUP_MEMORY,
+            flags.contains(crate::Barrier::WORK_GROUP),
+        );
+        semantics.set(
+            spirv::MemorySemantics::SUBGROUP_MEMORY,
+            flags.contains(crate::Barrier::SUB_GROUP),
+        );
+        semantics.set(
+            spirv::MemorySemantics::IMAGE_MEMORY,
+            flags.contains(crate::Barrier::TEXTURE),
+        );
+        let mem_scope_id = if flags.contains(crate::Barrier::STORAGE) {
+            self.get_index_constant(spirv::Scope::Device as u32)
+        } else if flags.contains(crate::Barrier::SUB_GROUP) {
+            self.get_index_constant(spirv::Scope::Subgroup as u32)
+        } else {
+            self.get_index_constant(spirv::Scope::Workgroup as u32)
+        };
+        let semantics_id = self.get_index_constant(semantics.bits());
+        block
+            .body
+            .push(Instruction::memory_barrier(mem_scope_id, semantics_id));
+    }
+
     fn generate_workgroup_vars_init_block(
         &mut self,
         entry_id: Word,
@@ -1844,7 +1881,7 @@ impl Writer {
 
         let mut post_if_block = Block::new(merge_id);
 
-        self.write_barrier(crate::Barrier::WORK_GROUP, &mut post_if_block);
+        self.write_control_barrier(crate::Barrier::WORK_GROUP, &mut post_if_block);
 
         let next_id = self.id_gen.next();
         function.consume(post_if_block, Instruction::branch(next_id));

naga/src/back/wgsl/writer.rs

@@ -830,7 +830,7 @@ impl<W: Write> Writer<W> {
             Statement::Continue => {
                 writeln!(self.out, "{level}continue;")?;
             }
-            Statement::Barrier(barrier) => {
+            Statement::ControlBarrier(barrier) | Statement::MemoryBarrier(barrier) => {
                 if barrier.contains(crate::Barrier::STORAGE) {
                     writeln!(self.out, "{level}storageBarrier();")?;
                 }

naga/src/compact/statements.rs

@@ -155,7 +155,8 @@ impl FunctionTracer<'_> {
                     St::Break
                     | St::Continue
                     | St::Kill
-                    | St::Barrier(_)
+                    | St::ControlBarrier(_)
+                    | St::MemoryBarrier(_)
                     | St::Return { value: None } => {}
                 }
             }
@@ -364,7 +365,8 @@ impl FunctionMap {
                     St::Break
                     | St::Continue
                     | St::Kill
-                    | St::Barrier(_)
+                    | St::ControlBarrier(_)
+                    | St::MemoryBarrier(_)
                     | St::Return { value: None } => {}
                 }
             }

naga/src/front/glsl/builtins.rs

@@ -2035,8 +2035,10 @@ impl MacroCall {
             )?,
             MacroCall::Barrier => {
                 ctx.emit_restart();
-                ctx.body
-                    .push(crate::Statement::Barrier(crate::Barrier::all()), meta);
+                ctx.body.push(
+                    crate::Statement::ControlBarrier(crate::Barrier::all()),
+                    meta,
+                );
                 return Ok(None);
             }
             MacroCall::SmoothStep { splatted } => {