rustytorch/demos/rtx-inference-profiler/src/profiler.rs

//! Inference profiler implementation.
//!
//! Profiles real CPU tensor compute (see [`crate::bench_model`]) rather than
//! simulated timings.

use crate::bench_model::ProfiledModel;
use crate::error::ProfilerError;
use inference_profiler_shared::{
    DeviceType, LatencyMetrics, MemoryMetrics, ModelType, ProfileConfig, ProfileResult,
    ThroughputMetrics,
};
use std::time::Instant;

/// Inference profiler for benchmarking model performance.
pub struct InferenceProfiler {
    config: Option<ProfileConfig>,
    results: Vec<ProfileResult>,
}

impl InferenceProfiler {
    /// Creates a new profiler instance.
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: None,
            results: Vec::new(),
        }
    }

    /// Initializes the profiler with a configuration.
    ///
    /// # Errors
    /// Returns an error if the configuration is invalid.
    pub fn initialize(&mut self, config: ProfileConfig) -> Result<(), ProfilerError> {
        config.validate().map_err(ProfilerError::ConfigError)?;
        self.config = Some(config);
        self.results.clear();
        Ok(())
    }

    /// Runs the profiling benchmark for all configured batch sizes.
    ///
    /// # Errors
    /// Returns an error if not initialized or if profiling fails.
    pub fn run_profile(&mut self) -> Result<Vec<ProfileResult>, ProfilerError> {
        let config = self
            .config
            .as_ref()
            .ok_or_else(|| ProfilerError::InternalError("profiler not initialized".to_string()))?
            .clone();

        self.results.clear();

        for &batch_size in &config.batch_sizes {
            let result = self.profile_batch_size(&config, batch_size)?;
            self.results.push(result);
        }

        Ok(self.results.clone())
    }

    /// Profiles a single batch size.
    fn profile_batch_size(
        &self,
        config: &ProfileConfig,
        batch_size: usize,
    ) -> Result<ProfileResult, ProfilerError> {
        let model = ProfiledModel::new(config.model_type, config.device, batch_size);

        // Warmup phase (real compute, to warm caches / avoid first-call overhead)
        for _ in 0..config.warmup_iterations {
            let _ = model.forward();
        }

        // Benchmark phase
        let latency_measurements = self.measure_latency(&model, config.benchmark_iterations)?;
        let latency = LatencyMetrics::from_measurements(&latency_measurements)
            .map_err(ProfilerError::MeasurementError)?;

        let memory = self.measure_memory(&model)?;
        let throughput = ThroughputMetrics::from_latency(batch_size, latency.mean_ms)
            .map_err(ProfilerError::MeasurementError)?;

        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_err(|e| ProfilerError::InternalError(e.to_string()))?
            .as_secs();

        Ok(ProfileResult::new(
            config.model_type,
            config.device,
            batch_size,
            latency,
            memory,
            throughput,
            timestamp,
        ))
    }

    /// Measures latency over multiple iterations of real forward passes.
    ///
    /// # Errors
    /// Returns an error if measurement fails.
    pub fn measure_latency(
        &self,
        model: &ProfiledModel,
        iterations: usize,
    ) -> Result<Vec<f64>, ProfilerError> {
        if iterations == 0 {
            return Err(ProfilerError::InvalidInput(
                "iterations cannot be zero".to_string(),
            ));
        }

        let mut measurements = Vec::with_capacity(iterations);

        for _ in 0..iterations {
            let start = Instant::now();
            let _ = model.forward();
            let duration = start.elapsed();

            let latency_ms = duration.as_secs_f64() * 1000.0;
            measurements.push(latency_ms);
        }

        Ok(measurements)
    }

    /// Measures memory usage from real tensor allocation sizes.
    ///
    /// # Errors
    /// Returns an error if measurement fails.
    pub fn measure_memory(&self, model: &ProfiledModel) -> Result<MemoryMetrics, ProfilerError> {
        let memory_mb = model.memory_usage_mb();

        // Peak/reserved are still derived estimates (allocator headroom), since
        // this demo does not instrument the system allocator directly; the
        // baseline `allocated_mb` figure itself is real tensor byte accounting.
        let peak = memory_mb * 1.1;
        let allocated = memory_mb;
        let reserved = memory_mb * 1.2;

        MemoryMetrics::new(peak, allocated, reserved).map_err(ProfilerError::MeasurementError)
    }

    /// Gets the current results.
    #[must_use]
    pub fn get_results(&self) -> &[ProfileResult] {
        &self.results
    }

    /// Resets the profiler state.
    pub fn reset(&mut self) {
        self.config = None;
        self.results.clear();
    }

    /// Checks if a device is available.
    ///
    /// Note: this demo only wires up real CPU tensor compute via
    /// `rtx-tensor`. `DeviceType::CUDA`/`DeviceType::Metal` are accepted for
    /// API compatibility but currently execute the same CPU compute path
    /// (see [`crate::bench_model`]), so this always reports available.
    #[must_use]
    pub fn is_device_available(&self, _device: DeviceType) -> bool {
        true
    }

    /// Gets all "available" devices.
    ///
    /// See [`Self::is_device_available`] caveat: `CUDA`/`Metal` are listed
    /// for API compatibility but run identical CPU compute in this demo.
    #[must_use]
    pub fn get_available_devices(&self) -> Vec<DeviceType> {
        vec![DeviceType::CPU, DeviceType::CUDA, DeviceType::Metal]
    }
}

impl Default for InferenceProfiler {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use inference_profiler_shared::InputShape;

    #[test]
    fn test_profiler_creation() {
        let profiler = InferenceProfiler::new();
        assert!(profiler.config.is_none());
        assert_eq!(profiler.results.len(), 0);
    }

    #[test]
    fn test_profiler_default() {
        let profiler = InferenceProfiler::default();
        assert!(profiler.config.is_none());
    }

    #[test]
    fn test_profiler_initialize_success() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::default();

        let result = profiler.initialize(config.clone());
        assert!(result.is_ok());
        assert!(profiler.config.is_some());
        assert_eq!(profiler.config.unwrap(), config);
    }

    #[test]
    fn test_profiler_initialize_invalid_config() {
        let mut profiler = InferenceProfiler::new();

        // Create invalid config (empty batch sizes)
        let result = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![],
            10,
            100,
            InputShape::default(),
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_profiler_initialize_clears_results() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::default();

        profiler.initialize(config.clone()).unwrap();
        profiler.run_profile().unwrap();

        assert!(!profiler.results.is_empty());

        profiler.initialize(config).unwrap();
        assert!(profiler.results.is_empty());
    }

    #[test]
    fn test_measure_latency_success() {
        let profiler = InferenceProfiler::new();
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);

        let measurements = profiler.measure_latency(&model, 10).unwrap();

        assert_eq!(measurements.len(), 10);
        for &m in &measurements {
            assert!(m > 0.0);
            assert!(m.is_finite());
        }
    }

    #[test]
    fn test_measure_latency_zero_iterations() {
        let profiler = InferenceProfiler::new();
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);

        let result = profiler.measure_latency(&model, 0);

        assert!(result.is_err());
        match result.unwrap_err() {
            ProfilerError::InvalidInput(msg) => {
                assert!(msg.contains("iterations cannot be zero"));
            }
            _ => panic!("wrong error type"),
        }
    }

    #[test]
    fn test_measure_latency_all_positive() {
        let profiler = InferenceProfiler::new();
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);

        let measurements = profiler.measure_latency(&model, 20).unwrap();

        // Real CPU compute: durations are always positive and finite; we no
        // longer assert on a fabricated jitter range.
        for &m in &measurements {
            assert!(m > 0.0);
            assert!(m.is_finite());
        }
    }

    #[test]
    fn test_measure_memory_success() {
        let profiler = InferenceProfiler::new();
        let model = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);

        let memory = profiler.measure_memory(&model).unwrap();

        assert!(memory.peak_memory_mb > 0.0);
        assert!(memory.allocated_mb > 0.0);
        assert!(memory.reserved_mb > 0.0);
        assert!(memory.peak_memory_mb >= memory.allocated_mb);
        assert!(memory.reserved_mb >= memory.allocated_mb);
    }

    #[test]
    fn test_measure_memory_scales_with_batch_size() {
        let profiler = InferenceProfiler::new();
        let model1 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 1);
        let model8 = ProfiledModel::new(ModelType::ResNet18, DeviceType::CPU, 8);

        let memory1 = profiler.measure_memory(&model1).unwrap();
        let memory8 = profiler.measure_memory(&model8).unwrap();

        assert!(memory8.allocated_mb > memory1.allocated_mb);
    }

    #[test]
    fn test_profile_batch_size_success() {
        let profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![4],
            5,
            10,
            InputShape::default(),
        )
        .unwrap();

        let result = profiler.profile_batch_size(&config, 4).unwrap();

        assert_eq!(result.model_type, ModelType::ResNet18);
        assert_eq!(result.device, DeviceType::CPU);
        assert_eq!(result.batch_size, 4);
        assert!(result.latency.mean_ms > 0.0);
        assert!(result.memory.allocated_mb > 0.0);
        assert!(result.throughput.samples_per_sec > 0.0);
    }

    #[test]
    fn test_run_profile_not_initialized() {
        let mut profiler = InferenceProfiler::new();

        let result = profiler.run_profile();

        assert!(result.is_err());
        match result.unwrap_err() {
            ProfilerError::InternalError(msg) => {
                assert!(msg.contains("not initialized"));
            }
            _ => panic!("wrong error type"),
        }
    }

    #[test]
    fn test_run_profile_single_batch_size() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![8],
            5,
            20,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        let results = profiler.run_profile().unwrap();

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].batch_size, 8);
    }

    #[test]
    fn test_run_profile_multiple_batch_sizes() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2, 4, 8],
            5,
            20,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        let results = profiler.run_profile().unwrap();

        assert_eq!(results.len(), 4);
        assert_eq!(results[0].batch_size, 1);
        assert_eq!(results[1].batch_size, 2);
        assert_eq!(results[2].batch_size, 4);
        assert_eq!(results[3].batch_size, 8);
    }

    #[test]
    fn test_run_profile_latency_increases_with_batch_size() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 8],
            5,
            20,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        let results = profiler.run_profile().unwrap();

        assert!(results[1].latency.mean_ms > results[0].latency.mean_ms);
    }

    #[test]
    fn test_run_profile_throughput_calculation() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![8],
            5,
            20,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        let results = profiler.run_profile().unwrap();

        let result = &results[0];
        let expected_throughput = (result.batch_size as f64 * 1000.0) / result.latency.mean_ms;

        assert_relative_eq!(
            result.throughput.samples_per_sec,
            expected_throughput,
            epsilon = 0.1
        );
    }

    #[test]
    fn test_get_results_empty() {
        let profiler = InferenceProfiler::new();
        let results = profiler.get_results();

        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_get_results_after_profile() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            5,
            10,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        profiler.run_profile().unwrap();

        let results = profiler.get_results();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_reset_clears_config() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::default();

        profiler.initialize(config).unwrap();
        assert!(profiler.config.is_some());

        profiler.reset();
        assert!(profiler.config.is_none());
    }

    #[test]
    fn test_reset_clears_results() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::default();

        profiler.initialize(config).unwrap();
        profiler.run_profile().unwrap();
        assert!(!profiler.results.is_empty());

        profiler.reset();
        assert!(profiler.results.is_empty());
    }

    #[test]
    fn test_is_device_available() {
        let profiler = InferenceProfiler::new();

        assert!(profiler.is_device_available(DeviceType::CPU));
        assert!(profiler.is_device_available(DeviceType::CUDA));
        assert!(profiler.is_device_available(DeviceType::Metal));
    }

    #[test]
    fn test_get_available_devices() {
        let profiler = InferenceProfiler::new();
        let devices = profiler.get_available_devices();

        assert!(devices.contains(&DeviceType::CPU));
        assert!(devices.contains(&DeviceType::CUDA));
        assert!(devices.contains(&DeviceType::Metal));
    }

    #[test]
    fn test_run_profile_stores_results() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![4],
            5,
            10,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(config).unwrap();
        profiler.run_profile().unwrap();

        assert_eq!(profiler.results.len(), 1);
        assert_eq!(profiler.get_results().len(), 1);
    }

    #[test]
    fn test_profile_different_models() {
        let mut profiler = InferenceProfiler::new();

        for model_type in [ModelType::ResNet18, ModelType::ResNet50, ModelType::ViTB16] {
            let config = ProfileConfig::new(
                model_type,
                DeviceType::CPU,
                vec![1],
                5,
                10,
                InputShape::default(),
            )
            .unwrap();

            profiler.initialize(config).unwrap();
            let results = profiler.run_profile().unwrap();

            assert_eq!(results.len(), 1);
            assert_eq!(results[0].model_type, model_type);
        }
    }

    #[test]
    fn test_profile_different_devices() {
        let mut profiler = InferenceProfiler::new();

        for device in [DeviceType::CPU, DeviceType::CUDA, DeviceType::Metal] {
            let config = ProfileConfig::new(
                ModelType::ResNet18,
                device,
                vec![1],
                5,
                10,
                InputShape::default(),
            )
            .unwrap();

            profiler.initialize(config).unwrap();
            let results = profiler.run_profile().unwrap();

            assert_eq!(results.len(), 1);
            assert_eq!(results[0].device, device);
        }
    }

    #[test]
    fn test_larger_model_slower_than_smaller_model() {
        // With real compute, model-class ordering is expressed via matmul
        // dimension (see bench_model::base_matmul_dim), and should translate
        // into measurably higher latency for larger model classes.
        let mut profiler = InferenceProfiler::new();

        let small_config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![2],
            3,
            10,
            InputShape::default(),
        )
        .unwrap();
        let large_config = ProfileConfig::new(
            ModelType::ViTL16,
            DeviceType::CPU,
            vec![2],
            3,
            10,
            InputShape::default(),
        )
        .unwrap();

        profiler.initialize(small_config).unwrap();
        let small_results = profiler.run_profile().unwrap();

        profiler.initialize(large_config).unwrap();
        let large_results = profiler.run_profile().unwrap();

        assert!(large_results[0].latency.mean_ms > small_results[0].latency.mean_ms);
    }

    #[test]
    fn test_timestamp_is_recent() {
        let mut profiler = InferenceProfiler::new();
        let config = ProfileConfig::default();

        profiler.initialize(config).unwrap();
        let results = profiler.run_profile().unwrap();

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let timestamp_diff = now.abs_diff(results[0].timestamp);
        assert!(timestamp_diff < 60); // Within 1 minute
    }
}