atlas_spark_bench/
gpu.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GPU utilities for kernel microbenchmarks.
4//!
5//! Initializes `AtlasRegistry`, allocates device buffers, and provides
6//! CUDA event-based timing for Criterion `iter_custom` benchmarks.
7
8use std::ffi::c_void;
9use std::sync::OnceLock;
10
11use anyhow::{Result, bail};
12use atlas_core::registry::{AtlasRegistry, RawCudaFunc};
13use cudarc::driver::LaunchConfig;
14
15// Raw CUDA driver API for benchmarks.
16unsafe extern "C" {
17    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
18    fn cuMemFree_v2(dptr: u64) -> i32;
19    fn cuMemsetD8Async(dst: u64, value: u8, n: usize, stream: u64) -> i32;
20    fn cuStreamSynchronize(stream: u64) -> i32;
21    fn cuEventCreate(event: *mut u64, flags: u32) -> i32;
22    fn cuEventRecord(event: u64, stream: u64) -> i32;
23    fn cuEventSynchronize(event: u64) -> i32;
24    fn cuEventElapsedTime(ms: *mut f32, start: u64, end: u64) -> i32;
25    fn cuEventDestroy_v2(event: u64) -> i32;
26}
27
28/// The bench process's registry. A benchmark run loads one model's kernels and
29/// keeps them for the process — leaking it deliberately is what buys the
30/// `&'static` the bench harnesses are written against. Production loads a
31/// registry per model and drops it; see `atlas_core::registry::release`.
32static INIT: OnceLock<&'static AtlasRegistry> = OnceLock::new();
33
34/// Ensure the registry is loaded (idempotent).
35pub fn ensure_registry() -> &'static AtlasRegistry {
36    INIT.get_or_init(|| {
37        let ptx = atlas_kernels::ptx_modules();
38        let registry =
39            AtlasRegistry::load(0, &ptx).expect("AtlasRegistry load failed — is GPU available?");
40        Box::leak(Box::new(registry)) as &'static AtlasRegistry
41    })
42}
43
44/// Allocate `bytes` of GPU memory, zero-initialized.
45pub fn gpu_alloc_zeroed(stream: u64, bytes: usize) -> Result<u64> {
46    let mut dptr: u64 = 0;
47    let status = unsafe { cuMemAlloc_v2(&mut dptr, bytes) };
48    if status != 0 {
49        bail!("cuMemAlloc_v2 failed: status {status}, {bytes} bytes");
50    }
51    let status = unsafe { cuMemsetD8Async(dptr, 0, bytes, stream) };
52    if status != 0 {
53        bail!("cuMemsetD8Async failed: status {status}");
54    }
55    Ok(dptr)
56}
57
58/// Free GPU memory.
59pub fn gpu_free(dptr: u64) {
60    if dptr != 0 {
61        unsafe { cuMemFree_v2(dptr) };
62    }
63}
64
65/// Synchronize the stream.
66pub fn gpu_sync(stream: u64) -> Result<()> {
67    let status = unsafe { cuStreamSynchronize(stream) };
68    if status != 0 {
69        bail!("cuStreamSynchronize failed: status {status}");
70    }
71    Ok(())
72}
73
74/// Resolve `module::func` for a bench.
75///
76/// Each bench function calls this ONCE, before its timed group — so the
77/// `OnceLock` the benches used to declare at file scope memoized a lookup that
78/// already happened exactly once. It is a local here, which keeps
79/// `raw_function_cached`'s signature satisfied without a process global per
80/// kernel per bench binary.
81pub fn get_kernel(registry: &'static AtlasRegistry, module: &str, func: &str) -> RawCudaFunc {
82    let cache = OnceLock::new();
83    registry
84        .raw_function_cached(&cache, module, func)
85        .unwrap_or_else(|e| panic!("Kernel {module}::{func} not found: {e}"))
86}
87
88/// Launch a kernel with raw parameters.
89///
90/// # Safety
91/// `params` must contain valid pointers matching the kernel signature.
92pub unsafe fn launch(
93    registry: &AtlasRegistry,
94    func: RawCudaFunc,
95    grid: (u32, u32, u32),
96    block: (u32, u32, u32),
97    shared_mem: u32,
98    stream: u64,
99    params: &mut [*mut c_void],
100) -> Result<()> {
101    let cfg = LaunchConfig {
102        grid_dim: grid,
103        block_dim: block,
104        shared_mem_bytes: shared_mem,
105    };
106    unsafe { registry.launch_on_stream(func, cfg, stream, params) }
107        .map_err(|e| anyhow::anyhow!("Kernel launch failed: {e}"))
108}
109
110/// Measure kernel execution time in milliseconds using CUDA events.
111/// Runs `warmup` warmup iterations, then `iters` timed iterations,
112/// returning the minimum time across 3 rounds.
113pub fn bench_kernel_ms(
114    stream: u64,
115    warmup: usize,
116    iters: usize,
117    mut kernel_fn: impl FnMut(),
118) -> f32 {
119    // Create CUDA events
120    let mut start: u64 = 0;
121    let mut end: u64 = 0;
122    unsafe {
123        cuEventCreate(&mut start, 0);
124        cuEventCreate(&mut end, 0);
125    }
126
127    let mut best_ms = f32::MAX;
128
129    for _ in 0..3 {
130        // Warmup
131        for _ in 0..warmup {
132            kernel_fn();
133        }
134        gpu_sync(stream).unwrap();
135
136        // Timed run
137        unsafe { cuEventRecord(start, stream) };
138        for _ in 0..iters {
139            kernel_fn();
140        }
141        unsafe { cuEventRecord(end, stream) };
142        unsafe { cuEventSynchronize(end) };
143
144        let mut elapsed_ms: f32 = 0.0;
145        unsafe { cuEventElapsedTime(&mut elapsed_ms, start, end) };
146        let per_iter = elapsed_ms / iters as f32;
147        if per_iter < best_ms {
148            best_ms = per_iter;
149        }
150    }
151
152    unsafe {
153        cuEventDestroy_v2(start);
154        cuEventDestroy_v2(end);
155    }
156
157    best_ms
158}
atlas_spark_bench/gpu.rs

atlas_spark_bench/
gpu.rs
