cubecl_ir/
tma.rs

1use crate::{Instruction, TypeHash};
2use alloc::{string::String, vec::Vec};
3use core::fmt::{Display, Write};
4
5use crate::OperationReflect;
6
7use super::Value;
8
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
11#[operation(opcode_name = TmaOpCode)]
12/// Operations available on a barrier
13pub enum TmaOps {
14    TmaStore {
15        #[args(allow_ptr, ptr_read)]
16        source: Value,
17        coordinates: Vec<Value>,
18    },
19    CommitGroup,
20    WaitGroup {
21        max_pending: u32,
22    },
23    WaitGroupRead {
24        max_pending: u32,
25    },
26}
27
28impl Display for TmaOps {
29    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
30        match self {
31            TmaOps::TmaStore {
32                source,
33                coordinates,
34            } => {
35                let rank = coordinates.len();
36                let coords = coordinates.iter().fold(String::new(), |mut s, coord| {
37                    let _ = write!(s, ", {coord}");
38                    s
39                });
40                write!(f, "tma_load::<{rank}>({source} {coords})")
41            }
42            TmaOps::CommitGroup => write!(f, "memcpy_async_bulk_commit_group()"),
43            TmaOps::WaitGroup { max_pending } => {
44                write!(f, "tma_wait_group::<{max_pending}>()")
45            }
46            TmaOps::WaitGroupRead { max_pending } => {
47                write!(f, "tma_wait_group_read::<{max_pending}>()")
48            }
49        }
50    }
51}
52
53impl From<TmaOps> for Instruction {
54    fn from(value: TmaOps) -> Self {
55        Instruction::no_out(value)
56    }
57}
cubecl_ir/tma.rs

cubecl_ir/
tma.rs
