cubecl_tile/space/partition/
compute_scope.rs

1//! [`ComputeScope`]: the hardware a level's tiles can go to.
2
3use cubecl::prelude::*;
4
5/// The hardware a level's tiles can go to, ordered inner to outer.
6#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
7pub enum ComputeScope {
8    Unit,
9    Plane,
10    /// `planes` consecutive planes, the first a multiple of `planes`, issuing one instruction
11    /// together: a warpgroup MMA's four.
12    PlaneGroup {
13        planes: usize,
14    },
15    Cube,
16}
17
18impl ComputeScope {
19    /// The four planes of a warpgroup, which issue each warpgroup MMA together
20    /// ([`Instruction::Wgmma`](crate::Instruction::Wgmma)).
21    pub const WARPGROUP: Self = Self::PlaneGroup { planes: 4 };
22
23    /// The planes one instance of this scope spans, where it spans whole planes.
24    pub fn planes(self) -> Option<usize> {
25        match self {
26            ComputeScope::Plane => Some(1),
27            ComputeScope::PlaneGroup { planes } => Some(planes),
28            ComputeScope::Unit | ComputeScope::Cube => None,
29        }
30    }
31}
32
33#[cube]
34impl ComputeScope {
35    /// This instance's position within `compute_scope`: which plane or plane group of the cube,
36    /// or which unit of the plane.
37    pub fn position(#[comptime] compute_scope: ComputeScope) -> usize {
38        match comptime!(compute_scope) {
39            ComputeScope::Plane => UNIT_POS_Y as usize,
40            ComputeScope::PlaneGroup { planes } => UNIT_POS_Y as usize / planes,
41            ComputeScope::Unit => UNIT_POS_X as usize,
42            ComputeScope::Cube => {
43                panic!(
44                    "ComputeScope::position: a cube has one position per grid dimension; say which"
45                )
46            }
47        }
48    }
49
50    /// This unit's position among the units of its instance of `compute_scope`: in the cube,
51    /// in its plane group, or in its plane. A plane is the launch's `x`
52    /// ([`Partitioning::cube_dim`](crate::Partitioning::cube_dim)), so a unit's position in it is
53    /// `UNIT_POS_X`, and a group's planes are consecutive `y`s.
54    pub fn unit(#[comptime] compute_scope: ComputeScope) -> usize {
55        match comptime!(compute_scope) {
56            ComputeScope::Cube => UNIT_POS as usize,
57            ComputeScope::PlaneGroup { planes } => {
58                (UNIT_POS_Y as usize % planes) * CUBE_DIM_X as usize + UNIT_POS_X as usize
59            }
60            ComputeScope::Plane => UNIT_POS_X as usize,
61            ComputeScope::Unit => comptime!(panic!(
62                "ComputeScope::unit: a unit is the only unit of its own scope"
63            )),
64        }
65    }
66}
cubecl_tile/space/partition/compute_scope.rs

cubecl_tile/space/partition/
compute_scope.rs
