Impl SyncStorage (#72)

tracel-ai · Aug 19, 2024 · b99e753 · b99e753
1 parent ccde038
commit b99e753
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 1 deletion.
diff --git a/crates/cubecl-core/src/frontend/synchronization.rs b/crates/cubecl-core/src/frontend/synchronization.rs
@@ -10,3 +10,13 @@ pub mod sync_units {
         context.register(Synchronization::SyncUnits)
     }
 }
+
+pub fn sync_storage() {}
+
+pub mod sync_storage {
+    use super::*;
+
+    pub fn __expand(context: &mut CubeContext) {
+        context.register(Synchronization::SyncStorage)
+    }
+}
diff --git a/crates/cubecl-core/src/ir/synchronization.rs b/crates/cubecl-core/src/ir/synchronization.rs
@@ -6,4 +6,5 @@ use serde::{Deserialize, Serialize};
 pub enum Synchronization {
     // Synchronizize units in a cube.
     SyncUnits,
+    SyncStorage
 }
diff --git a/crates/cubecl-cuda/src/compiler/base.rs b/crates/cubecl-cuda/src/compiler/base.rs
@@ -135,6 +135,7 @@ impl CudaCompiler {
             gpu::Operation::Branch(val) => self.compile_branch(instructions, val),
             gpu::Operation::Synchronization(val) => match val {
                 gpu::Synchronization::SyncUnits => instructions.push(Instruction::SyncThreads),
+                gpu::Synchronization::SyncStorage => instructions.push(Instruction::ThreadFence),
             },
             gpu::Operation::Subcube(op) => {
                 self.wrap_size_checked = true;

diff --git a/crates/cubecl-cuda/src/compiler/instruction.rs b/crates/cubecl-cuda/src/compiler/instruction.rs
@@ -114,6 +114,7 @@ pub enum Instruction {
         out: Variable,
     },
     SyncThreads,
+    ThreadFence,
     Ceil(UnaryInstruction),
     Floor(UnaryInstruction),
     Wrap(WarpInstruction),
@@ -264,6 +265,7 @@ for (uint {i} = {start}; {i} < {end}; {increment}) {{
                 out,
             } => Clamp::format(f, input, min_value, max_value, out),
             Instruction::SyncThreads => f.write_str("__syncthreads();\n"),
+            Instruction::ThreadFence => f.write_str("__threadfence();\n"),
             Instruction::Ceil(it) => Ceil::format(f, &it.input, &it.out),
             Instruction::Floor(it) => Floor::format(f, &it.input, &it.out),
             Instruction::SliceLength { input, out } => {

diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs
@@ -398,6 +398,9 @@ impl WgslCompiler {
         match synchronization {
             cube::Synchronization::SyncUnits => {
                 instructions.push(wgsl::Instruction::WorkgroupBarrier)
+            },
+            cube::Synchronization::SyncStorage=> {
+                instructions.push(wgsl::Instruction::StorageBarrier)
             }
         };
     }

diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs
@@ -44,6 +44,7 @@ pub enum Instruction {
     Return,
     Break,
     WorkgroupBarrier,
+    StorageBarrier,
     // Index handles casting to correct local variable.
     Index {
         lhs: Variable,
@@ -578,6 +579,7 @@ for (var {i}: u32 = {start}; {i} < {end}; {increment}) {{
             Instruction::Return => f.write_str("return;\n"),
             Instruction::Break => f.write_str("break;\n"),
             Instruction::WorkgroupBarrier => f.write_str("workgroupBarrier();\n"),
+            Instruction::StorageBarrier => f.write_str("storageBarrier();\n"),
             Instruction::Length { var, out } => match var {
                 Variable::Slice { .. } => f.write_fmt(format_args!("{out} = {var}_length;\n")),
                 _ => f.write_fmt(format_args!("{out} = arrayLength(&{var});\n")),
@@ -615,7 +617,7 @@ for (var {i}: u32 = {start}; {i} < {end}; {increment}) {{
                 f.write_fmt(format_args!("{out} = atomicLoad({input});\n"))
             }
             Instruction::AtomicStore { input, out } => {
-                f.write_fmt(format_args!("{out} = atomicStore({input});\n"))
+                f.write_fmt(format_args!("atomicStore({out},{input});\n"))
             }
             Instruction::AtomicSwap { lhs, rhs, out } => {
                 f.write_fmt(format_args!("{out} = atomicExchange({lhs}, {rhs});"))