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
//! [`ComputeScope`]: the hardware a level's tiles can go to.
use cubecl::prelude::*;
/// The hardware a level's tiles can go to, ordered inner to outer.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
pub enum ComputeScope {
Unit,
Plane,
/// `planes` consecutive planes, the first a multiple of `planes`, issuing one instruction
/// together: a warpgroup MMA's four.
PlaneGroup {
planes: usize,
},
Cube,
}
impl ComputeScope {
/// The four planes of a warpgroup, which issue each warpgroup MMA together
/// ([`Instruction::Wgmma`](crate::Instruction::Wgmma)).
pub const WARPGROUP: Self = Self::PlaneGroup { planes: 4 };
/// The planes one instance of this scope spans, where it spans whole planes.
pub fn planes(self) -> Option<usize> {
match self {
ComputeScope::Plane => Some(1),
ComputeScope::PlaneGroup { planes } => Some(planes),
ComputeScope::Unit | ComputeScope::Cube => None,
}
}
}
#[cube]
impl ComputeScope {
/// This instance's position within `compute_scope`: which plane or plane group of the cube,
/// or which unit of the plane.
pub fn position(#[comptime] compute_scope: ComputeScope) -> usize {
match comptime!(compute_scope) {
ComputeScope::Plane => UNIT_POS_Y as usize,
ComputeScope::PlaneGroup { planes } => UNIT_POS_Y as usize / planes,
ComputeScope::Unit => UNIT_POS_X as usize,
ComputeScope::Cube => {
panic!(
"ComputeScope::position: a cube has one position per grid dimension; say which"
)
}
}
}
/// This unit's position among the units of its instance of `compute_scope`: in the cube,
/// in its plane group, or in its plane. A plane is the launch's `x`
/// ([`Partitioning::cube_dim`](crate::Partitioning::cube_dim)), so a unit's position in it is
/// `UNIT_POS_X`, and a group's planes are consecutive `y`s.
pub fn unit(#[comptime] compute_scope: ComputeScope) -> usize {
match comptime!(compute_scope) {
ComputeScope::Cube => UNIT_POS as usize,
ComputeScope::PlaneGroup { planes } => {
(UNIT_POS_Y as usize % planes) * CUBE_DIM_X as usize + UNIT_POS_X as usize
}
ComputeScope::Plane => UNIT_POS_X as usize,
ComputeScope::Unit => comptime!(panic!(
"ComputeScope::unit: a unit is the only unit of its own scope"
)),
}
}
}