1#[cfg(target_arch = "x86_64")]
4use core::arch::x86_64::{_rdtsc, _mm_lfence, _mm_mfence};
5
6#[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}