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_storage` is the same but change "cube address space(shared memory)" to "storage address space(input args)". But the set of invocations that are collaborating is still only the invocations in the same cube.
48///
49/// * There is no guarantee about using barriers alone to make the writes to storage buffer in one cube become visible to invocations in a different cube.
50pub fn sync_storage() {}
51
52pub mod sync_storage {
53    use super::*;
54
55    pub fn expand(scope: &Scope) {
56        scope.register(&SyncOp::new(scope.ctx_mut(), SyncScope::Device));
57    }
58}
59
60/// `sync_async_proxy_shared` is a synchronization fence for the experimental SM 9.0+ copy
61/// functions, applying bidirectionally between the async proxy (i.e. TMA) and shared memory.
62/// Should be used after initializing the barriers, and before the copy operation.
63/// PTX: `fence.proxy.async.shared::cta`
64/// Experimental and subject to change.
65pub fn sync_async_proxy_shared() {
66    unexpanded!()
67}
68
69pub mod sync_async_proxy_shared {
70    use super::*;
71
72    pub fn expand(scope: &Scope) {
73        scope.register(&SyncAsyncProxyOp::new(scope.ctx_mut()))
74    }
75}
76
77/// Barrier, then load `reference` with the result marked workgroup-uniform —
78/// mirrors WGSL's `workgroupUniformLoad`. Lets a workgroup-shared value gate
79/// control flow that contains barriers. Non-WGSL backends lower it to
80/// [`sync_cube`] plus a plain load.
81///
82/// Use [`workgroup_uniform_load_atomic`] for `Atomic<E>`.
83#[allow(unused_variables)]
84pub fn workgroup_uniform_load<E: CubePrimitive>(reference: &E) -> E {
85    unexpanded!()
86}
87
88/// Module containing the expand function for [`workgroup_uniform_load()`].
89pub mod workgroup_uniform_load {
90    use cubecl_ir::{
91        dialect::plane::UniformLoadOp, pliron::builtin::op_interfaces::OneResultInterface,
92    };
93
94    use crate::frontend::HasValue;
95
96    use super::*;
97
98    /// Expand method of [`workgroup_uniform_load()`].
99    pub fn expand<E: CubePrimitive>(scope: &Scope, reference: &NativeExpand<E>) -> NativeExpand<E> {
100        let ptr = reference.value(scope);
101        let op = UniformLoadOp::new(scope.ctx_mut(), ptr);
102        scope.register(&op);
103        op.get_result(scope.ctx()).into()
104    }
105}
106
107/// Atomic counterpart of [`workgroup_uniform_load`]: barrier + atomic load,
108/// returning the underlying numeric (WGSL's atomic `workgroupUniformLoad`
109/// overload).
110#[allow(unused_variables)]
111pub fn workgroup_uniform_load_atomic<E: CubePrimitive<Scalar: Numeric>>(
112    reference: &Atomic<E>,
113) -> E {
114    unexpanded!()
115}
116
117/// Module containing the expand function for [`workgroup_uniform_load_atomic()`].
118pub mod workgroup_uniform_load_atomic {
119    use cubecl_ir::{
120        dialect::plane::AtomicUniformLoadOp, pliron::builtin::op_interfaces::OneResultInterface,
121    };
122
123    use crate::frontend::HasValue;
124
125    use super::*;
126
127    /// Expand method of [`workgroup_uniform_load_atomic()`].
128    pub fn expand<E: CubePrimitive<Scalar: Numeric>>(
129        scope: &Scope,
130        reference: &NativeExpand<Atomic<E>>,
131    ) -> NativeExpand<E> {
132        let ptr = reference.value(scope);
133        let op = AtomicUniformLoadOp::new(scope.ctx_mut(), ptr);
134        scope.register(&op);
135        op.get_result(scope.ctx()).into()
136    }
137}
cubecl_core/frontend/synchronization.rs

cubecl_core/frontend/
synchronization.rs
