scirs2_fft/

auto_tuning.rs

//! Auto-tuning for hardware-specific FFT optimizations
//!
//! This module provides functionality to automatically tune FFT parameters
//! for optimal performance on the current hardware. It includes:
//!
//! - Benchmarking different FFT configurations
//! - Selecting optimal parameters based on timing results
//! - Persisting tuning results for future use
//! - Detecting CPU features and adapting algorithms accordingly

use rustfft::FftPlanner;
use scirs2_core::numeric::Complex64;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{BufReader, BufWriter};
use std::path::{Path, PathBuf};
use std::time::Instant;

use crate::error::{FFTError, FFTResult};
use crate::plan_serialization::PlanSerializationManager;

/// A range of FFT sizes to benchmark
#[derive(Debug, Clone)]
pub struct SizeRange {
    /// Minimum size to test
    pub min: usize,
    /// Maximum size to test
    pub max: usize,
    /// Step between sizes (can be multiplication factor)
    pub step: SizeStep,
}

/// Step type for size range
#[derive(Debug, Clone)]
pub enum SizeStep {
    /// Add a constant value
    Linear(usize),
    /// Multiply by a factor
    Exponential(f64),
    /// Use powers of two
    PowersOfTwo,
    /// Use specific sizes
    Custom(Vec<usize>),
}

/// FFT algorithm variant to benchmark
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FftVariant {
    /// Standard FFT
    Standard,
    /// In-place FFT
    InPlace,
    /// Cached-plan FFT
    Cached,
    /// Split-radix FFT
    SplitRadix,
}

/// Configuration for auto-tuning
#[derive(Debug, Clone)]
pub struct AutoTuneConfig {
    /// Sizes to benchmark
    pub sizes: SizeRange,
    /// Number of repetitions per test
    pub repetitions: usize,
    /// Warm-up iterations (not timed)
    pub warmup: usize,
    /// FFT variants to test
    pub variants: Vec<FftVariant>,
    /// Path to save tuning results
    pub database_path: PathBuf,
}

impl Default for AutoTuneConfig {
    fn default() -> Self {
        Self {
            sizes: SizeRange {
                min: 16,
                max: 8192,
                step: SizeStep::PowersOfTwo,
            },
            repetitions: 10,
            warmup: 3,
            variants: vec![FftVariant::Standard, FftVariant::Cached],
            database_path: PathBuf::from(".fft_tuning_db.json"),
        }
    }
}

/// Results from a single benchmark
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
    /// FFT size
    pub size: usize,
    /// FFT variant
    pub variant: FftVariant,
    /// Whether this is forward or inverse FFT
    pub forward: bool,
    /// Average execution time in nanoseconds
    pub avg_time_ns: u64,
    /// Minimum execution time in nanoseconds
    pub min_time_ns: u64,
    /// Standard deviation in nanoseconds
    pub std_dev_ns: f64,
    /// System information when the benchmark was run
    pub system_info: SystemInfo,
}

/// System information for result matching
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemInfo {
    /// CPU model
    pub cpu_model: String,
    /// Number of cores
    pub num_cores: usize,
    /// Architecture
    pub architecture: String,
    /// CPU features (SIMD instruction sets, etc.)
    pub cpu_features: Vec<String>,
}

/// Database of tuning results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TuningDatabase {
    /// Benchmark results
    pub results: Vec<BenchmarkResult>,
    /// Last updated timestamp
    pub last_updated: u64,
    /// Best algorithm for each size
    pub best_algorithms: HashMap<(usize, bool), FftVariant>,
}

/// Auto-tuning manager
pub struct AutoTuner {
    /// Configuration
    config: AutoTuneConfig,
    /// Database of results
    database: TuningDatabase,
    /// Whether to use tuning
    enabled: bool,
}

impl Default for AutoTuner {
    fn default() -> Self {
        Self::with_config(AutoTuneConfig::default())
    }
}

impl AutoTuner {
    /// Create a new auto-tuner with default configuration
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a new auto-tuner with custom configuration
    pub fn with_config(config: AutoTuneConfig) -> Self {
        let database =
            Self::load_database(&config.database_path).unwrap_or_else(|_| TuningDatabase {
                results: Vec::new(),
                last_updated: std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs(),
                best_algorithms: HashMap::new(),
            });

        Self {
            config,
            database,
            enabled: true,
        }
    }

    /// Load the tuning database from disk
    fn load_database(path: &Path) -> FFTResult<TuningDatabase> {
        if !path.exists() {
            return Err(FFTError::IOError(format!(
                "Tuning database file not found: {}",
                path.display()
            )));
        }

        let file = File::open(path)
            .map_err(|e| FFTError::IOError(format!("Failed to open tuning database: {e}")))?;

        let reader = BufReader::new(file);
        let database: TuningDatabase = serde_json::from_reader(reader)
            .map_err(|e| FFTError::ValueError(format!("Failed to parse tuning database: {e}")))?;

        Ok(database)
    }

    /// Save the tuning database to disk
    pub fn save_database(&self) -> FFTResult<()> {
        // Create parent directories if they don't exist
        if let Some(parent) = self.config.database_path.parent() {
            fs::create_dir_all(parent).map_err(|e| {
                FFTError::IOError(format!(
                    "Failed to create directory for tuning database: {e}"
                ))
            })?;
        }

        let file = File::create(&self.config.database_path).map_err(|e| {
            FFTError::IOError(format!("Failed to create tuning database file: {e}"))
        })?;

        let writer = BufWriter::new(file);
        serde_json::to_writer_pretty(writer, &self.database)
            .map_err(|e| FFTError::IOError(format!("Failed to serialize tuning database: {e}")))?;

        Ok(())
    }

    /// Enable or disable auto-tuning
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Check if auto-tuning is enabled
    pub fn is_enabled(&self) -> bool {
        self.enabled
    }

    /// Run benchmarks for all configured FFT variants and sizes
    pub fn run_benchmarks(&mut self) -> FFTResult<()> {
        if !self.enabled {
            return Ok(());
        }

        let sizes = self.generate_sizes();
        let mut results = Vec::new();

        for size in sizes {
            for &variant in &self.config.variants {
                // Benchmark forward transform
                let forward_result = self.benchmark_variant(size, variant, true)?;
                results.push(forward_result);

                // Benchmark inverse transform
                let inverse_result = self.benchmark_variant(size, variant, false)?;
                results.push(inverse_result);
            }
        }

        // Update database
        self.database.results.extend(results);
        self.update_best_algorithms();
        self.save_database()?;

        Ok(())
    }

    /// Generate the list of sizes to benchmark
    fn generate_sizes(&self) -> Vec<usize> {
        let mut sizes = Vec::new();

        match &self.config.sizes.step {
            SizeStep::Linear(step) => {
                let mut size = self.config.sizes.min;
                while size <= self.config.sizes.max {
                    sizes.push(size);
                    size += step;
                }
            }
            SizeStep::Exponential(factor) => {
                let mut size = self.config.sizes.min as f64;
                while size <= self.config.sizes.max as f64 {
                    sizes.push(size as usize);
                    size *= factor;
                }
            }
            SizeStep::PowersOfTwo => {
                let mut size = 1;
                while size < self.config.sizes.min {
                    size *= 2;
                }
                while size <= self.config.sizes.max {
                    sizes.push(size);
                    size *= 2;
                }
            }
            SizeStep::Custom(custom_sizes) => {
                for &size in custom_sizes {
                    if size >= self.config.sizes.min && size <= self.config.sizes.max {
                        sizes.push(size);
                    }
                }
            }
        }

        sizes
    }

    /// Benchmark a specific FFT variant for a given size
    fn benchmark_variant(
        &self,
        size: usize,
        variant: FftVariant,
        forward: bool,
    ) -> FFTResult<BenchmarkResult> {
        // Create test data
        let mut data = vec![Complex64::new(0.0, 0.0); size];
        for (i, val) in data.iter_mut().enumerate().take(size) {
            *val = Complex64::new(i as f64, (i * 2) as f64);
        }

        // Warm-up phase
        for _ in 0..self.config.warmup {
            match variant {
                FftVariant::Standard => {
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    let mut buffer = data.clone();
                    fft.process(&mut buffer);
                }
                FftVariant::InPlace => {
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    // Use in-place processing with scratch buffer
                    let mut buffer = data.clone();
                    let mut scratch = vec![Complex64::new(0.0, 0.0); fft.get_inplace_scratch_len()];
                    fft.process_with_scratch(&mut buffer, &mut scratch);
                }
                FftVariant::Cached => {
                    // Create a plan via the serialization manager
                    let manager = PlanSerializationManager::new(&self.config.database_path);
                    let plan_info = manager.create_plan_info(size, forward);
                    let (_, time) = crate::plan_serialization::create_and_time_plan(size, forward);
                    manager.record_plan_usage(&plan_info, time).unwrap_or(());
                }
                FftVariant::SplitRadix => {
                    // For now, this is just an example variant
                    // In a real implementation, we'd use a specific split-radix algorithm
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    let mut buffer = data.clone();
                    fft.process(&mut buffer);
                }
            }
        }

        // Timing phase
        let mut times = Vec::with_capacity(self.config.repetitions);

        for _ in 0..self.config.repetitions {
            let start = Instant::now();

            match variant {
                FftVariant::Standard => {
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    let mut buffer = data.clone();
                    fft.process(&mut buffer);
                }
                FftVariant::InPlace => {
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    // Use in-place processing with scratch buffer
                    let mut buffer = data.clone();
                    let mut scratch = vec![Complex64::new(0.0, 0.0); fft.get_inplace_scratch_len()];
                    fft.process_with_scratch(&mut buffer, &mut scratch);
                }
                FftVariant::Cached => {
                    // Use the plan cache
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    let mut buffer = data.clone();
                    fft.process(&mut buffer);
                }
                FftVariant::SplitRadix => {
                    // Placeholder for split-radix implementation
                    let mut planner = FftPlanner::new();
                    let fft = if forward {
                        planner.plan_fft_forward(size)
                    } else {
                        planner.plan_fft_inverse(size)
                    };
                    let mut buffer = data.clone();
                    fft.process(&mut buffer);
                }
            }

            let elapsed = start.elapsed();
            times.push(elapsed.as_nanos() as u64);
        }

        // Calculate statistics
        let avg_time = times.iter().sum::<u64>() / times.len() as u64;
        let min_time = *times.iter().min().unwrap_or(&0);

        // Calculate standard deviation
        let variance = times
            .iter()
            .map(|&t| {
                let diff = t as f64 - avg_time as f64;
                diff * diff
            })
            .sum::<f64>()
            / times.len() as f64;
        let std_dev = variance.sqrt();

        Ok(BenchmarkResult {
            size,
            variant,
            forward,
            avg_time_ns: avg_time,
            min_time_ns: min_time,
            std_dev_ns: std_dev,
            system_info: self.detect_system_info(),
        })
    }

    /// Detect system information for result matching
    fn detect_system_info(&self) -> SystemInfo {
        // This is a simplified version - a real implementation would
        // detect actual CPU model, features, etc.
        SystemInfo {
            cpu_model: String::from("Unknown"),
            num_cores: num_cpus::get(),
            architecture: std::env::consts::ARCH.to_string(),
            cpu_features: detect_cpu_features(),
        }
    }

    /// Update the best algorithms based on benchmark results
    fn update_best_algorithms(&mut self) {
        // Clear existing best algorithms
        self.database.best_algorithms.clear();

        // Group results by size and direction
        let mut grouped: HashMap<(usize, bool), Vec<&BenchmarkResult>> = HashMap::new();
        for result in &self.database.results {
            grouped
                .entry((result.size, result.forward))
                .or_default()
                .push(result);
        }

        // Find the best algorithm for each group
        for ((size, forward), results) in grouped {
            if let Some(best) = results.iter().min_by_key(|r| r.avg_time_ns) {
                self.database
                    .best_algorithms
                    .insert((size, forward), best.variant);
            }
        }
    }

    /// Get the best FFT variant for the given size and direction
    pub fn get_best_variant(&self, size: usize, forward: bool) -> FftVariant {
        if !self.enabled {
            return FftVariant::Standard;
        }

        // Look for exact size match
        if let Some(&variant) = self.database.best_algorithms.get(&(size, forward)) {
            return variant;
        }

        // Look for closest size match
        let mut closest_size = 0;
        let mut min_diff = usize::MAX;

        for &(s, f) in self.database.best_algorithms.keys() {
            if f == forward {
                let diff = s.abs_diff(size);
                if diff < min_diff {
                    min_diff = diff;
                    closest_size = s;
                }
            }
        }

        if closest_size > 0 {
            if let Some(&variant) = self.database.best_algorithms.get(&(closest_size, forward)) {
                return variant;
            }
        }

        // Default to standard FFT if no match
        FftVariant::Standard
    }

    /// Run FFT with optimal algorithm selection
    pub fn run_optimal_fft<T>(
        &self,
        input: &[T],
        size: Option<usize>,
        forward: bool,
    ) -> FFTResult<Vec<Complex64>>
    where
        T: Clone + Into<Complex64>,
    {
        let actual_size = size.unwrap_or(input.len());
        let variant = self.get_best_variant(actual_size, forward);

        // Convert input to complex
        let mut buffer: Vec<Complex64> = input.iter().map(|x| x.clone().into()).collect();
        // Pad if necessary
        if buffer.len() < actual_size {
            buffer.resize(actual_size, Complex64::new(0.0, 0.0));
        }

        match variant {
            FftVariant::Standard => {
                let mut planner = FftPlanner::new();
                let fft = if forward {
                    planner.plan_fft_forward(actual_size)
                } else {
                    planner.plan_fft_inverse(actual_size)
                };
                fft.process(&mut buffer);
            }
            FftVariant::InPlace => {
                let mut planner = FftPlanner::new();
                let fft = if forward {
                    planner.plan_fft_forward(actual_size)
                } else {
                    planner.plan_fft_inverse(actual_size)
                };
                let mut scratch = vec![Complex64::new(0.0, 0.0); fft.get_inplace_scratch_len()];
                fft.process_with_scratch(&mut buffer, &mut scratch);
            }
            FftVariant::Cached => {
                // Use the plan cache via PlanSerializationManager
                // Create a plan directly - manager is not needed here
                let (plan_, _) =
                    crate::plan_serialization::create_and_time_plan(actual_size, forward);
                plan_.process(&mut buffer);
            }
            FftVariant::SplitRadix => {
                // Placeholder for split-radix FFT
                let mut planner = FftPlanner::new();
                let fft = if forward {
                    planner.plan_fft_forward(actual_size)
                } else {
                    planner.plan_fft_inverse(actual_size)
                };
                fft.process(&mut buffer);
            }
        }

        // Scale inverse FFT by 1/N if required
        if !forward {
            let scale = 1.0 / (actual_size as f64);
            for val in &mut buffer {
                *val *= scale;
            }
        }

        Ok(buffer)
    }
}

/// Detect CPU features for result matching
#[allow(dead_code)]
fn detect_cpu_features() -> Vec<String> {
    let mut features = Vec::new();

    // Target-specific feature detection
    #[cfg(target_arch = "x86_64")]
    {
        #[cfg(target_feature = "sse")]
        features.push("sse".to_string());

        #[cfg(target_feature = "sse2")]
        features.push("sse2".to_string());

        #[cfg(target_feature = "sse3")]
        features.push("sse3".to_string());

        #[cfg(target_feature = "sse4.1")]
        features.push("sse4.1".to_string());

        #[cfg(target_feature = "sse4.2")]
        features.push("sse4.2".to_string());

        #[cfg(target_feature = "avx")]
        features.push("avx".to_string());

        #[cfg(target_feature = "avx2")]
        features.push("avx2".to_string());

        #[cfg(target_feature = "fma")]
        features.push("fma".to_string());
    }

    // ARM-specific features
    #[cfg(target_arch = "aarch64")]
    {
        #[cfg(target_feature = "neon")]
        features.push("neon".to_string());
    }

    // Add more architecture-specific features if needed

    features
}

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

    #[test]
    fn test_size_generation() {
        // Test powers of two
        let config = AutoTuneConfig {
            sizes: SizeRange {
                min: 8,
                max: 64,
                step: SizeStep::PowersOfTwo,
            },
            ..Default::default()
        };
        let tuner = AutoTuner::with_config(config);
        let sizes = tuner.generate_sizes();
        assert_eq!(sizes, vec![8, 16, 32, 64]);

        // Test linear steps
        let config = AutoTuneConfig {
            sizes: SizeRange {
                min: 10,
                max: 30,
                step: SizeStep::Linear(5),
            },
            ..Default::default()
        };
        let tuner = AutoTuner::with_config(config);
        let sizes = tuner.generate_sizes();
        assert_eq!(sizes, vec![10, 15, 20, 25, 30]);

        // Test exponential steps
        let config = AutoTuneConfig {
            sizes: SizeRange {
                min: 10,
                max: 100,
                step: SizeStep::Exponential(2.0),
            },
            ..Default::default()
        };
        let tuner = AutoTuner::with_config(config);
        let sizes = tuner.generate_sizes();
        assert_eq!(sizes, vec![10, 20, 40, 80]);

        // Test custom sizes
        let config = AutoTuneConfig {
            sizes: SizeRange {
                min: 10,
                max: 100,
                step: SizeStep::Custom(vec![5, 15, 25, 50, 150]),
            },
            ..Default::default()
        };
        let tuner = AutoTuner::with_config(config);
        let sizes = tuner.generate_sizes();
        assert_eq!(sizes, vec![15, 25, 50]);
    }

    #[test]
    fn test_auto_tuner_basic() {
        // Create a temporary directory for test
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_tuning_db.json");

        // Create configuration with minimal benchmarking
        let config = AutoTuneConfig {
            sizes: SizeRange {
                min: 16,
                max: 32,
                step: SizeStep::PowersOfTwo,
            },
            repetitions: 2,
            warmup: 1,
            variants: vec![FftVariant::Standard, FftVariant::InPlace],
            database_path: db_path.clone(),
        };

        let mut tuner = AutoTuner::with_config(config);

        // Run minimal benchmarks (this is fast enough for a test)
        match tuner.run_benchmarks() {
            Ok(_) => {
                // Verify database file was created
                assert!(db_path.exists());

                // Test getting a best variant
                let variant = tuner.get_best_variant(16, true);
                assert!(matches!(
                    variant,
                    FftVariant::Standard | FftVariant::InPlace
                ));
            }
            Err(e) => {
                // Benchmark may fail in some environments, just log and continue
                println!("Benchmark failed: {e}");
            }
        }
    }
}