Skip to main content

pfx_gpu/pace/
trace.rs

1use std::cell::Cell;
2use std::io::Write;
3use std::panic::Location;
4use std::path::PathBuf;
5use std::sync::{Mutex, OnceLock};
6use std::time::Instant;
7
8use super::{Frozen, Moment, Timing};
9
10thread_local! {
11    static NEXT_SLICE: Cell<u8> = const { Cell::new(0) };
12}
13
14pub const TRACE_VAR: &str = "PFX_TURN_TRACE";
15
16pub const COLUMNS: &str =
17    "at_ms\tkind\tsite\twork_ms\tlongest_ms\tgap_ms\tslices\tgpu\twall\tgiven\tuntimed\ttest";
18
19static EPOCH: OnceLock<Instant> = OnceLock::new();
20static LAST: Mutex<Option<Moment>> = Mutex::new(None);
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub enum Source {
24    Gpu,
25    Wall,
26    Given,
27    Untimed,
28}
29
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct Sources {
32    pub gpu: u32,
33    pub wall: u32,
34    pub given: u32,
35    pub untimed: u32,
36}
37
38impl Sources {
39    pub fn add(&mut self, source: Source) {
40        match source {
41            Source::Gpu => self.gpu += 1,
42            Source::Wall => self.wall += 1,
43            Source::Given => self.given += 1,
44            Source::Untimed => self.untimed += 1,
45        }
46    }
47
48    pub fn slices(&self) -> u32 {
49        self.gpu + self.wall + self.given + self.untimed
50    }
51}
52
53fn path() -> Option<PathBuf> {
54    std::env::var_os(TRACE_VAR)
55        .filter(|path| !path.is_empty())
56        .map(PathBuf::from)
57}
58
59pub(super) fn arm() {
60    if path().is_none() {
61        return;
62    }
63    EPOCH.get_or_init(Instant::now);
64    let mut last = LAST.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
65    if last.is_none() {
66        *last = Some(Moment::now(&Frozen::from_env()));
67    }
68}
69
70pub(super) fn turned() {
71    if path().is_none() {
72        return;
73    }
74    let mut last = LAST.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
75    *last = Some(Moment::now(&Frozen::from_env()));
76}
77
78pub(super) fn created(site: &Location<'_>) {
79    record("new", site, 0.0, 0.0, Sources::default());
80}
81
82pub(super) fn slice_timed(timing: Timing) {
83    NEXT_SLICE.set(match timing {
84        Timing::Gpu => 1,
85        Timing::Cpu => 2,
86    });
87}
88
89pub(super) fn take_slice() -> Option<Timing> {
90    match NEXT_SLICE.replace(0) {
91        1 => Some(Timing::Gpu),
92        2 => Some(Timing::Cpu),
93        _ => None,
94    }
95}
96
97pub(super) fn record(
98    kind: &str,
99    site: &Location<'_>,
100    work_ms: f64,
101    longest_ms: f64,
102    sources: Sources,
103) {
104    let Some(path) = path() else {
105        return;
106    };
107    record_to(&path, kind, site, work_ms, longest_ms, sources);
108}
109
110fn record_to(
111    path: &std::path::Path,
112    kind: &str,
113    site: &Location<'_>,
114    work_ms: f64,
115    longest_ms: f64,
116    sources: Sources,
117) {
118    let epoch = *EPOCH.get_or_init(Instant::now);
119    let now = Moment::now(&Frozen::from_env());
120    let gap_ms = {
121        let last = LAST.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
122        last.map_or(0.0, |earlier| now.ms_since(earlier))
123    };
124    let at_ms = now.at.saturating_duration_since(epoch).as_secs_f64() * 1000.0;
125    let file = std::fs::OpenOptions::new()
126        .create(true)
127        .append(true)
128        .open(path);
129    let Ok(mut file) = file else {
130        return;
131    };
132    let line = format!(
133        "{at_ms:.1}\t{kind}\t{}:{}\t{work_ms:.3}\t{longest_ms:.3}\t{gap_ms:.1}\t{}\t{}\t{}\t{}\t{}\t{}\n",
134        site.file(),
135        site.line(),
136        sources.slices(),
137        sources.gpu,
138        sources.wall,
139        sources.given,
140        sources.untimed,
141        std::thread::current().name().unwrap_or("-"),
142    );
143    let _ = file.write_all(line.as_bytes());
144}
145
146#[derive(Clone, Debug, PartialEq)]
147pub struct Row {
148    pub at_ms: f64,
149    pub kind: String,
150    pub site: String,
151    pub work_ms: f64,
152    pub longest_ms: f64,
153    pub gap_ms: f64,
154    pub slices: u32,
155    pub sources: Sources,
156    pub test: String,
157}
158
159pub fn parse(text: &str) -> Vec<Row> {
160    text.lines()
161        .filter_map(|line| {
162            let fields: Vec<&str> = line.split('\t').collect();
163            if fields.len() != 12 {
164                return None;
165            }
166            let number = |index: usize| fields[index].parse::<f64>().ok();
167            let count = |index: usize| fields[index].parse::<u32>().ok();
168            Some(Row {
169                at_ms: number(0)?,
170                kind: fields[1].to_owned(),
171                site: fields[2].to_owned(),
172                work_ms: number(3)?,
173                longest_ms: number(4)?,
174                gap_ms: number(5)?,
175                slices: count(6)?,
176                sources: Sources {
177                    gpu: count(7)?,
178                    wall: count(8)?,
179                    given: count(9)?,
180                    untimed: count(10)?,
181                },
182                test: fields[11].to_owned(),
183            })
184        })
185        .collect()
186}
187
188#[derive(Clone, Debug, PartialEq)]
189pub struct Finding {
190    pub site: String,
191    pub test: String,
192    pub longest_ms: f64,
193    pub gap_ms: f64,
194}
195
196pub fn findings(rows: &[Row], slice_ms: f64, gap_ms: f64) -> Vec<Finding> {
197    rows.iter()
198        .filter(|row| row.longest_ms > slice_ms || row.gap_ms > gap_ms)
199        .map(|row| Finding {
200            site: row.site.clone(),
201            test: row.test.clone(),
202            longest_ms: row.longest_ms,
203            gap_ms: row.gap_ms,
204        })
205        .collect()
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    const TEXT: &str = "0.0\tturns\ta.rs:1\t12.000\t4.000\t0.0\t3\t3\t0\t0\t0\tone\n\
213        90.0\tturns\tb.rs:2\t600.000\t590.000\t90.0\t1\t0\t1\t0\t0\ttwo\n\
214        900.0\tpace\tc.rs:3\t1.000\t1.000\t810.0\t1\t1\t0\t0\t0\tthree\n\
215        not a row\n";
216
217    #[test]
218    fn rows_parse_and_junk_is_skipped() {
219        let rows = parse(TEXT);
220        assert_eq!(rows.len(), 3);
221        assert_eq!(rows[1].site, "b.rs:2");
222        assert_eq!(rows[1].sources.wall, 1);
223        assert_eq!(rows[2].test, "three");
224        assert_eq!(rows[0].sources.slices(), 3);
225    }
226
227    #[test]
228    fn findings_name_long_slices_and_long_gaps() {
229        let found = findings(&parse(TEXT), 300.0, 500.0);
230        assert_eq!(found.len(), 2);
231        assert_eq!(found[0].site, "b.rs:2");
232        assert_eq!(found[1].site, "c.rs:3");
233        assert!(findings(&parse(TEXT), 1000.0, 1000.0).is_empty());
234    }
235
236    fn scratch(name: &str) -> PathBuf {
237        let dir =
238            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../tmp/turn-trace-tests");
239        std::fs::create_dir_all(&dir).unwrap();
240        let path = dir.join(name);
241        let _ = std::fs::remove_file(&path);
242        path
243    }
244
245    #[test]
246    fn a_recorded_turn_reads_back_with_its_site() {
247        let path = scratch("record.tsv");
248        let mut sources = Sources::default();
249        sources.add(Source::Gpu);
250        sources.add(Source::Wall);
251        record_to(&path, "turns", Location::caller(), 7.5, 4.25, sources);
252        let rows = parse(&std::fs::read_to_string(&path).unwrap());
253        assert_eq!(rows.len(), 1);
254        assert_eq!(rows[0].kind, "turns");
255        assert!(rows[0].site.contains("trace.rs:"));
256        assert_eq!((rows[0].work_ms, rows[0].longest_ms), (7.5, 4.25));
257        assert_eq!((rows[0].slices, rows[0].sources), (2, sources));
258    }
259
260    #[test]
261    fn a_paced_add_is_classified_by_the_pacers_timing() {
262        let mut turns = super::super::Turns::new(1e9, false).with_wall_ms(1e9);
263        slice_timed(Timing::Gpu);
264        turns.add(1.0);
265        slice_timed(Timing::Cpu);
266        turns.add(1.0);
267        turns.add(1.0);
268        turns.add(f64::NAN);
269        slice_timed(Timing::Gpu);
270        take_slice();
271        turns.add(1.0);
272        assert_eq!(
273            turns.sources,
274            Sources {
275                gpu: 1,
276                wall: 1,
277                given: 2,
278                untimed: 1
279            }
280        );
281    }
282
283    #[test]
284    fn a_poll_turns_only_once_the_wall_time_is_up() {
285        let mut turns = super::super::Turns::new(1e9, false).with_wall_ms(1e9);
286        assert!(!turns.poll());
287        assert_eq!(turns.taken(), 0);
288        let mut turns = super::super::Turns::new(1e9, false).with_wall_ms(1.0);
289        std::thread::sleep(std::time::Duration::from_millis(5));
290        assert!(turns.poll());
291        assert_eq!(turns.taken(), 1);
292    }
293
294    #[test]
295    fn sources_count_by_kind() {
296        let mut sources = Sources::default();
297        sources.add(Source::Gpu);
298        sources.add(Source::Gpu);
299        sources.add(Source::Untimed);
300        assert_eq!((sources.gpu, sources.untimed, sources.slices()), (2, 1, 3));
301    }
302
303    const CHILD_VAR: &str = "PFX_TURN_TRACE_CHILD";
304
305    #[test]
306    #[ignore = "needs a GPU; run under pgpu"]
307    fn a_paced_gpu_run_traces_no_long_slice_and_no_long_gap() {
308        if std::env::var_os(CHILD_VAR).is_some() {
309            traced_workload();
310            return;
311        }
312        let path = scratch("paced_run.tsv");
313        let status = std::process::Command::new(std::env::current_exe().unwrap())
314            .args([
315                &format!(
316                    "{}::a_paced_gpu_run_traces_no_long_slice_and_no_long_gap",
317                    module_path!().trim_start_matches("pfx_gpu::")
318                ),
319                "--exact",
320                "--ignored",
321                "--test-threads=1",
322            ])
323            .env(TRACE_VAR, &path)
324            .env(CHILD_VAR, "1")
325            .status()
326            .unwrap();
327        assert!(status.success());
328        assert!(path.exists(), "the child ran no test");
329        let rows = parse(&std::fs::read_to_string(&path).unwrap());
330        assert!(rows.len() >= 2, "the run traced {} turns", rows.len());
331        assert!(
332            rows.iter()
333                .any(|row| row.sources.gpu + row.sources.wall > 0)
334        );
335        let found = findings(&rows, 300.0, 500.0);
336        assert!(found.is_empty(), "{found:#?}");
337    }
338
339    fn traced_workload() {
340        use super::super::{Pacer, Slice, Turns, Work};
341        const GROUPS: u32 = 600;
342        let gpu = pollster::block_on(crate::Gpu::headless()).unwrap();
343        let device = &gpu.device;
344        let mut turns = Turns::default();
345        let shader = device.create_shader_module(crate::wgpu::ShaderModuleDescriptor {
346            label: Some("trace test"),
347            source: crate::wgpu::ShaderSource::Wgsl(
348                "@group(0) @binding(0) var<storage, read_write> out: array<u32>;
349                 @compute @workgroup_size(64)
350                 fn main(@builtin(global_invocation_id) id: vec3<u32>) {
351                     var a = id.x;
352                     for (var i = 0u; i < 60000u; i = i + 1u) {
353                         a = a * 1664525u + 1013904223u;
354                     }
355                     out[id.x] = a;
356                 }"
357                .into(),
358            ),
359        });
360        let pipeline = device.create_compute_pipeline(&crate::wgpu::ComputePipelineDescriptor {
361            label: Some("trace test"),
362            layout: None,
363            module: &shader,
364            entry_point: Some("main"),
365            compilation_options: Default::default(),
366            cache: None,
367        });
368        let out = device.create_buffer(&crate::wgpu::BufferDescriptor {
369            label: Some("out"),
370            size: u64::from(GROUPS) * 64 * 4,
371            usage: crate::wgpu::BufferUsages::STORAGE,
372            mapped_at_creation: false,
373        });
374        let group = device.create_bind_group(&crate::wgpu::BindGroupDescriptor {
375            label: None,
376            layout: &pipeline.get_bind_group_layout(0),
377            entries: &[crate::wgpu::BindGroupEntry {
378                binding: 0,
379                resource: out.as_entire_binding(),
380            }],
381        });
382        for _ in 0..3 {
383            let mut pacer = Pacer::default();
384            pacer
385                .run(
386                    "trace test",
387                    device,
388                    &gpu.queue,
389                    Work::Dispatches { count: GROUPS },
390                    |encoder, slice| {
391                        let Slice::Dispatches { count, .. } = slice else {
392                            unreachable!()
393                        };
394                        let mut pass = encoder.begin_compute_pass(&Default::default());
395                        pass.set_pipeline(&pipeline);
396                        pass.set_bind_group(0, &group, &[]);
397                        pass.dispatch_workgroups(count, 1, 1);
398                    },
399                    |ms| turns.add(ms),
400                )
401                .unwrap();
402            turns.time(|| {
403                let mut encoder = device.create_command_encoder(&Default::default());
404                {
405                    let mut pass = encoder.begin_compute_pass(&Default::default());
406                    pass.set_pipeline(&pipeline);
407                    pass.set_bind_group(0, &group, &[]);
408                    pass.dispatch_workgroups(16, 1, 1);
409                }
410                let index = gpu.queue.submit(Some(encoder.finish()));
411                let _ = device.poll(crate::wgpu::PollType::Wait {
412                    submission_index: Some(index),
413                    timeout: None,
414                });
415            });
416        }
417        turns.turn();
418    }
419}