oxicuda 0.1.4

OxiCUDA - Pure Rust CUDA replacement for the COOLJAPAN ecosystem (95% performance target)
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
101
102
103
104
105
106
107
108
109
110
111
112

use std::hint::black_box;

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use oxicuda_blas::level3::gemm::dispatch::{GemmDispatcher, GemmProblem};
use oxicuda_blas::{AlgorithmSelector, MathMode, Transpose};
use oxicuda_ptx::arch::SmVersion;
use oxicuda_ptx::ir::PtxType;

fn make_problem(m: u32, n: u32, k: u32) -> GemmProblem {
    GemmProblem {
        m,
        n,
        k,
        trans_a: Transpose::NoTrans,
        trans_b: Transpose::NoTrans,
        input_type: PtxType::F32,
        output_type: PtxType::F32,
        math_mode: MathMode::Default,
    }
}

fn bench_gemm_classify(c: &mut Criterion) {
    let mut group = c.benchmark_group("blas_gemm_classify");

    let dispatcher = GemmDispatcher::new(SmVersion::Sm80);

    let sizes: &[(&str, u32, u32, u32)] = &[
        ("standard_512", 512, 512, 512),
        ("skinny_8x512", 8, 512, 256),
        ("split_k_64x64x8192", 64, 64, 8192),
        ("large_4096", 4096, 4096, 4096),
    ];

    for &(name, m, n, k) in sizes {
        let problem = make_problem(m, n, k);
        group.bench_with_input(
            BenchmarkId::new("classify", name),
            &problem,
            |b, problem| {
                b.iter(|| black_box(dispatcher.classify(problem)));
            },
        );
    }
    group.finish();
}

fn bench_gemm_heuristic_tile(c: &mut Criterion) {
    let mut group = c.benchmark_group("blas_gemm_heuristic_tile");

    let dispatcher = GemmDispatcher::new(SmVersion::Sm80);

    let sizes: &[(&str, u32, u32, u32)] = &[
        ("standard_512", 512, 512, 512),
        ("skinny_8x512", 8, 512, 256),
        ("split_k_64x64x8192", 64, 64, 8192),
        ("large_4096", 4096, 4096, 4096),
    ];

    for &(name, m, n, k) in sizes {
        let problem = make_problem(m, n, k);
        let category = dispatcher.classify(&problem);
        group.bench_with_input(
            BenchmarkId::new("heuristic_tile", name),
            &(problem, category),
            |b, (problem, category)| {
                b.iter(|| black_box(dispatcher.heuristic_tile_config(problem, category)));
            },
        );
    }
    group.finish();
}

fn bench_algorithm_selector(c: &mut Criterion) {
    let mut group = c.benchmark_group("blas_algorithm_selector");

    let selector = AlgorithmSelector::new(SmVersion::Sm80);

    let sizes: &[(&str, u32, u32, u32)] = &[
        ("small_128", 128, 128, 128),
        ("medium_512", 512, 512, 512),
        ("large_1024", 1024, 1024, 1024),
    ];

    for &(name, m, n, k) in sizes {
        let problem = make_problem(m, n, k);
        group.bench_with_input(
            BenchmarkId::new("enumerate_algorithms", name),
            &problem,
            |b, problem| {
                b.iter(|| black_box(selector.enumerate_algorithms(problem)));
            },
        );
    }

    group.bench_function("algorithm_count", |b| {
        b.iter(|| black_box(selector.algorithm_count()));
    });

    group.bench_function("sm_version", |b| {
        b.iter(|| black_box(selector.sm_version()));
    });

    group.finish();
}

criterion_group!(
    blas_benches,
    bench_gemm_classify,
    bench_gemm_heuristic_tile,
    bench_algorithm_selector
);
criterion_main!(blas_benches);