fastnn 0.3.3

A robust deep learning library in Rust with CUDA 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
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Hot path throughput on the GPU if there is one, otherwise the CPU.
//! GPU timings include the launch.

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use fastnn::prelude::*;

fn device() -> (Device, &'static str) {
    match Device::cuda(0) {
        Ok(device) => (device, "cuda"),
        Err(_) => (Device::Cpu, "cpu"),
    }
}

fn matmul(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("matmul/{name}"));

    for size in [128usize, 512, 1024] {
        let a = Tensor::randn(&[size, size]).to(device);
        let b = Tensor::randn(&[size, size]).to(device);

        // Two flops per multiply-add.
        group.throughput(Throughput::Elements(2 * (size * size * size) as u64));
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |bench, _| {
            bench.iter(|| a.matmul(&b))
        });
    }
    group.finish();
}

/// Should track plain matmul, if not a transposed copy is being materialised.
fn matmul_fused(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("matmul_fused/{name}"));

    let a = Tensor::randn(&[512, 512]).to(device);
    let b = Tensor::randn(&[512, 512]).to(device);

    group.bench_function("plain", |bench| bench.iter(|| a.matmul(&b)));
    group.bench_function("nt", |bench| bench.iter(|| a.matmul_nt(&b)));
    group.bench_function("tn", |bench| bench.iter(|| a.matmul_tn(&b)));
    group.bench_function("transpose_then_matmul", |bench| {
        bench.iter(|| a.matmul(&b.transpose()))
    });
    group.finish();
}

fn elementwise(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("elementwise/{name}"));

    let a = Tensor::randn(&[1024, 1024]).to(device);
    let b = Tensor::randn(&[1024, 1024]).to(device);
    group.throughput(Throughput::Elements(1024 * 1024));

    group.bench_function("add", |bench| bench.iter(|| a.add(&b)));
    group.bench_function("mul", |bench| bench.iter(|| a.mul(&b)));
    group.bench_function("relu", |bench| bench.iter(|| a.relu()));
    group.bench_function("gelu", |bench| bench.iter(|| a.gelu()));
    group.finish();
}

fn reductions(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("reduce/{name}"));

    let x = Tensor::randn(&[1024, 1024]).to(device);
    group.throughput(Throughput::Elements(1024 * 1024));

    group.bench_function("sum", |bench| bench.iter(|| x.sum()));
    group.bench_function("sum_axis_0", |bench| bench.iter(|| x.sum_axis(0)));
    group.bench_function("sum_axis_1", |bench| bench.iter(|| x.sum_axis(1)));
    group.bench_function("softmax", |bench| bench.iter(|| x.softmax()));
    group.finish();
}

fn layout(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("layout/{name}"));

    let heads = Tensor::randn(&[32, 8, 128, 64]).to(device);
    group.bench_function("permute_heads", |bench| {
        bench.iter(|| heads.permute(&[0, 2, 1, 3]))
    });

    let row = Tensor::randn(&[1, 512]).to(device);
    group.bench_function("expand", |bench| bench.iter(|| row.expand(&[1024, 512])));
    group.finish();
}

fn training_step(c: &mut Criterion) {
    let (device, name) = device();
    let mut group = c.benchmark_group(format!("step/{name}"));

    let model = Sequential::new()
        .add(Linear::new(784, 256))
        .add(ReLU)
        .add(Linear::new(256, 10));
    model.to_device(device);

    let mut opt = Adam::new(model.parameters(), 1e-3);
    let inputs = Tensor::randn(&[128, 784]).to(device);
    let targets: Vec<usize> = (0..128).map(|i| i % 10).collect();

    group.bench_function("mlp_batch128", |bench| {
        bench.iter(|| {
            let loss = cross_entropy(&model.forward(&inputs), &targets);
            opt.zero_grad();
            loss.backward();
            opt.step();
        })
    });

    let block = TransformerBlock::causal(256, 8, 1024, 0.0);
    block.to_device(device);
    block.eval();
    let sequence = Tensor::randn(&[8, 128, 256]).to(device);

    group.bench_function("transformer_block_forward", |bench| {
        bench.iter(|| no_grad(|| block.forward(&sequence)))
    });
    group.finish();
}

criterion_group!(
    benches,
    matmul,
    matmul_fused,
    elementwise,
    reductions,
    layout,
    training_step
);
criterion_main!(benches);