cubecl_core/frontend/
pending.rs

1//! Work a kernel issued that completes later.
2//!
3//! Some operations run asynchronously: a warpgroup MMA, a TMA store. The kernel goes on while
4//! they run, and must not touch what they read or write until they complete. Such an operation
5//! returns a [`Pending`], which holds what the work produces, or `()`, and gives it back from
6//! [`Pending::wait`] once the work is done with what the kernel must not touch:
7//!
8//! ```rust, ignore
9//! let stored = tma_store_2d(&tile, &mut output, row, col);
10//! // ... other work, while the store reads `tile` ...
11//! stored.wait();
12//! // `tile` may be written again.
13//! ```
14//!
15//! The hardware tracks such work in commit groups, and waits for a number of groups rather than
16//! for a given one. The compiler derives that number for each wait, from the groups committed
17//! between the operation and the wait on every path that leads to it: a wait in a pipelined loop
18//! lets the newer groups run on.
19//!
20//! What "done" means depends on the work. A warpgroup MMA is done once it wrote its accumulator.
21//! A TMA store is done once it read its shared memory source: its writes to global memory may
22//! still be in flight, and are visible once the kernel ends. [`Pending::wait_complete`] also
23//! waits for those writes, for a kernel that hands them to another cube.
24
25use crate::{self as cubecl, prelude::*};
26use cubecl_ir::{
27    dialect::pending::{CommitOp, ReadyOp, WaitOp},
28    types::pending::{AsyncGroup, WaitUntil},
29};
30use cubecl_macros::{CubeTypeMut, intrinsic};
31use pliron::value::Value;
32
33/// Work that completes later, and the `T` it produces then.
34#[derive(CubeType, CubeTypeMut)]
35#[must_use = "the work runs on until it is waited on"]
36pub struct Pending<T: PendingValue> {
37    #[allow(unused)]
38    pub(crate) value: T,
39    pub(crate) token: GroupToken,
40}
41
42/// A value [`Pending`] work produces.
43#[cube]
44pub trait PendingValue: CubeType + Sized {
45    /// Waits for `pending` to complete, and returns its value.
46    fn wait_for(pending: Pending<Self>) -> Self;
47}
48
49#[cube]
50impl<T: PendingValue> Pending<T> {
51    /// Waits for the work to complete, and returns what it produced.
52    pub fn wait(self) -> T {
53        T::wait_for(self)
54    }
55}
56
57#[cube]
58impl PendingValue for () {
59    fn wait_for(pending: Pending<Self>) -> Self {
60        pending.token.wait();
61    }
62}
63
64#[cube]
65impl Pending<()> {
66    /// Waits for the work to fully complete. A TMA store has then also performed its writes to
67    /// global memory, which a fence or an atomic after the wait publishes to other cubes.
68    /// Anything else completes as [`Pending::wait`] waits for it.
69    pub fn wait_complete(self) {
70        self.token.wait_until(comptime![WaitUntil::Complete]);
71    }
72}
73
74// A completion carries no value, so waiting on it leaves nothing behind: a loop may wait on it
75// and replace it in the same iteration.
76impl Copy for PendingExpand<()> {}
77// The same work, waited on through either: a handle, as the value's own expand clones are.
78impl<T: PendingValue<ExpandType: Clone>> Clone for PendingExpand<T> {
79    fn clone(&self) -> Self {
80        Self {
81            value: self.value.clone(),
82            token: self.token,
83        }
84    }
85}
86impl Copy for Pending<()> {}
87impl<T: PendingValue + Clone> Clone for Pending<T> {
88    fn clone(&self) -> Self {
89        Self {
90            value: self.value.clone(),
91            token: self.token,
92        }
93    }
94}
95
96impl<T: PendingValue> PendingExpand<T> {
97    /// Commits the operations of `group` issued since its last commit, which produce `value`.
98    pub(crate) fn commit(scope: &Scope, group: AsyncGroup, value: T::ExpandType) -> Self {
99        Self {
100            value,
101            token: GroupToken::commit(scope, group),
102        }
103    }
104
105    /// `value`, with nothing of `group` to wait for.
106    pub(crate) fn ready(scope: &Scope, group: AsyncGroup, value: T::ExpandType) -> Self {
107        let op = ReadyOp::new(scope.ctx_mut(), group);
108        Self {
109            value,
110            token: scope.register_with_result(&op).into(),
111        }
112    }
113}
114
115/// The token of a commit group, which a wait resolves to a count of groups.
116#[derive(Clone, Copy)]
117pub(crate) struct GroupToken;
118pub(crate) type GroupTokenExpand = NativeExpand<GroupToken>;
119
120impl GroupToken {
121    pub(crate) fn commit(scope: &Scope, group: AsyncGroup) -> NativeExpand<GroupToken> {
122        let op = CommitOp::new(scope.ctx_mut(), group);
123        scope.register_with_result(&op).into()
124    }
125}
126
127#[cube]
128impl GroupToken {
129    /// Waits until the group this token was committed into is done with what the kernel must
130    /// not touch.
131    pub(crate) fn wait(&self) {
132        self.wait_until(comptime![WaitUntil::Released]);
133    }
134
135    /// Waits until the group this token was committed into reaches `until`.
136    #[allow(unused_variables)]
137    pub(crate) fn wait_until(&self, #[comptime] until: WaitUntil) {
138        intrinsic!(|scope| {
139            let token = self.read_value(scope);
140            scope.register(&WaitOp::new(scope.ctx_mut(), token, until));
141        })
142    }
143}
144
145impl CubeType for GroupToken {
146    type ExpandType = NativeExpand<GroupToken>;
147}
148
149impl CubeDebug for GroupToken {}
150
151impl ReadValue for NativeExpand<GroupToken> {
152    fn read_value(&self, scope: &Scope) -> Value {
153        self.expand.read_value(scope)
154    }
155}
156
157impl NativeAssign for GroupToken {}
158
159impl AsMutExpand for NativeExpand<GroupToken> {
160    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
161        self
162    }
163}
cubecl_core/frontend/pending.rs

cubecl_core/frontend/
pending.rs
