cubecl-ir 0.12.0-pre.2

Intermediate representation for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

use cubecl_macros_internal::cube_op;
use derive_more::From;
use derive_new::new;
use pliron::derive::{format, op_interface_impl, pliron_attr};

use crate::{CanMaterialize, HasSideEffects, interfaces::Synchronizes, prelude::*};

/// Scope that the synchronization should apply to. This is a *minimum*, when fine-grained control
/// is not available it should synchronize at the smallest scope that includes this scope
/// (i.e. `SyncScope::Plane` may be implemented by a `workgroupBarrier()`)
///
/// Every scope up to [`Cube`](SyncScope::Cube) synchronizes the units that share it and orders the
/// memory they wrote for each other. [`Device`](SyncScope::Device) contains
/// [`Cube`](SyncScope::Cube) — the units of the cube meet at it and their shared memory is
/// ordered — and reaches past it: it is also a release and an acquire at device scope, so a write
/// one cube published before it is visible to any other cube that synchronizes at this scope
/// afterwards. Only a runtime whose
/// [`device_memory_scope`](crate::Features::device_memory_scope) is set promises that second
/// half; the others give the cube barrier alone.
///
/// [`System`](SyncScope::System) reaches past the device: what a cube wrote before it is visible to
/// the host and to kernels on other devices that bind this memory, and what they published before
/// it is visible here after it. That is what kernels on different devices hand each other data
/// through. A runtime that shares no memory between devices gives it the meaning of
/// [`Device`](SyncScope::Device), since nothing outside the device can see its memory.
#[format]
#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)]
pub enum SyncScope {
    Unit,
    Plane,
    Cube,
    Device,
    System,
}

#[pliron_attr(name = "cube.sync_scope", format = "$0", verifier = "succ")]
#[derive(new, From, PartialEq, Eq, Clone, Debug, Hash, PartialOrd, Ord)]
pub struct SyncScopeAttr(pub SyncScope);

#[cube_op(name = "sync.sync")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct SyncOp {
    pub scope: SyncScopeAttr,
}

#[op_interface_impl]
impl Synchronizes for SyncOp {
    fn minimum_scope(&self, ctx: &Context) -> SyncScope {
        self.scope(ctx).0
    }
    fn maximum_scope(&self, ctx: &Context) -> SyncScope {
        self.scope(ctx).0
    }
}

/// A plane barrier that also orders storage memory: the units of the plane meet, as at a
/// [`SyncOp`] at [`Plane`](SyncScope::Plane) scope, and every write one of them made to storage
/// before it is visible to the others after it. A plane barrier alone orders shared memory only.
/// It reaches no further than the plane: other planes and other cubes are not ordered by it.
#[cube_op(name = "sync.sync_plane_storage")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct SyncPlaneStorageOp {}

#[op_interface_impl]
impl Synchronizes for SyncPlaneStorageOp {
    fn minimum_scope(&self, _ctx: &Context) -> SyncScope {
        SyncScope::Plane
    }
    fn maximum_scope(&self, _ctx: &Context) -> SyncScope {
        SyncScope::Plane
    }
}

/// Fences the async proxy in CUDA, to make shared memory available to it. Does not implement
/// `Synchronizes`, because it works only as a memory availability barrier with an outside chip.
/// It does not synchronize the actual threads, and is typically called only by the TMA leader.
#[cube_op(name = "sync.sync_async_proxy")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct SyncAsyncProxyOp {}