iro_cuda_ffi_profile/

timer.rs

//! GPU timing utilities with reusable events.
//!
//! This module provides low-overhead timing primitives that avoid repeated
//! event allocation in hot loops.

use iro_cuda_ffi::event::{Event, EventKind};
use iro_cuda_ffi::prelude::{Result, Stream};

/// Reusable GPU timer for measuring kernel execution time.
///
/// Unlike creating new events for each timing, `GpuTimer` reuses its internal
/// events to avoid allocation overhead in hot loops.
///
/// # Example
///
/// ```ignore
/// use iro_cuda_ffi_profile::GpuTimer;
///
/// let timer = GpuTimer::new()?;
/// let mut times = Vec::new();
///
/// for _ in 0..100 {
///     timer.start(&stream)?;
///     my_kernel(&stream, ...)?;
///     times.push(timer.stop_sync(&stream)?);
/// }
/// ```
#[derive(Debug)]
pub struct GpuTimer {
    start: Event,
    end: Event,
}

impl GpuTimer {
    /// Creates a new GPU timer with reusable timed events.
    pub fn new() -> Result<Self> {
        Ok(Self {
            start: Event::new(EventKind::Timed)?,
            end: Event::new(EventKind::Timed)?,
        })
    }

    /// Records the start timestamp in the given stream.
    ///
    /// Call this before the operation(s) you want to time.
    #[inline]
    pub fn start(&self, stream: &Stream) -> Result<()> {
        stream.record_event(&self.start)
    }

    /// Records the end timestamp in the given stream.
    ///
    /// Call this after the operation(s) you want to time.
    /// Does not synchronize - use `elapsed()` after manual sync,
    /// or use `stop_sync()` for convenience.
    #[inline]
    pub fn stop(&self, stream: &Stream) -> Result<()> {
        stream.record_event(&self.end)
    }

    /// Records end timestamp, synchronizes the end event, and returns elapsed milliseconds.
    ///
    /// This is the most common usage pattern for timing a single operation.
    #[inline]
    pub fn stop_sync(&self, stream: &Stream) -> Result<f32> {
        stream.record_event(&self.end)?;
        self.end.synchronize()?;
        self.end.elapsed_since(&self.start)
    }

    /// Returns elapsed milliseconds between start and end events.
    ///
    /// Both events must have been recorded and completed before calling this.
    /// If you haven't synchronized, call `stream.synchronize()` or
    /// `self.end.synchronize()` first.
    #[inline]
    pub fn elapsed(&self) -> Result<f32> {
        self.end.elapsed_since(&self.start)
    }

    /// Times a closure, returning its result and elapsed milliseconds.
    ///
    /// Equivalent to:
    /// ```ignore
    /// timer.start(&stream)?;
    /// let result = f()?;
    /// let ms = timer.stop_sync(&stream)?;
    /// ```
    #[inline]
    pub fn time<T, F>(&self, stream: &Stream, f: F) -> Result<(T, f32)>
    where
        F: FnOnce() -> Result<T>,
    {
        self.start(stream)?;
        let result = f()?;
        let ms = self.stop_sync(stream)?;
        Ok((result, ms))
    }

    /// Times a closure that doesn't return a Result.
    ///
    /// Useful for timing infallible operations.
    #[inline]
    pub fn time_infallible<T, F>(&self, stream: &Stream, f: F) -> Result<(T, f32)>
    where
        F: FnOnce() -> T,
    {
        self.start(stream)?;
        let result = f();
        let ms = self.stop_sync(stream)?;
        Ok((result, ms))
    }
}

/// Extension trait for convenient one-shot timing on streams.
///
/// For repeated timing in loops, prefer `GpuTimer` to avoid event allocation overhead.
pub trait StreamTimingExt {
    /// Times a closure, returning its result and elapsed milliseconds.
    ///
    /// Creates temporary events for timing. For hot loops, use `GpuTimer` instead.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use iro_cuda_ffi_profile::StreamTimingExt;
    ///
    /// let (_, ms) = stream.timed(|| {
    ///     my_kernel(&stream, ...)?;
    ///     Ok(())
    /// })?;
    /// println!("Kernel took {ms:.3} ms");
    /// ```
    fn timed<T, F>(&self, f: F) -> Result<(T, f32)>
    where
        F: FnOnce() -> Result<T>;

    /// Times a closure, discarding the result and returning only elapsed milliseconds.
    fn timed_ms<F>(&self, f: F) -> Result<f32>
    where
        F: FnOnce() -> Result<()>;
}

impl StreamTimingExt for Stream {
    fn timed<T, F>(&self, f: F) -> Result<(T, f32)>
    where
        F: FnOnce() -> Result<T>,
    {
        let start = self.record_timed_event()?;
        let result = f()?;
        let end = self.record_timed_event()?;
        // Synchronize only the end event, not the entire stream.
        // This avoids waiting for unrelated work that may have been
        // queued after the end event was recorded.
        end.synchronize()?;
        let ms = end.elapsed_since(&start)?;
        Ok((result, ms))
    }

    fn timed_ms<F>(&self, f: F) -> Result<f32>
    where
        F: FnOnce() -> Result<()>,
    {
        let ((), ms) = self.timed(f)?;
        Ok(ms)
    }
}

/// Accumulator for collecting timing samples with minimal overhead.
///
/// Pre-allocates capacity and provides efficient sample collection.
#[derive(Debug)]
pub struct TimingSamples {
    samples: Vec<f64>,
}

impl TimingSamples {
    /// Creates a new accumulator with pre-allocated capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            samples: Vec::with_capacity(capacity),
        }
    }

    /// Adds a timing sample in milliseconds.
    #[inline]
    pub fn push(&mut self, ms: f32) {
        self.samples.push(ms as f64);
    }

    /// Adds a timing sample in milliseconds (f64).
    #[inline]
    pub fn push_f64(&mut self, ms: f64) {
        self.samples.push(ms);
    }

    /// Returns the collected samples.
    #[inline]
    pub fn as_slice(&self) -> &[f64] {
        &self.samples
    }

    /// Returns the number of samples collected.
    #[inline]
    pub fn len(&self) -> usize {
        self.samples.len()
    }

    /// Returns true if no samples have been collected.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.samples.is_empty()
    }

    /// Clears all samples, keeping allocated capacity.
    #[inline]
    pub fn clear(&mut self) {
        self.samples.clear();
    }

    /// Computes statistics from the collected samples.
    ///
    /// # Panics
    ///
    /// Panics if no samples have been collected.
    pub fn stats(&self) -> crate::Stats {
        crate::Stats::from_samples(&self.samples)
    }
}

impl From<TimingSamples> for Vec<f64> {
    fn from(samples: TimingSamples) -> Self {
        samples.samples
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_timing_samples() {
        let mut samples = TimingSamples::with_capacity(10);
        assert!(samples.is_empty());

        samples.push(1.0);
        samples.push(2.0);
        samples.push(3.0);

        assert_eq!(samples.len(), 3);
        assert_eq!(samples.as_slice(), &[1.0, 2.0, 3.0]);

        let stats = samples.stats();
        assert_eq!(stats.count, 3);
        assert!((stats.mean - 2.0).abs() < 1e-10);
    }
}