cubecl_ir/dialect/
synchronization.rs

1use cubecl_macros_internal::cube_op;
2use derive_more::From;
3use derive_new::new;
4use pliron::{
5    derive::{format, op_interface_impl, pliron_attr},
6    opts::dce::SideEffects,
7};
8
9use crate::{CanMaterialize, NoMemoryEffect, interfaces::Synchronizes, prelude::*};
10
11/// Scope that the synchronization should apply to. This is a *minimum*, when fine-grained control
12/// is not available it should synchronize at the smallest scope that includes this scope
13/// (i.e. `SyncScope::Plane` may be implemented by a `workgroupBarrier()`)
14#[format]
15#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)]
16pub enum SyncScope {
17    Unit,
18    Plane,
19    Cube,
20    Device,
21}
22
23#[pliron_attr(name = "cube.sync_scope", format = "$0", verifier = "succ")]
24#[derive(new, From, PartialEq, Eq, Clone, Debug, Hash, PartialOrd, Ord)]
25pub struct SyncScopeAttr(pub SyncScope);
26
27#[cube_op(name = "sync.sync")]
28#[result_ty(none)]
29#[op_traits(CanMaterialize, NoMemoryEffect)]
30pub struct SyncOp {
31    pub scope: SyncScopeAttr,
32}
33
34#[op_interface_impl]
35impl Synchronizes for SyncOp {
36    fn minimum_scope(&self, ctx: &Context) -> SyncScope {
37        self.scope(ctx).0
38    }
39    fn maximum_scope(&self, ctx: &Context) -> SyncScope {
40        self.scope(ctx).0
41    }
42}
43
44#[op_interface_impl]
45impl SideEffects for SyncOp {
46    fn has_side_effects(&self, _ctx: &Context) -> bool {
47        true
48    }
49}
50
51/// Fences the async proxy in CUDA, to make shared memory available to it. Does not implement
52/// `Synchronizes`, because it works only as a memory availability barrier with an outside chip.
53/// It does not synchronize the actual threads, and is typically called only by the TMA leader.
54#[cube_op(name = "sync.sync_async_proxy")]
55#[result_ty(none)]
56#[op_traits(CanMaterialize, NoMemoryEffect)]
57pub struct SyncAsyncProxyOp {}
cubecl_ir/dialect/synchronization.rs

cubecl_ir/dialect/
synchronization.rs
