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///
15/// Every scope up to [`Cube`](SyncScope::Cube) synchronizes the units that share it and orders the
16/// memory they wrote for each other. [`Device`](SyncScope::Device) contains
17/// [`Cube`](SyncScope::Cube) — the units of the cube meet at it and their shared memory is
18/// ordered — and reaches past it: it is also a release and an acquire at device scope, so a write
19/// one cube published before it is visible to any other cube that synchronizes at this scope
20/// afterwards. Only a runtime whose
21/// [`device_memory_scope`](crate::Features::device_memory_scope) is set promises that second
22/// half; the others give the cube barrier alone.
23#[format]
24#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)]
25pub enum SyncScope {
26    Unit,
27    Plane,
28    Cube,
29    Device,
30}
31
32#[pliron_attr(name = "cube.sync_scope", format = "$0", verifier = "succ")]
33#[derive(new, From, PartialEq, Eq, Clone, Debug, Hash, PartialOrd, Ord)]
34pub struct SyncScopeAttr(pub SyncScope);
35
36#[cube_op(name = "sync.sync")]
37#[result_ty(none)]
38#[op_traits(CanMaterialize, NoMemoryEffect)]
39pub struct SyncOp {
40    pub scope: SyncScopeAttr,
41}
42
43#[op_interface_impl]
44impl Synchronizes for SyncOp {
45    fn minimum_scope(&self, ctx: &Context) -> SyncScope {
46        self.scope(ctx).0
47    }
48    fn maximum_scope(&self, ctx: &Context) -> SyncScope {
49        self.scope(ctx).0
50    }
51}
52
53#[op_interface_impl]
54impl SideEffects for SyncOp {
55    fn has_side_effects(&self, _ctx: &Context) -> bool {
56        true
57    }
58}
59
60/// Fences the async proxy in CUDA, to make shared memory available to it. Does not implement
61/// `Synchronizes`, because it works only as a memory availability barrier with an outside chip.
62/// It does not synchronize the actual threads, and is typically called only by the TMA leader.
63#[cube_op(name = "sync.sync_async_proxy")]
64#[result_ty(none)]
65#[op_traits(CanMaterialize, NoMemoryEffect)]
66pub struct SyncAsyncProxyOp {}
cubecl_ir/dialect/synchronization.rs

cubecl_ir/dialect/
synchronization.rs
