aria_inference/
profile.rs

1//! Optional load / generate timings for `--profile`.
2
3use serde::Serialize;
4use std::cell::{Cell, RefCell};
5use std::time::Instant;
6
7thread_local! {
8    static LOAD_ON: Cell<bool> = const { Cell::new(false) };
9    static LOAD: RefCell<LoadProfile> = const { RefCell::new(LoadProfile::empty()) };
10}
11
12#[derive(Debug, Clone, Serialize)]
13pub struct LoadProfile {
14    pub mmap_ms: f64,
15    pub dequant_ms: f64,
16    pub unrotate_ms: f64,
17    pub materialize_ms: f64,
18    pub cuda_upload_ms: f64,
19}
20
21impl LoadProfile {
22    const fn empty() -> Self {
23        Self {
24            mmap_ms: 0.0,
25            dequant_ms: 0.0,
26            unrotate_ms: 0.0,
27            materialize_ms: 0.0,
28            cuda_upload_ms: 0.0,
29        }
30    }
31}
32
33#[derive(Debug, Clone, Serialize, Default)]
34pub struct GenerateProfile {
35    pub prefill_ms: f64,
36    pub decode_ms: f64,
37    pub gemm_attn_ms: f64,
38    pub gemm_ffn_ms: f64,
39    pub gemm_lm_head_ms: f64,
40}
41
42#[derive(Debug, Clone, Serialize)]
43pub struct EngineProfile {
44    pub compute: String,
45    pub load: LoadProfile,
46    pub generate: Option<GenerateProfile>,
47    pub ci_fail: bool,
48}
49
50pub fn load_profile_begin(enabled: bool) {
51    LOAD_ON.with(|c| c.set(enabled));
52    LOAD.with(|l| *l.borrow_mut() = LoadProfile::empty());
53}
54
55pub fn load_profile_enabled() -> bool {
56    LOAD_ON.with(|c| c.get())
57}
58
59pub fn load_profile_add_dequant(ms: f64) {
60    if !load_profile_enabled() {
61        return;
62    }
63    LOAD.with(|l| l.borrow_mut().dequant_ms += ms);
64}
65
66pub fn load_profile_add_unrotate(ms: f64) {
67    if !load_profile_enabled() {
68        return;
69    }
70    LOAD.with(|l| l.borrow_mut().unrotate_ms += ms);
71}
72
73pub fn load_profile_set_mmap(ms: f64) {
74    LOAD.with(|l| l.borrow_mut().mmap_ms = ms);
75}
76
77pub fn load_profile_set_materialize(ms: f64) {
78    LOAD.with(|l| l.borrow_mut().materialize_ms = ms);
79}
80
81pub fn load_profile_set_cuda_upload(ms: f64) {
82    LOAD.with(|l| l.borrow_mut().cuda_upload_ms = ms);
83}
84
85pub fn load_profile_take() -> LoadProfile {
86    LOAD.with(|l| l.borrow().clone())
87}
88
89pub fn elapsed_ms(t0: Instant) -> f64 {
90    t0.elapsed().as_secs_f64() * 1000.0
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn load_accum_when_enabled() {
99        load_profile_begin(true);
100        load_profile_add_dequant(1.5);
101        load_profile_add_unrotate(2.5);
102        let p = load_profile_take();
103        assert!((p.dequant_ms - 1.5).abs() < 1e-9);
104        assert!((p.unrotate_ms - 2.5).abs() < 1e-9);
105        load_profile_begin(false);
106        load_profile_add_dequant(9.0);
107        assert!(load_profile_take().dequant_ms < 1e-9);
108    }
109}
aria_inference/profile.rs

aria_inference/
profile.rs
