mircuda 0.1.0

Native, explicit-stream Rust gateway to NVIDIA CUDA
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
135
136
137
138
139

use std::{
    io::{self, Write},
    time::{Duration, Instant},
};

use mircuda::{DenseMatmulPlan, DenseMatmulSpec, DeviceBuffer, Driver, Graph, Stream, bf16};

const M: usize = 1;
const N: usize = 4_096;
const K: usize = 4_096;
const NODES: u16 = 32;
const REPLAYS: u16 = 200;

struct Report {
    name: &'static str,
    enqueue: Duration,
    device_ms: f32,
}

struct Sequence<'a> {
    stream: &'a Stream,
    plan: &'a mut DenseMatmulPlan<bf16>,
    a: &'a DeviceBuffer<bf16>,
    b: &'a DeviceBuffer<bf16>,
    c: &'a mut DeviceBuffer<bf16>,
}

fn capture_sequence(sequence: &mut Sequence<'_>) -> mircuda::Result<()> {
    for _ in 0..NODES {
        sequence
            .plan
            .execute(sequence.stream, sequence.a, sequence.b, sequence.c, 1.0, 0.0)?;
    }
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let driver = Driver::initialize()?;
    let device = driver.devices()?.into_iter().next().ok_or(mircuda::Error::InvalidLaunch)?;
    let context = driver.create_context(device)?;
    let info = context.device_info()?;
    let stream = context.create_stream()?;
    let pool = context.default_memory_pool()?;
    let a = pool.allocate_zeroed::<bf16>(&stream, M * K)?;
    let b = pool.allocate_zeroed::<bf16>(&stream, N * K)?;
    let mut c = pool.allocate_zeroed::<bf16>(&stream, M * N)?;
    let mut plan = DenseMatmulPlan::new(&context, &stream, DenseMatmulSpec::new(M, N, K)?)?;
    let direct = measure_direct(&context, &stream, &mut plan, &a, &b, &mut c)?;
    let sequence = Sequence {
        stream: &stream,
        plan: &mut plan,
        a: &a,
        b: &b,
        c: &mut c,
    };
    let mut graph = stream.capture(sequence, capture_sequence)?;
    let replay = measure_graph(&context, &stream, &mut graph)?;
    let mut output = io::stdout().lock();
    writeln!(output, "device: {}", info.name)?;
    writeln!(output, "sequence: {NODES} BF16 GEMMs, m={M} n={N} k={K}")?;
    write_report(&mut output, &direct)?;
    write_report(&mut output, &replay)?;
    Ok(())
}

fn measure_direct(
    context: &mircuda::Context,
    stream: &Stream,
    plan: &mut DenseMatmulPlan<bf16>,
    a: &DeviceBuffer<bf16>,
    b: &DeviceBuffer<bf16>,
    c: &mut DeviceBuffer<bf16>,
) -> mircuda::Result<Report> {
    for _ in 0..NODES {
        plan.execute(stream, a, b, c, 1.0, 0.0)?;
    }
    stream.synchronize()?;
    let started = context.create_event(true)?;
    let completed = context.create_event(true)?;
    started.record(stream)?;
    let host_started = Instant::now();
    for _ in 0..REPLAYS {
        for _ in 0..NODES {
            plan.execute(stream, a, b, c, 1.0, 0.0)?;
        }
    }
    let enqueue = host_started.elapsed();
    completed.record(stream)?;
    completed.synchronize()?;
    Ok(Report {
        name: "direct",
        enqueue,
        device_ms: started.elapsed_ms(&completed)?,
    })
}

fn measure_graph<Resources>(
    context: &mircuda::Context,
    stream: &Stream,
    graph: &mut Graph<Resources>,
) -> mircuda::Result<Report> {
    for _ in 0..10 {
        graph.launch(stream)?;
    }
    stream.synchronize()?;
    let started = context.create_event(true)?;
    let completed = context.create_event(true)?;
    started.record(stream)?;
    let host_started = Instant::now();
    for _ in 0..REPLAYS {
        graph.launch(stream)?;
    }
    let enqueue = host_started.elapsed();
    completed.record(stream)?;
    completed.synchronize()?;
    Ok(Report {
        name: "graph",
        enqueue,
        device_ms: started.elapsed_ms(&completed)?,
    })
}

fn write_report(output: &mut impl Write, report: &Report) -> io::Result<()> {
    let replays = f64::from(REPLAYS);
    let nodes = f64::from(NODES);
    writeln!(output, "{}:", report.name)?;
    writeln!(
        output,
        "  enqueue: {:.3} us/replay, {:.3} ns/GEMM",
        report.enqueue.as_secs_f64() * 1_000_000.0 / replays,
        report.enqueue.as_secs_f64() * 1_000_000_000.0 / replays / nodes
    )?;
    writeln!(
        output,
        "  device: {:.3} us/replay, {:.3} us/GEMM",
        f64::from(report.device_ms) * 1_000.0 / replays,
        f64::from(report.device_ms) * 1_000.0 / replays / nodes
    )
}