cubek-convolution 0.3.0

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

//! Forced-blueprint tests for non-default knobs (swizzle, partition
//! buffering). Each knob is exercised once on a representative routine;
//! per-routine combinations live in `full/`.

use cubek_convolution::ConvAlgorithm;
use cubek_matmul::multi_level::{
    PartitionSize, StageSize, components::stage::PartitionBuffering, stage::SwizzleMode,
    stage::SwizzleModes,
};

use super::common::{default_size, default_tile_size, f16_dtypes, tiling_scheme};
use crate::convolution::launcher_strategy::test_algo;

fn small_partition() -> PartitionSize {
    PartitionSize { m: 1, n: 1, k: 1 }
}

fn small_stage() -> StageSize {
    StageSize { m: 2, n: 2, k: 1 }
}

// -- Swizzle -----------------------------------------------------------------

fn run_swizzle(swizzle: SwizzleModes) {
    test_algo(
        ConvAlgorithm::SimpleSyncCyclic,
        f16_dtypes(),
        tiling_scheme(default_tile_size(), small_partition(), small_stage()),
        swizzle,
        PartitionBuffering::Single,
        default_size(),
    );
}

#[test]
fn swizzle_b32() {
    run_swizzle(SwizzleModes {
        lhs: SwizzleMode::B32,
        rhs: SwizzleMode::B32,
        ..Default::default()
    });
}

#[test]
fn swizzle_b64() {
    run_swizzle(SwizzleModes {
        lhs: SwizzleMode::B64,
        rhs: SwizzleMode::B64,
        ..Default::default()
    });
}

#[test]
fn swizzle_b128() {
    run_swizzle(SwizzleModes {
        lhs: SwizzleMode::B128,
        rhs: SwizzleMode::B128,
        ..Default::default()
    });
}

// -- Partition buffering -----------------------------------------------------
//
// Partition buffering pipelines inside a partition along n, so it needs at
// least two tiles along n (partition.n >= 2).

#[test]
fn partition_buffering_double() {
    test_algo(
        ConvAlgorithm::SimpleSyncTilewise,
        f16_dtypes(),
        tiling_scheme(
            default_tile_size(),
            PartitionSize { m: 1, n: 2, k: 1 },
            small_stage(),
        ),
        SwizzleModes::default(),
        PartitionBuffering::Double,
        default_size(),
    );
}