cubek-convolution 0.2.0

CubeK: Convolution Kernels
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

//! Shared helpers for the basic tier.

use cubecl::prelude::CubePrimitive;
use cubek_matmul::{
    components::stage::PartitionBuffering,
    definition::{MatmulElems, MatmulGlobalElems, TilingScheme},
};
use cubek_std::{PartitionSize, StageSize, SwizzleModes, TileSize};

use crate::suite::launcher_strategy::ConvolutionSize;

/// `MatmulElems` with f16 globals (and the f32 accumulator picked by
/// `from_globals`).
pub(crate) fn f16_dtypes() -> MatmulElems {
    let f16 = half::f16::as_type_native_unchecked().storage_type();
    MatmulElems::from_globals(&MatmulGlobalElems {
        lhs: f16,
        rhs: f16,
        out: f16,
    })
}

/// Tile size that works on both macOS (where CMMA is 8x8x8) and elsewhere.
#[cfg(target_os = "macos")]
pub(crate) fn default_tile_size() -> TileSize {
    TileSize { m: 8, n: 8, k: 8 }
}

#[cfg(not(target_os = "macos"))]
pub(crate) fn default_tile_size() -> TileSize {
    TileSize {
        m: 16,
        n: 16,
        k: 16,
    }
}

/// Smallest viable partition for a smoke test.
pub(crate) fn small_partition() -> PartitionSize {
    PartitionSize { m: 1, n: 1, k: 1 }
}

/// Two-plane stage along m,n. Default for most basic-tier tests.
pub(crate) fn small_stage() -> StageSize {
    StageSize { m: 2, n: 2, k: 1 }
}

pub(crate) fn default_tiling_scheme() -> TilingScheme {
    TilingScheme::builder()
        .with_tile_size(default_tile_size())
        .with_partition_size(small_partition())
        .with_stage_size(small_stage())
        .build()
        .unwrap()
}

pub(crate) fn default_swizzle() -> SwizzleModes {
    SwizzleModes::default()
}

pub(crate) fn default_partition_buffering() -> PartitionBuffering {
    PartitionBuffering::Single
}

/// Smoke-sized 2D problem (16x16 spatial, 16 in_c, 32 out_c).
pub(crate) fn small_size() -> ConvolutionSize {
    ConvolutionSize {
        h: 16,
        w: 16,
        c: 16,
        out_c: 32,
    }
}

/// Wider mid-sized problem, used for `basic`-feature-gated tests.
pub(crate) fn medium_size() -> ConvolutionSize {
    ConvolutionSize {
        h: 32,
        w: 32,
        c: 32,
        out_c: 16,
    }
}