cubecl_ir/dialect/
synchronization.rs

1use cubecl_macros_internal::cube_op;
2use derive_more::From;
3use derive_new::new;
4use pliron::derive::{format, op_interface_impl, pliron_attr};
5
6use crate::{CanMaterialize, HasSideEffects, interfaces::Synchronizes, prelude::*};
7
8/// Scope that the synchronization should apply to. This is a *minimum*, when fine-grained control
9/// is not available it should synchronize at the smallest scope that includes this scope
10/// (i.e. `SyncScope::Plane` may be implemented by a `workgroupBarrier()`)
11///
12/// Every scope up to [`Cube`](SyncScope::Cube) synchronizes the units that share it and orders the
13/// memory they wrote for each other. [`Device`](SyncScope::Device) contains
14/// [`Cube`](SyncScope::Cube) — the units of the cube meet at it and their shared memory is
15/// ordered — and reaches past it: it is also a release and an acquire at device scope, so a write
16/// one cube published before it is visible to any other cube that synchronizes at this scope
17/// afterwards. Only a runtime whose
18/// [`device_memory_scope`](crate::Features::device_memory_scope) is set promises that second
19/// half; the others give the cube barrier alone.
20#[format]
21#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)]
22pub enum SyncScope {
23    Unit,
24    Plane,
25    Cube,
26    Device,
27}
28
29#[pliron_attr(name = "cube.sync_scope", format = "$0", verifier = "succ")]
30#[derive(new, From, PartialEq, Eq, Clone, Debug, Hash, PartialOrd, Ord)]
31pub struct SyncScopeAttr(pub SyncScope);
32
33#[cube_op(name = "sync.sync")]
34#[result_ty(none)]
35#[op_traits(CanMaterialize, HasSideEffects)]
36pub struct SyncOp {
37    pub scope: SyncScopeAttr,
38}
39
40#[op_interface_impl]
41impl Synchronizes for SyncOp {
42    fn minimum_scope(&self, ctx: &Context) -> SyncScope {
43        self.scope(ctx).0
44    }
45    fn maximum_scope(&self, ctx: &Context) -> SyncScope {
46        self.scope(ctx).0
47    }
48}
49
50/// Fences the async proxy in CUDA, to make shared memory available to it. Does not implement
51/// `Synchronizes`, because it works only as a memory availability barrier with an outside chip.
52/// It does not synchronize the actual threads, and is typically called only by the TMA leader.
53#[cube_op(name = "sync.sync_async_proxy")]
54#[result_ty(none)]
55#[op_traits(CanMaterialize, HasSideEffects)]
56pub struct SyncAsyncProxyOp {}
cubecl_ir/dialect/synchronization.rs

cubecl_ir/dialect/
synchronization.rs
