hft_benchmarks/
timing.rs

1//! High-resolution timing utilities using CPU timestamp counter
2
3#[cfg(target_arch = "x86_64")]
4use core::arch::x86_64::{_rdtsc, _mm_lfence, _mm_mfence};
5
6/// Read timestamp counter for the current architecture
7#[cfg(target_arch = "x86_64")]
8#[inline(always)]
9unsafe fn read_timestamp_with_fences() -> u64 {
10    _mm_mfence();
11    let tsc = _rdtsc();
12    _mm_lfence();
13    tsc
14}
15
16#[cfg(target_arch = "aarch64")]
17#[inline(always)]
18unsafe fn read_timestamp_with_fences() -> u64 {
19    read_timestamp()
20}
21
22#[cfg(target_arch = "aarch64")]
23#[inline(always)]
24pub fn read_timestamp() -> u64 {
25    read_virtual_counter()
26}
27
28#[cfg(target_arch = "aarch64")]
29#[inline(always)]
30fn read_virtual_counter() -> u64 {
31    unsafe {
32        let counter: u64;
33        std::arch::asm!("mrs {}, cntvct_el0", out(reg) counter, options(nomem, nostack));
34        counter
35    }
36}
37
38#[cfg(target_arch = "aarch64")]
39fn get_counter_frequency() -> u64 {
40    unsafe {
41        let freq: u64;
42        std::arch::asm!("mrs {}, cntfrq_el0", out(reg) freq, options(nomem, nostack));
43        freq
44    }
45}
46
47pub struct PrecisionTimer {
48    start: u64,
49    frequency_mhz: u64,
50}
51
52impl PrecisionTimer {
53    #[inline(always)]
54    pub fn start() -> Self {
55        unsafe {
56            let start = read_timestamp_with_fences();
57            
58            Self {
59                start,
60                frequency_mhz: crate::mock_core::cpu_frequency_mhz(),
61            }
62        }
63    }
64    
65    #[inline(always)]
66    pub fn stop(self) -> u64 {
67        unsafe {
68            let end = read_timestamp_with_fences();
69            
70            let cycles = end - self.start;
71            if self.frequency_mhz == 0 {
72                cycles
73            } else {
74                #[cfg(target_arch = "x86_64")]
75                {
76                    (cycles * 1000) / self.frequency_mhz
77                }
78                #[cfg(target_arch = "aarch64")]
79                {
80                    let counter_freq = get_counter_frequency();
81                    if counter_freq > 0 {
82                        if cycles == 0 {
83                            0
84                        } else {
85                            std::cmp::max(1, (cycles * 1_000_000_000) / counter_freq)
86                        }
87                    } else {
88                        (cycles * 1000) / self.frequency_mhz
89                    }
90                }
91            }
92        }
93    }
94}
95
96pub fn time_function<F, R>(f: F) -> (R, u64)
97where
98    F: FnOnce() -> R,
99{
100    let timer = PrecisionTimer::start();
101    let result = f();
102    let elapsed = timer.stop();
103    (result, elapsed)
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::calibration::calibrate_tsc_frequency;
110    
111    #[test]
112    fn test_precision_timer() {
113        calibrate_tsc_frequency();
114        
115        let timer = PrecisionTimer::start();
116        let _ = (0..10).sum::<i32>();
117        let elapsed = timer.stop();
118        
119        #[cfg(target_arch = "x86_64")]
120        {
121            assert!(elapsed < 1000, "Elapsed time too high: {}ns", elapsed);
122        }
123        #[cfg(target_arch = "aarch64")]
124        {
125            assert!(elapsed < 5000, "Elapsed time too high: {}ns", elapsed);
126        }
127    }
128    
129    #[test]
130    fn test_time_function() {
131        calibrate_tsc_frequency();
132        
133        let (result, elapsed) = time_function(|| {
134            (0..100).sum::<i32>()
135        });
136        
137        assert_eq!(result, 4950);
138        assert!(elapsed < 10000, "Function took too long: {}ns", elapsed);
139    }
140}