Skip to main content

cherenkov/qwen4_exp/gpu/activity/
reads.rs

1//! Bounded read counters. Reader threads publish once per completed record.
2
3use serde::{Deserialize, Serialize};
4use std::sync::{
5    Arc, Mutex,
6    atomic::{AtomicBool, Ordering},
7};
8use std::time::Instant;
9
10#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
11pub struct ReadStats {
12    pub requested_reads: u64,
13    pub requested_bytes: u64,
14    pub completed_reads: u64,
15    pub failed_reads: u64,
16    pub completed_bytes: u64,
17    pub read_seconds: f64,
18    pub max_read_seconds: f64,
19}
20
21impl ReadStats {
22    pub fn add(&mut self, other: Self) {
23        self.requested_reads += other.requested_reads;
24        self.requested_bytes += other.requested_bytes;
25        self.completed_reads += other.completed_reads;
26        self.failed_reads += other.failed_reads;
27        self.completed_bytes += other.completed_bytes;
28        self.read_seconds += other.read_seconds;
29        self.max_read_seconds = self.max_read_seconds.max(other.max_read_seconds);
30    }
31}
32
33#[derive(Debug, Clone, Copy)]
34pub enum ReadSource {
35    Prefill,
36    Demand,
37    Prefetch,
38}
39
40#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
41pub struct ReadSources {
42    pub prefill: ReadStats,
43    pub demand: ReadStats,
44    pub prefetch: ReadStats,
45}
46
47impl ReadSources {
48    fn source(&mut self, source: ReadSource) -> &mut ReadStats {
49        match source {
50            ReadSource::Prefill => &mut self.prefill,
51            ReadSource::Demand => &mut self.demand,
52            ReadSource::Prefetch => &mut self.prefetch,
53        }
54    }
55
56    pub fn add(&mut self, other: Self) {
57        self.prefill.add(other.prefill);
58        self.demand.add(other.demand);
59        self.prefetch.add(other.prefetch);
60    }
61
62    pub fn total(self) -> ReadStats {
63        let mut total = self.prefill;
64
65        total.add(self.demand);
66        total.add(self.prefetch);
67
68        total
69    }
70}
71
72#[derive(Debug, Clone)]
73pub struct ReadTracker(Arc<Mutex<Vec<[ReadSources; 3]>>>);
74
75impl ReadTracker {
76    pub fn new(layers: usize) -> Self {
77        Self(Arc::new(Mutex::new(vec![
78            [ReadSources::default(); 3];
79            layers
80        ])))
81    }
82
83    pub fn snapshot(&self) -> Vec<[ReadSources; 3]> {
84        self.0.lock().unwrap().clone()
85    }
86
87    pub fn ticket(&self, layer: usize, kind: u8, source: ReadSource, bytes: usize) -> ReadTicket {
88        let mut layers = self.0.lock().unwrap();
89        let stats = layers[layer][kind as usize].source(source);
90        stats.requested_reads += 1;
91        stats.requested_bytes += bytes as u64;
92
93        ReadTicket {
94            tracker: self.clone(),
95            layer,
96            kind,
97            source,
98            bytes,
99            finished: Arc::new(AtomicBool::new(false)),
100        }
101    }
102}
103
104#[derive(Debug, Clone)]
105pub struct ReadTicket {
106    tracker: ReadTracker,
107    layer: usize,
108    kind: u8,
109    source: ReadSource,
110    bytes: usize,
111    finished: Arc<AtomicBool>,
112}
113
114impl ReadTicket {
115    pub fn done(&self) -> bool {
116        self.finished.load(Ordering::Acquire)
117    }
118
119    pub fn measure(&self, read: impl FnOnce() -> usize) {
120        let started = Instant::now();
121        let bytes = read();
122
123        self.finish(bytes, started.elapsed().as_secs_f64());
124    }
125
126    fn finish(&self, bytes: usize, seconds: f64) {
127        let mut layers = self.tracker.0.lock().unwrap();
128        let stats = layers[self.layer][self.kind as usize].source(self.source);
129        stats.completed_reads += u64::from(bytes == self.bytes);
130        stats.failed_reads += u64::from(bytes != self.bytes);
131        stats.completed_bytes += bytes as u64;
132        stats.read_seconds += seconds;
133        stats.max_read_seconds = stats.max_read_seconds.max(seconds);
134
135        self.finished.store(true, Ordering::Release);
136    }
137}
138
139#[cfg(test)]
140#[path = "../../../../tests/unit/qwen4_exp/gpu/read_stats.rs"]
141mod tests;