cubecl_core/frontend/
synchronization.rs

1use crate::{
2    frontend::{NativeExpand, element::Atomic},
3    ir::{
4        Scope,
5        dialect::synchronization::{SyncAsyncProxyOp, SyncOp, SyncScope},
6    },
7    prelude::{CubePrimitive, Numeric},
8    unexpanded,
9};
10
11// Among all backends, the memory order guarantee of WebGPU is the weakest
12// So Cubecl's memory order cannot be stronger than that of WebGPU
13
14/// # Coordinates the following among all invocations in the current cube:
15///
16/// * Memory writes to variables in cube address space(shared memory) complete,
17///   e.g. writes that were initiated actually land in the cube address space memory.
18///
19/// * Then all the invocations in the cube wait for each other to arrive at the barrier, i.e. this step.
20///
21/// * Then all the invocations int the cube begin executing after the barrier, and all writes to cube address space made before the barrier are now visible to any invocation in this cube.
22pub fn sync_cube() {}
23
24pub mod sync_cube {
25    use super::*;
26
27    pub fn expand(scope: &Scope) {
28        scope.register(&SyncOp::new(scope.ctx_mut(), SyncScope::Cube));
29    }
30}
31
32/// Synchronizes units within their plane (e.g., warp or SIMD group).
33///
34/// Warning: not all targets support plane-level synchronization.
35pub fn sync_plane() {
36    unexpanded!()
37}
38
39pub mod sync_plane {
40    use super::*;
41
42    pub fn expand(scope: &Scope) {
43        scope.register(&SyncOp::new(scope.ctx_mut(), SyncScope::Plane));
44    }
45}
46
47/// [`sync_cube`] over storage memory (input args) rather than shared memory, and the one
48/// synchronization that reaches past the cube.
49///
50/// Every unit of the cube must reach it, and it is a [`sync_cube`] as well as what follows, so a
51/// value one unit leaves in shared memory before it is readable by the rest after it. On top of
52/// that cube barrier it is a release and an acquire at device scope: every write this cube made to storage before it is
53/// visible to any other cube that calls it afterwards, and every write another cube published
54/// before its own call is visible here after it. That is what lets cubes hand each other partial
55/// results — publish, `sync_storage`, then announce through an atomic — with no second dispatch.
56///
57/// **Only where the runtime says so.** WebGPU's memory model promises nothing across workgroups,
58/// so on a backend whose
59/// [`device_memory_scope`](cubecl_ir::Features::device_memory_scope) is `false` this is a cube
60/// barrier and no more, and a cube reading another's writes may read them stale. Ask before
61/// launching a kernel that depends on the promise.
62pub fn sync_storage() {}
63
64pub mod sync_storage {
65    use super::*;
66
67    pub fn expand(scope: &Scope) {
68        scope.register(&SyncOp::new(scope.ctx_mut(), SyncScope::Device));
69    }
70}
71
72/// `sync_async_proxy_shared` is a synchronization fence for the experimental SM 9.0+ copy
73/// functions, applying bidirectionally between the async proxy (i.e. TMA) and shared memory.
74/// Should be used after initializing the barriers, and before the copy operation.
75/// PTX: `fence.proxy.async.shared::cta`
76/// Experimental and subject to change.
77pub fn sync_async_proxy_shared() {
78    unexpanded!()
79}
80
81pub mod sync_async_proxy_shared {
82    use super::*;
83
84    pub fn expand(scope: &Scope) {
85        scope.register(&SyncAsyncProxyOp::new(scope.ctx_mut()))
86    }
87}
88
89/// Barrier, then load `reference` with the result marked workgroup-uniform —
90/// mirrors WGSL's `workgroupUniformLoad`. Lets a workgroup-shared value gate
91/// control flow that contains barriers. Non-WGSL backends lower it to
92/// [`sync_cube`] plus a plain load.
93///
94/// Use [`workgroup_uniform_load_atomic`] for `Atomic<E>`.
95#[allow(unused_variables)]
96pub fn workgroup_uniform_load<E: CubePrimitive>(reference: &E) -> E {
97    unexpanded!()
98}
99
100/// Module containing the expand function for [`workgroup_uniform_load()`].
101pub mod workgroup_uniform_load {
102    use cubecl_ir::{
103        dialect::plane::UniformLoadOp, pliron::builtin::op_interfaces::OneResultInterface,
104    };
105
106    use crate::frontend::HasValue;
107
108    use super::*;
109
110    /// Expand method of [`workgroup_uniform_load()`].
111    pub fn expand<E: CubePrimitive>(scope: &Scope, reference: &NativeExpand<E>) -> NativeExpand<E> {
112        let ptr = reference.value(scope);
113        let op = UniformLoadOp::new(scope.ctx_mut(), ptr);
114        scope.register(&op);
115        op.get_result(scope.ctx()).into()
116    }
117}
118
119/// Atomic counterpart of [`workgroup_uniform_load`]: barrier + atomic load,
120/// returning the underlying numeric (WGSL's atomic `workgroupUniformLoad`
121/// overload).
122#[allow(unused_variables)]
123pub fn workgroup_uniform_load_atomic<E: CubePrimitive<Scalar: Numeric>>(
124    reference: &Atomic<E>,
125) -> E {
126    unexpanded!()
127}
128
129/// Module containing the expand function for [`workgroup_uniform_load_atomic()`].
130pub mod workgroup_uniform_load_atomic {
131    use cubecl_ir::{
132        dialect::plane::AtomicUniformLoadOp, pliron::builtin::op_interfaces::OneResultInterface,
133    };
134
135    use crate::frontend::HasValue;
136
137    use super::*;
138
139    /// Expand method of [`workgroup_uniform_load_atomic()`].
140    pub fn expand<E: CubePrimitive<Scalar: Numeric>>(
141        scope: &Scope,
142        reference: &NativeExpand<Atomic<E>>,
143    ) -> NativeExpand<E> {
144        let ptr = reference.value(scope);
145        let op = AtomicUniformLoadOp::new(scope.ctx_mut(), ptr);
146        scope.register(&op);
147        op.get_result(scope.ctx()).into()
148    }
149}
cubecl_core/frontend/synchronization.rs

cubecl_core/frontend/
synchronization.rs
