cubek-tile 0.3.0

CubeK: Tile
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

//! A scale level a launch may or may not have bound, as the tile it serves.

use cubecl::prelude::*;

use crate::*;

/// A scale level a launch may or may not have bound, as the tile it serves.
#[cube]
pub fn scale_tile<S: Numeric>(
    level: &ComptimeOption<TileArg<'static, u32, Const<1>>>,
    #[comptime] space: Partitioning,
) -> ComptimeOption<Tile<S>> {
    #[comptime]
    match level {
        ComptimeOption::Some(level) => {
            ComptimeOption::new_Some(level.tile_as::<S>(comptime!(space.clone())))
        }
        ComptimeOption::None => ComptimeOption::new_None(),
    }
}

/// A scale level windowed the way the factor it multiplies is, or nothing where none was bound.
#[cube]
pub trait MaybeTile: CubeType {
    /// The element the level is served at.
    type E: Numeric;

    /// This level at `region`, descending it as [`Tile::at`] descends a factor.
    fn at(&self, region: &Region) -> ComptimeOption<Tile<Self::E>>;

    /// This level staged per region of `level`, or as it lies where `level` is `None`.
    fn staged(
        &self,
        #[comptime] level: Option<Level>,
        #[comptime] storage: StageStorage,
    ) -> ComptimeOption<Tile<Self::E>>;

    /// This level filled from `src`, where both are there; nothing where either is not.
    fn copy_from(&mut self, src: &ComptimeOption<Tile<Self::E>>);
}

#[cube]
impl<E: Numeric> MaybeTile for ComptimeOption<Tile<E>> {
    type E = E;

    fn at(&self, region: &Region) -> ComptimeOption<Tile<E>> {
        #[comptime]
        match self {
            ComptimeOption::Some(level) => ComptimeOption::new_Some(level.at(region)),
            ComptimeOption::None => ComptimeOption::new_None(),
        }
    }

    fn staged(
        &self,
        #[comptime] level: Option<Level>,
        #[comptime] storage: StageStorage,
    ) -> ComptimeOption<Tile<E>> {
        match comptime!(level) {
            None => self.clone(),
            Some(level) =>
            {
                #[comptime]
                match self {
                    ComptimeOption::Some(scales) => ComptimeOption::new_Some(
                        scales.stage(comptime!(level.clone()), comptime!(storage.clone())),
                    ),
                    ComptimeOption::None => ComptimeOption::new_None(),
                }
            }
        }
    }

    fn copy_from(&mut self, src: &ComptimeOption<Tile<E>>) {
        #[comptime]
        if let (ComptimeOption::Some(stage), ComptimeOption::Some(src)) = (self, src) {
            stage.copy_from(src);
        }
    }
}