cubecl-tile 0.12.0-pre.1

CubeCL Tile: a space cut into tiles, and the operations on them.
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

//! One axis's size ([`Extent`]) and every axis's at once ([`Shape`]).

use crate::{Axis, AxisMap};

/// One axis's size: `Static` is comptime, `Dynamic` is resolved in-kernel from the tensor shape.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub(crate) enum Extent {
    Static(usize),
    Dynamic,
}

impl Extent {
    /// The comptime size; panics on `Dynamic`.
    pub(crate) fn get(self) -> usize {
        match self {
            Extent::Static(n) => n,
            Extent::Dynamic => {
                panic!("Extent::get: this axis is Dynamic; its size is only known at runtime")
            }
        }
    }

    pub(crate) fn is_dynamic(self) -> bool {
        matches!(self, Extent::Dynamic)
    }

    /// The extent two spaces sharing an axis broadcast to when they are
    /// [`merge`](crate::Space::merge)d.
    pub(crate) fn broadcast(self, other: Extent) -> Extent {
        match (self, other) {
            (Extent::Static(1), b) => b,
            (a, Extent::Static(1)) => a,
            (Extent::Dynamic, _) | (_, Extent::Dynamic) => Extent::Dynamic,
            (Extent::Static(a), Extent::Static(b)) if a == b => Extent::Static(a),
            _ => panic!("Space::merge: axis appears with conflicting extents"),
        }
    }
}

/// Every axis with its comptime extent, in canonical order.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Shape {
    extents: AxisMap<Extent>,
}

impl Shape {
    pub(crate) fn new(extents: &[(Axis, Extent)]) -> Self {
        Shape {
            extents: AxisMap::new(extents),
        }
    }

    pub fn rank(&self) -> usize {
        self.extents.len()
    }

    pub fn axis_at(&self, i: usize) -> Axis {
        self.extents.axis_at(i)
    }

    pub fn position(&self, axis: Axis) -> usize {
        self.extents.position(axis)
    }

    pub fn contains(&self, axis: Axis) -> bool {
        self.extents.contains(axis)
    }

    pub fn axes(&self) -> impl Iterator<Item = Axis> + '_ {
        self.extents.axes()
    }

    pub(crate) fn extent_raw(&self, axis: Axis) -> Extent {
        self.extents.get(axis)
    }

    /// The axis's comptime size; panics on a dynamic axis.
    pub fn extent(&self, axis: Axis) -> usize {
        self.extent_raw(axis).get()
    }

    pub fn extent_at(&self, i: usize) -> usize {
        self.extent(self.axis_at(i))
    }

    pub fn is_dynamic(&self, axis: Axis) -> bool {
        self.extent_raw(axis).is_dynamic()
    }

    /// Every axis is [`Static`](Extent::Static), so a walk over it is fully comptime.
    pub(crate) fn is_static(&self) -> bool {
        self.axes().all(|axis| !self.is_dynamic(axis))
    }

    /// The listed axes with their extents, in the order listed.
    pub(crate) fn subspace(&self, axes: &[Axis]) -> Shape {
        let entries: Vec<_> = axes.iter().map(|&a| (a, self.extent_raw(a))).collect();
        Shape::new(&entries)
    }
}