scirs2_core/

jit.rs

//! Just-In-Time (JIT) Compilation Framework for Dynamic Kernel Generation
//!
//! This module provides a comprehensive JIT compilation system for generating optimized
//! kernels at runtime. It supports multiple backends including LLVM IR generation,
//! GPU kernel compilation, and adaptive optimization based on runtime characteristics.
//!
//! Features:
//! - LLVM-based code generation for CPU and GPU
//! - Runtime optimization and specialization
//! - Adaptive compilation based on execution patterns
//! - Multi-backend support (CUDA, OpenCL, CPU)
//! - Kernel caching and reuse
//! - Performance profiling and auto-tuning

use crate::error::{CoreError, ErrorContext, ErrorLocation};
#[cfg(feature = "gpu")]
#[allow(unused_imports)]
use crate::gpu::{GpuBackend, GpuContext, GpuError};
use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};
use thiserror::Error;

#[cfg(feature = "parallel")]
#[allow(unused_imports)]
use crate::parallel_ops::*;

/// JIT compilation error types
#[derive(Error, Debug)]
pub enum JitError {
    /// Compilation failed
    #[error("JIT compilation failed: {0}")]
    CompilationError(String),

    /// Code generation error
    #[error("Code generation error: {0}")]
    CodeGenerationError(String),

    /// Optimization error
    #[error("Optimization error: {0}")]
    OptimizationError(String),

    /// Backend not supported
    #[error("Backend not supported: {backend}")]
    BackendNotSupported { backend: String },

    /// Invalid kernel source
    #[error("Invalid kernel source: {0}")]
    InvalidKernelSource(String),

    /// Runtime execution error
    #[error("Runtime execution error: {0}")]
    RuntimeError(String),

    /// Cache error
    #[error("Kernel cache error: {0}")]
    CacheError(String),

    /// Profiling error
    #[error("Profiling error: {0}")]
    ProfilingError(String),

    /// Underlying GPU error
    #[cfg(feature = "gpu")]
    #[error("GPU error: {0}")]
    GpuError(#[from] GpuError),
}

impl From<JitError> for CoreError {
    fn from(err: JitError) -> Self {
        match err {
            JitError::CompilationError(msg) => CoreError::ComputationError(
                ErrorContext::new(format!("{msg}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
            JitError::CodeGenerationError(msg) => CoreError::ComputationError(
                ErrorContext::new(format!("{msg}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
            JitError::OptimizationError(msg) => CoreError::ComputationError(
                ErrorContext::new(format!("{msg}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
            JitError::BackendNotSupported { backend } => CoreError::NotImplementedError(
                ErrorContext::new(format!("{backend}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
            JitError::RuntimeError(msg) => CoreError::ComputationError(
                ErrorContext::new(format!("{msg}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
            _ => CoreError::ComputationError(
                ErrorContext::new(format!("{err}"))
                    .with_location(ErrorLocation::new(file!(), line!())),
            ),
        }
    }
}

/// JIT compilation backends
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum JitBackend {
    /// LLVM-based compilation
    Llvm,
    /// GPU-specific backends
    Cuda,
    OpenCl,
    Metal,
    WebGpu,
    /// Interpreter-based execution
    Interpreter,
    /// Native code generation
    NativeCode,
    /// Custom backend
    Custom(&'static str),
}

impl fmt::Display for JitBackend {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            JitBackend::Llvm => write!(f, "LLVM"),
            JitBackend::Cuda => write!(f, "CUDA"),
            JitBackend::OpenCl => write!(f, "OpenCL"),
            JitBackend::Metal => write!(f, "Metal"),
            JitBackend::WebGpu => write!(f, "WebGPU"),
            JitBackend::Interpreter => write!(f, "Interpreter"),
            JitBackend::NativeCode => write!(f, "NativeCode"),
            JitBackend::Custom(name) => write!(f, "Custom({})", name),
        }
    }
}

/// JIT compilation target architectures
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TargetArchitecture {
    /// x86-64 CPU
    X86_64,
    /// ARM64 CPU
    Arm64,
    /// NVIDIA GPU (CUDA)
    NvidiaGpu,
    /// AMD GPU (ROCm)
    AmdGpu,
    /// Intel GPU
    IntelGpu,
    /// Apple GPU (Metal)
    AppleGpu,
    /// WebGPU
    WebGpu,
}

/// Optimization levels for JIT compilation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OptimizationLevel {
    /// No optimization
    None,
    /// Basic optimizations
    O1,
    /// Standard optimizations
    O2,
    /// Aggressive optimizations
    O3,
    /// Size optimizations
    Os,
    /// Fast math optimizations
    Ofast,
    /// Adaptive optimization based on profiling
    Adaptive,
}

/// JIT compilation configuration
#[derive(Debug, Clone)]
pub struct JitConfig {
    /// Target backend
    pub backend: JitBackend,
    /// Target architecture
    pub target_arch: TargetArchitecture,
    /// Optimization level
    pub optimization_level: OptimizationLevel,
    /// Enable caching
    pub enable_caching: bool,
    /// Enable profiling
    pub enable_profiling: bool,
    /// Maximum cache size
    pub max_cache_size: usize,
    /// Compilation timeout
    pub compilation_timeout: Duration,
    /// Enable adaptive optimization
    pub adaptive_optimization: bool,
    /// Custom compilation flags
    pub custom_flags: Vec<String>,
}

impl Default for JitConfig {
    fn default() -> Self {
        Self {
            backend: JitBackend::Llvm,
            target_arch: TargetArchitecture::X86_64,
            optimization_level: OptimizationLevel::O2,
            enable_caching: true,
            enable_profiling: true,
            max_cache_size: 256 * 1024 * 1024, // 256MB
            compilation_timeout: Duration::from_secs(30),
            adaptive_optimization: true,
            custom_flags: Vec::new(),
        }
    }
}

/// Kernel source code abstraction
#[derive(Debug, Clone)]
pub struct KernelSource {
    /// Unique identifier for the kernel
    pub id: String,
    /// Source code
    pub source: String,
    /// Kernel language/dialect
    pub language: KernelLanguage,
    /// Entry point function name
    pub entry_point: String,
    /// Input parameter types
    pub input_types: Vec<DataType>,
    /// Output parameter types
    pub output_types: Vec<DataType>,
    /// Compilation hints
    pub hints: CompilationHints,
}

/// Kernel programming languages/dialects
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KernelLanguage {
    /// LLVM IR
    LlvmIr,
    /// CUDA C/C++
    Cuda,
    /// OpenCL C
    OpenCl,
    /// HLSL (DirectX)
    Hlsl,
    /// Metal Shading Language
    Metal,
    /// WGSL (WebGPU)
    Wgsl,
    /// High-level DSL
    HighLevel,
    /// Assembly language
    Assembly,
}

/// Data types for kernel parameters
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DataType {
    /// 8-bit signed integer
    I8,
    /// 16-bit signed integer
    I16,
    /// 32-bit signed integer
    I32,
    /// 64-bit signed integer
    I64,
    /// 8-bit unsigned integer
    U8,
    /// 16-bit unsigned integer
    U16,
    /// 32-bit unsigned integer
    U32,
    /// 64-bit unsigned integer
    U64,
    /// 16-bit floating point
    F16,
    /// 32-bit floating point
    F32,
    /// 64-bit floating point
    F64,
    /// Boolean
    Bool,
    /// Pointer to memory
    Ptr(Box<DataType>),
    /// Array of fixed size
    Array(Box<DataType>, usize),
    /// Vector types
    Vec2(Box<DataType>),
    Vec3(Box<DataType>),
    Vec4(Box<DataType>),
}

/// Compilation hints for optimization
#[derive(Debug, Clone, Default)]
pub struct CompilationHints {
    /// Expected workload size
    pub workload_size: Option<usize>,
    /// Memory access pattern
    pub memory_pattern: Option<MemoryPattern>,
    /// Computational intensity
    pub compute_intensity: Option<ComputeIntensity>,
    /// Parallelization hints
    pub parallelization: Option<ParallelizationHints>,
    /// Target-specific hints
    pub target_hints: HashMap<String, String>,
}

/// Memory access patterns
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryPattern {
    /// Sequential access
    Sequential,
    /// Random access
    Random,
    /// Strided access
    Strided,
    /// Coalesced access
    Coalesced,
    /// Scattered access
    Scattered,
}

/// Computational intensity levels
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ComputeIntensity {
    /// Memory-bound operations
    MemoryBound,
    /// Compute-bound operations
    ComputeBound,
    /// Balanced compute and memory
    Balanced,
    /// Bandwidth-intensive
    BandwidthIntensive,
}

impl Default for ComputeIntensity {
    fn default() -> Self {
        ComputeIntensity::Balanced
    }
}

/// Parallelization hints
#[derive(Debug, Clone)]
pub struct ParallelizationHints {
    /// Preferred work group size
    pub work_group_size: Option<[usize; 3]>,
    /// Vectorization width
    pub vector_width: Option<usize>,
    /// Loop unrolling factor
    pub unroll_factor: Option<usize>,
    /// Enable auto-vectorization
    pub auto_vectorize: bool,
}

impl Default for ParallelizationHints {
    fn default() -> Self {
        Self {
            work_group_size: None,
            vector_width: None,
            unroll_factor: None,
            auto_vectorize: true,
        }
    }
}

/// Compiled kernel representation
#[derive(Debug, Clone)]
pub struct CompiledKernel {
    /// Kernel identifier
    pub id: String,
    /// Compiled binary/bytecode
    pub binary: Vec<u8>,
    /// Backend used for compilation
    pub backend: JitBackend,
    /// Target architecture
    pub target_arch: TargetArchitecture,
    /// Compilation metadata
    pub metadata: KernelMetadata,
    /// Performance characteristics
    pub performance: KernelPerformance,
}

/// Kernel compilation metadata
#[derive(Debug, Clone)]
pub struct KernelMetadata {
    /// Compilation timestamp
    pub compiled_at: Instant,
    /// Compilation time
    pub compilation_time: Duration,
    /// Optimization level used
    pub optimization_level: OptimizationLevel,
    /// Binary size
    pub binary_size: usize,
    /// Register usage (GPU kernels)
    pub register_usage: Option<usize>,
    /// Shared memory usage (GPU kernels)
    pub shared_memory_usage: Option<usize>,
    /// Compiler version/info
    pub compiler_info: String,
}

/// Kernel performance characteristics
#[derive(Debug, Clone, Default)]
pub struct KernelPerformance {
    /// Execution count
    pub execution_count: usize,
    /// Total execution time
    pub totalexecution_time: Duration,
    /// Average execution time
    pub avgexecution_time: Duration,
    /// Best execution time
    pub bestexecution_time: Duration,
    /// Worst execution time
    pub worstexecution_time: Duration,
    /// Throughput (operations per second)
    pub throughput: f64,
    /// Energy efficiency (operations per joule)
    pub energy_efficiency: Option<f64>,
}

/// JIT compiler interface
pub struct JitCompiler {
    /// Configuration
    config: JitConfig,
    /// Backend implementations
    backends: HashMap<JitBackend, Box<dyn JitBackendImpl>>,
    /// Kernel cache
    cache: Arc<RwLock<KernelCache>>,
    /// Performance profiler
    profiler: Arc<Mutex<KernelProfiler>>,
    /// Adaptive optimizer
    adaptive_optimizer: Arc<Mutex<AdaptiveOptimizer>>,
}

/// Kernel cache for compiled kernels
#[derive(Debug)]
pub struct KernelCache {
    /// Cached kernels
    kernels: HashMap<String, CompiledKernel>,
    /// Cache size in bytes
    current_size: usize,
    /// Maximum cache size
    maxsize: usize,
    /// Access frequency tracking
    access_counts: HashMap<String, usize>,
    /// Last access times
    last_accessed: HashMap<String, Instant>,
}

/// Kernel performance profiler
#[derive(Debug)]
pub struct KernelProfiler {
    /// Execution profiles
    profiles: HashMap<String, Vec<ExecutionProfile>>,
    /// Hardware performance counters
    hw_counters: HardwareCounters,
    /// Profiling enabled
    enabled: bool,
}

/// Individual execution profile
#[derive(Debug, Clone)]
pub struct ExecutionProfile {
    /// Execution timestamp
    pub timestamp: Instant,
    /// Execution time
    pub execution_time: Duration,
    /// Memory bandwidth utilized
    pub memorybandwidth: f64,
    /// Compute utilization
    pub compute_utilization: f64,
    /// Cache hit rates
    pub cache_hit_rates: Vec<f64>,
    /// Power consumption
    pub power_consumption: Option<f64>,
}

/// Hardware performance counters
#[derive(Debug, Default)]
pub struct HardwareCounters {
    /// CPU cycles
    pub cpu_cycles: u64,
    /// Instructions executed
    pub instructions: u64,
    /// Cache misses
    pub cache_misses: u64,
    /// Memory transactions
    pub memory_transactions: u64,
    /// GPU-specific counters
    pub gpu_counters: HashMap<String, u64>,
}

/// Adaptive optimizer for runtime optimization
#[derive(Debug)]
pub struct AdaptiveOptimizer {
    /// Optimization history
    optimization_history: HashMap<String, Vec<OptimizationResult>>,
    /// Learning model for optimization decisions
    learning_model: Option<Box<dyn OptimizationModel>>,
    /// Optimization strategies
    strategies: Vec<OptimizationStrategy>,
}

/// Optimization result tracking
#[derive(Debug, Clone)]
pub struct OptimizationResult {
    /// Strategy used
    pub strategy: OptimizationStrategy,
    /// Performance improvement
    pub improvement: f64,
    /// Compilation overhead
    pub compilation_overhead: Duration,
    /// Success flag
    pub success: bool,
}

/// Optimization strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OptimizationStrategy {
    /// Loop unrolling
    LoopUnrolling,
    /// Vectorization
    Vectorization,
    /// Memory prefetching
    MemoryPrefetching,
    /// Register allocation optimization
    RegisterAllocation,
    /// Instruction scheduling
    InstructionScheduling,
    /// Constant folding
    ConstantFolding,
    /// Dead code elimination
    DeadCodeElimination,
    /// Function inlining
    FunctionInlining,
}

/// Machine learning model for optimization decisions
pub trait OptimizationModel: Send + Sync + std::fmt::Debug {
    /// Predict optimal strategy for a kernel
    fn predict(&self, features: &KernelFeatures) -> OptimizationStrategy;

    /// Update model with feedback
    fn update_model(&mut self, features: &KernelFeatures, result: &OptimizationResult);
}

/// Kernel feature extraction for ML optimization
#[derive(Debug, Clone)]
pub struct KernelFeatures {
    /// Source code metrics
    pub source_metrics: SourceMetrics,
    /// Runtime characteristics
    pub runtime_metrics: RuntimeMetrics,
    /// Target characteristics
    pub target_metrics: TargetMetrics,
}

/// Source code metrics
#[derive(Debug, Clone, Default)]
pub struct SourceMetrics {
    /// Lines of code
    pub lines_ofcode: usize,
    /// Loop count
    pub loop_count: usize,
    /// Branching factor
    pub branching_factor: f64,
    /// Memory operations count
    pub memory_ops_count: usize,
    /// Arithmetic operations count
    pub arithmetic_ops_count: usize,
    /// Function call count
    pub function_call_count: usize,
}

/// Runtime characteristics
#[derive(Debug, Clone, Default)]
pub struct RuntimeMetrics {
    /// Typical input sizes
    pub typical_input_sizes: Vec<usize>,
    /// Execution frequency
    pub execution_frequency: f64,
    /// Memory access patterns
    pub memory_patterns: Vec<MemoryPattern>,
    /// Computational intensity
    pub compute_intensity: ComputeIntensity,
}

/// Target platform metrics
#[derive(Debug, Clone, Default)]
pub struct TargetMetrics {
    /// Available compute units
    pub compute_units: usize,
    /// Memory bandwidth
    pub memorybandwidth: f64,
    /// Cache sizes
    pub cache_sizes: Vec<usize>,
    /// Vector width
    pub vector_width: usize,
}

/// JIT backend implementation trait
pub trait JitBackendImpl: Send + Sync {
    /// Compile kernel source to binary
    fn compile_kernel(
        &self,
        source: &KernelSource,
        config: &JitConfig,
    ) -> Result<CompiledKernel, JitError>;

    /// Execute compiled kernel
    fn execute_kernel(
        &self,
        kernel: &CompiledKernel,
        inputs: &[&dyn std::any::Any],
        outputs: &mut [&mut dyn std::any::Any],
    ) -> Result<ExecutionProfile, JitError>;

    /// Check if backend is available
    fn is_available(&self) -> bool;

    /// Get backend capabilities
    fn get_capabilities(&self) -> BackendCapabilities;
}

/// Backend capabilities
#[derive(Debug, Clone)]
pub struct BackendCapabilities {
    /// Supported data types
    pub supported_types: Vec<DataType>,
    /// Supported optimization levels
    pub optimization_levels: Vec<OptimizationLevel>,
    /// Maximum kernel size
    pub max_kernel_size: Option<usize>,
    /// Supports debugging
    pub supports_debugging: bool,
    /// Supports profiling
    pub supports_profiling: bool,
    /// Target architectures
    pub target_architectures: Vec<TargetArchitecture>,
}

impl JitCompiler {
    /// Create a new JIT compiler
    pub fn new(config: JitConfig) -> Result<Self, JitError> {
        let mut backends = HashMap::new();

        // Initialize available backends
        if config.backend == JitBackend::Llvm || config.backend == JitBackend::NativeCode {
            backends.insert(
                JitBackend::Llvm,
                Box::new(LlvmBackend::new()?) as Box<dyn JitBackendImpl>,
            );
        }

        backends.insert(
            JitBackend::Interpreter,
            Box::new(InterpreterBackend::new()) as Box<dyn JitBackendImpl>,
        );

        let cache = Arc::new(RwLock::new(KernelCache::size(config.max_cache_size)));
        let profiler = Arc::new(Mutex::new(KernelProfiler::new(config.enable_profiling)));
        let adaptive_optimizer = Arc::new(Mutex::new(AdaptiveOptimizer::new()));

        Ok(Self {
            config,
            backends,
            cache,
            profiler,
            adaptive_optimizer,
        })
    }

    /// Compile a kernel from source
    pub fn compile_kernel(&self, source: KernelSource) -> Result<String, JitError> {
        let kernel_id = source.id.clone();

        // Check cache first
        if self.config.enable_caching {
            let cache = self.cache.read().expect("Operation failed");
            if cache.contains_kernel(&kernel_id) {
                return Ok(kernel_id);
            }
        }

        // Get backend
        let backend = self.backends.get(&self.config.backend).ok_or_else(|| {
            JitError::BackendNotSupported {
                backend: format!("{:?}", self.config.backend),
            }
        })?;

        // Compile kernel
        let compiled_kernel = backend.compile_kernel(&source, &self.config)?;

        // Cache compiled kernel
        if self.config.enable_caching {
            let mut cache = self.cache.write().expect("Operation failed");
            cache.insert(compiled_kernel);
        }

        Ok(kernel_id)
    }

    /// Execute a compiled kernel
    pub fn execute_kernel(
        &self,
        kernel_id: &str,
        inputs: &[&dyn std::any::Any],
        outputs: &mut [&mut dyn std::any::Any],
    ) -> Result<(), JitError> {
        // Get compiled kernel from cache
        let kernel = {
            let cache = self.cache.read().expect("Operation failed");
            cache
                .get_readonly(kernel_id)
                .ok_or_else(|| JitError::CacheError(format!("{kernel_id}")))?
                .clone()
        };

        // Get backend
        let backend =
            self.backends
                .get(&kernel.backend)
                .ok_or_else(|| JitError::BackendNotSupported {
                    backend: format!("{:?}", kernel.backend),
                })?;

        // Execute kernel
        let profile = backend.execute_kernel(&kernel, inputs, outputs)?;

        // Update profiling data
        if self.config.enable_profiling {
            let mut profiler = self.profiler.lock().expect("Operation failed");
            profiler.record_execution(kernel_id, profile);
        }

        // Update adaptive optimization
        if self.config.adaptive_optimization {
            let mut optimizer = self.adaptive_optimizer.lock().expect("Operation failed");
            optimizer.update_performance_data(&kernel.performance);
        }

        Ok(())
    }

    /// Get kernel performance statistics
    pub fn get_kernel_performance(&self, kernel_id: &str) -> Option<KernelPerformance> {
        let mut cache = self.cache.write().expect("Operation failed");
        cache.get(kernel_id).map(|k| k.performance.clone())
    }

    /// Get compilation statistics
    pub fn get_compilation_stats(&self) -> CompilationStats {
        let cache = self.cache.read().expect("Operation failed");
        cache.get_stats()
    }

    /// Clear kernel cache
    pub fn clear_cache(&self) {
        let mut cache = self.cache.write().expect("Operation failed");
        cache.clear();
    }

    /// Optimize existing kernel
    pub fn optimize_kernel(&self, kernel_id: &str) -> Result<String, JitError> {
        let optimizer = self.adaptive_optimizer.lock().expect("Operation failed");
        optimizer.optimize_kernel(kernel_id, &self.config)
    }
}

/// Compilation statistics
#[derive(Debug, Clone, Default)]
pub struct CompilationStats {
    /// Total kernels compiled
    pub total_compiled: usize,
    /// Cache hit rate
    pub cache_hit_rate: f64,
    /// Average compilation time
    pub avg_compilation_time: Duration,
    /// Total cache size
    pub cache_size: usize,
    /// Most frequently used kernels
    pub top_kernels: Vec<(String, usize)>,
}

impl KernelCache {
    /// Create a new kernel cache
    pub fn size(value: usize) -> Self {
        Self {
            kernels: HashMap::new(),
            current_size: 0,
            maxsize: value,
            access_counts: HashMap::new(),
            last_accessed: HashMap::new(),
        }
    }

    /// Check if kernel is cached
    pub fn contains_kernel(&self, kernel_id: &str) -> bool {
        self.kernels.contains_key(kernel_id)
    }

    /// Get kernel from cache
    pub fn get(&mut self, kernel_id: &str) -> Option<&CompiledKernel> {
        if let Some(kernel) = self.kernels.get(kernel_id) {
            // Update access tracking
            *self.access_counts.entry(kernel_id.to_string()).or_insert(0) += 1;
            self.last_accessed
                .insert(kernel_id.to_string(), Instant::now());
            Some(kernel)
        } else {
            None
        }
    }

    /// Get a kernel from the cache without updating access tracking
    pub fn get_readonly(&self, kernel_id: &str) -> Option<&CompiledKernel> {
        self.kernels.get(kernel_id)
    }

    /// Insert kernel into cache
    pub fn insert(&mut self, kernel: CompiledKernel) {
        let kernel_id = kernel.id.clone();
        let kernel_size = kernel.binary.len();

        // Check if we need to evict
        while self.current_size + kernel_size > self.maxsize && !self.kernels.is_empty() {
            self.evict_lru();
        }

        self.current_size += kernel_size;
        self.kernels.insert(kernel_id.clone(), kernel);
        self.access_counts.insert(kernel_id.clone(), 1);
        self.last_accessed.insert(kernel_id, Instant::now());
    }

    /// Evict least recently used kernel
    fn evict_lru(&mut self) {
        if let Some((lru_id, _)) = self.last_accessed.iter().min_by_key(|(_, &time)| time) {
            let lru_id = lru_id.clone();
            if let Some(kernel) = self.kernels.remove(&lru_id) {
                self.current_size -= kernel.binary.len();
                self.access_counts.remove(&lru_id);
                self.last_accessed.remove(&lru_id);
            }
        }
    }

    /// Clear all cached kernels
    pub fn clear(&mut self) {
        self.kernels.clear();
        self.access_counts.clear();
        self.last_accessed.clear();
        self.current_size = 0;
    }

    /// Get cache statistics
    pub fn get_stats(&self) -> CompilationStats {
        let total_accesses: usize = self.access_counts.values().sum();
        let cache_hit_rate = if total_accesses > 0 {
            self.access_counts.len() as f64 / total_accesses as f64
        } else {
            0.0
        };

        let mut top_kernels: Vec<_> = self
            .access_counts
            .iter()
            .map(|(id, count)| (id.clone(), *count))
            .collect();
        top_kernels.sort_by_key(|b| std::cmp::Reverse(b.1));
        top_kernels.truncate(10);

        CompilationStats {
            total_compiled: self.kernels.len(),
            cache_hit_rate,
            avg_compilation_time: Duration::from_millis(100), // Placeholder
            cache_size: self.current_size,
            top_kernels,
        }
    }
}

impl KernelProfiler {
    /// Create a new profiler
    pub fn new(enabled: bool) -> Self {
        Self {
            profiles: HashMap::new(),
            hw_counters: HardwareCounters::default(),
            enabled,
        }
    }

    /// Record kernel execution
    pub fn record_execution(&mut self, kernel_id: &str, profile: ExecutionProfile) {
        if !self.enabled {
            return;
        }

        self.profiles
            .entry(kernel_id.to_string())
            .or_insert_with(Vec::new)
            .push(profile);
    }

    /// Get profiling data for a kernel
    pub fn id_2(&self, kernelid: &str) -> Option<&Vec<ExecutionProfile>> {
        self.profiles.get(kernelid)
    }
}

impl AdaptiveOptimizer {
    /// Create a new adaptive optimizer
    pub fn new() -> Self {
        Self {
            optimization_history: HashMap::new(),
            learning_model: None,
            strategies: vec![
                OptimizationStrategy::LoopUnrolling,
                OptimizationStrategy::Vectorization,
                OptimizationStrategy::MemoryPrefetching,
                OptimizationStrategy::RegisterAllocation,
            ],
        }
    }

    /// Update performance data by recording an `OptimizationResult` synthesised
    /// from the observed `KernelPerformance`.  The result is stored under the
    /// special key `"__perf_trends__"` so that `optimize_kernel` can later
    /// inspect aggregate patterns when no per-kernel history is available.
    ///
    /// Strategy selection heuristic:
    /// - High throughput (> 1 GOp/s) → prefer `Vectorization`
    /// - Low execution count (< 10 executions) → prefer `ConstantFolding`
    /// - Otherwise → cycle through the available strategies in insertion order
    pub fn update_performance_data(&mut self, data: &KernelPerformance) {
        // Derive a rough "improvement" signal: normalised throughput (clamped 0..1).
        // A throughput of 1e9 ops/s is treated as "perfect" (improvement = 1.0).
        let improvement = (data.throughput / 1.0e9).clamp(0.0, 1.0);

        // Pick a strategy based on the current performance snapshot.
        let strategy = if data.throughput > 1.0e9 {
            OptimizationStrategy::Vectorization
        } else if data.execution_count < 10 {
            OptimizationStrategy::ConstantFolding
        } else {
            // Round-robin over available strategies.
            let existing = self
                .optimization_history
                .get("__perf_trends__")
                .map(|v| v.len())
                .unwrap_or(0);
            let idx = existing % self.strategies.len();
            self.strategies[idx]
        };

        let result = OptimizationResult {
            strategy,
            improvement,
            compilation_overhead: data.avgexecution_time,
            success: improvement > 0.1,
        };

        self.optimization_history
            .entry("__perf_trends__".to_string())
            .or_default()
            .push(result);
    }

    /// Suggest an optimization directive for the kernel identified by `kernel_id`.
    ///
    /// If the kernel has dedicated history (populated by callers that use
    /// `kernel_id` as the history key), the most successful `OptimizationResult`
    /// is selected.  Otherwise the method falls back to the aggregate performance
    /// trends recorded by `update_performance_data`.  The returned `String` is a
    /// human-readable directive (e.g. `"vectorize"`, `"unroll"`) that the JIT
    /// back-end may interpret as a compilation hint.
    pub fn optimize_kernel(&self, kernel_id: &str, config: &JitConfig) -> Result<String, JitError> {
        // Try per-kernel history first.
        let history = self
            .optimization_history
            .get(kernel_id)
            .or_else(|| self.optimization_history.get("__perf_trends__"));

        if let Some(records) = history {
            // Find the record with the highest improvement that succeeded.
            let best = records.iter().filter(|r| r.success).max_by(|a, b| {
                a.improvement
                    .partial_cmp(&b.improvement)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });

            if let Some(best_result) = best {
                let directive = match best_result.strategy {
                    OptimizationStrategy::LoopUnrolling => "unroll",
                    OptimizationStrategy::Vectorization => "vectorize",
                    OptimizationStrategy::MemoryPrefetching => "prefetch",
                    OptimizationStrategy::RegisterAllocation => "regalloc",
                    OptimizationStrategy::InstructionScheduling => "schedule",
                    OptimizationStrategy::ConstantFolding => "constfold",
                    OptimizationStrategy::DeadCodeElimination => "dce",
                    OptimizationStrategy::FunctionInlining => "inline",
                };
                let level_flag = optimization_level_flag(config.optimization_level);
                return Ok(format!("{directive} {level_flag}"));
            }
        }

        // No history yet — derive a default directive from the config.
        let default_directive = match config.optimization_level {
            OptimizationLevel::None => "none",
            OptimizationLevel::O1 => "constfold",
            OptimizationLevel::O2 => "vectorize",
            OptimizationLevel::O3 => "unroll vectorize prefetch",
            OptimizationLevel::Os => "constfold dce",
            OptimizationLevel::Ofast => "unroll vectorize prefetch inline",
            OptimizationLevel::Adaptive => "vectorize",
        };
        let level_flag = optimization_level_flag(config.optimization_level);
        Ok(format!("{default_directive} {level_flag}"))
    }
}

/// Return a short flag string that names the optimisation level, e.g. `"-O2"`.
fn optimization_level_flag(level: OptimizationLevel) -> &'static str {
    match level {
        OptimizationLevel::None => "-O0",
        OptimizationLevel::O1 => "-O1",
        OptimizationLevel::O2 => "-O2",
        OptimizationLevel::O3 => "-O3",
        OptimizationLevel::Os => "-Os",
        OptimizationLevel::Ofast => "-Ofast",
        OptimizationLevel::Adaptive => "-O2",
    }
}

/// LLVM-based backend implementation
pub struct LlvmBackend {
    /// LLVM context
    context: Option<()>, // Placeholder for LLVM context
}

impl LlvmBackend {
    /// Create new LLVM backend
    pub fn new() -> Result<Self, JitError> {
        // In a real implementation, this would initialize LLVM
        Ok(Self { context: Some(()) })
    }
}

impl JitBackendImpl for LlvmBackend {
    fn compile_kernel(
        &self,
        source: &KernelSource,
        config: &JitConfig,
    ) -> Result<CompiledKernel, JitError> {
        // Placeholder implementation
        let compilation_start = Instant::now();

        // In a real implementation, this would:
        // 1. Parse the source code
        // 2. Generate LLVM IR
        // 3. Apply optimizations
        // 4. Generate machine code

        let compilation_time = compilation_start.elapsed();

        Ok(CompiledKernel {
            id: source.id.clone(),
            binary: vec![0; 1024], // Placeholder binary
            backend: config.backend,
            target_arch: config.target_arch,
            metadata: KernelMetadata {
                compiled_at: Instant::now(),
                compilation_time,
                optimization_level: config.optimization_level,
                binary_size: 1024,
                register_usage: Some(32),
                shared_memory_usage: Some(1024),
                compiler_info: "LLVM 15.0".to_string(),
            },
            performance: KernelPerformance::default(),
        })
    }

    fn execute_kernel(
        &self,
        kernel: &CompiledKernel,
        inputs: &[&dyn std::any::Any],
        outputs: &mut [&mut dyn std::any::Any],
    ) -> Result<ExecutionProfile, JitError> {
        // Placeholder implementation
        let start = Instant::now();

        // Simulate execution
        std::thread::sleep(Duration::from_micros(100));

        Ok(ExecutionProfile {
            timestamp: start,
            execution_time: start.elapsed(),
            memorybandwidth: 100.0, // GB/s
            compute_utilization: 0.8,
            cache_hit_rates: vec![0.95, 0.87, 0.72],
            power_consumption: Some(50.0), // Watts
        })
    }

    fn is_available(&self) -> bool {
        self.context.is_some()
    }

    fn get_capabilities(&self) -> BackendCapabilities {
        BackendCapabilities {
            supported_types: vec![
                DataType::I32,
                DataType::I64,
                DataType::F32,
                DataType::F64,
                DataType::Vec4(Box::new(DataType::F32)),
            ],
            optimization_levels: vec![
                OptimizationLevel::None,
                OptimizationLevel::O1,
                OptimizationLevel::O2,
                OptimizationLevel::O3,
            ],
            max_kernel_size: None,
            supports_debugging: true,
            supports_profiling: true,
            target_architectures: vec![TargetArchitecture::X86_64, TargetArchitecture::Arm64],
        }
    }
}

/// Interpreter-based backend for debugging and fallback
pub struct InterpreterBackend;

impl InterpreterBackend {
    /// Create new interpreter backend
    pub fn new() -> Self {
        Self
    }
}

impl JitBackendImpl for InterpreterBackend {
    fn compile_kernel(
        &self,
        source: &KernelSource,
        config: &JitConfig,
    ) -> Result<CompiledKernel, JitError> {
        // For interpreter, "compilation" is just validation
        let compilation_start = Instant::now();

        // Basic validation
        if source.source.is_empty() {
            return Err(JitError::InvalidKernelSource("Empty source".to_string()));
        }

        let compilation_time = compilation_start.elapsed();

        Ok(CompiledKernel {
            id: source.id.clone(),
            binary: source.source.as_bytes().to_vec(),
            backend: config.backend,
            target_arch: config.target_arch,
            metadata: KernelMetadata {
                compiled_at: Instant::now(),
                compilation_time,
                optimization_level: OptimizationLevel::None,
                binary_size: source.source.len(),
                register_usage: None,
                shared_memory_usage: None,
                compiler_info: JitBackend::Interpreter.to_string(),
            },
            performance: KernelPerformance::default(),
        })
    }

    fn execute_kernel(
        &self,
        kernel: &CompiledKernel,
        inputs: &[&dyn std::any::Any],
        outputs: &mut [&mut dyn std::any::Any],
    ) -> Result<ExecutionProfile, JitError> {
        // Placeholder interpreter execution
        let start = Instant::now();

        // Simulate interpretation
        std::thread::sleep(Duration::from_micros(500));

        Ok(ExecutionProfile {
            timestamp: start,
            execution_time: start.elapsed(),
            memorybandwidth: 10.0, // Lower bandwidth for interpreter
            compute_utilization: 0.1,
            cache_hit_rates: vec![1.0], // Perfect cache hit for interpreter
            power_consumption: Some(5.0), // Low power
        })
    }

    fn is_available(&self) -> bool {
        true // Interpreter is always available
    }

    fn get_capabilities(&self) -> BackendCapabilities {
        BackendCapabilities {
            supported_types: vec![DataType::I32, DataType::F32, DataType::F64, DataType::Bool],
            optimization_levels: vec![OptimizationLevel::None],
            max_kernel_size: Some(1024 * 1024), // 1MB limit for interpreter
            supports_debugging: true,
            supports_profiling: false,
            target_architectures: vec![TargetArchitecture::X86_64],
        }
    }
}

/// Convenience functions for common JIT operations
pub mod jit_dsl {
    use super::*;

    /// Create a simple arithmetic kernel
    pub fn create_arithmetic_kernel(
        operation: &str,
        input_type: DataType,
        output_type: DataType,
    ) -> KernelSource {
        let input_type_str = format!("{input_type:?}").to_lowercase();
        let output_type_str = format!("{output_type:?}").to_lowercase();

        let source = format!(
            r#"
kernel void arithmetic_op(global {input_type}* input, global {output_type}* output, int size) {{
    int idx = get_global_id(0);
    if (idx < size) {{
        output[idx] = {operation}(input[idx]);
    }}
}}
"#,
            input_type = input_type_str,
            output_type = output_type_str,
            operation = operation
        );

        KernelSource {
            id: format!("arithmetic_{operation}"),
            source,
            language: KernelLanguage::OpenCl,
            entry_point: "arithmetic_op".to_string(),
            input_types: vec![input_type],
            output_types: vec![output_type],
            hints: CompilationHints::default(),
        }
    }

    /// Create a reduction kernel
    pub fn create_reduction_kernel(operation: &str, datatype: DataType) -> KernelSource {
        let datatype_str = format!("{datatype:?}").to_lowercase();

        let source = format!(
            r#"
kernel void reduction_op(global {datatype}* input, global {datatype}* output, int size) {{
    local {datatype} shared_data[256];
    int tid = get_local_id(0);
    int gid = get_global_id(0);

    // Load data into shared memory
    shared_data[tid] = (gid < size) ? input[gid] : 0;
    barrier(CLK_LOCAL_MEM_FENCE);

    // Perform reduction
    for (int stride = get_local_size(0) / 2; stride > 0; stride /= 2) {{
        if (tid < stride) {{
            shared_data[tid] = {operation}(shared_data[tid], shared_data[tid + stride]);
        }}
        barrier(CLK_LOCAL_MEM_FENCE);
    }}

    // Write result
    if (tid == 0) {{
        output[get_group_id(0)] = shared_data[0];
    }}
}}
"#,
            datatype = datatype_str,
            operation = operation
        );

        KernelSource {
            id: format!("reduction_{operation}"),
            source,
            language: KernelLanguage::OpenCl,
            entry_point: "reduction_op".to_string(),
            input_types: vec![datatype.clone()],
            output_types: vec![datatype.clone()],
            hints: CompilationHints {
                workload_size: Some(1024),
                memory_pattern: Some(MemoryPattern::Sequential),
                compute_intensity: Some(ComputeIntensity::ComputeBound),
                parallelization: Some(ParallelizationHints {
                    work_group_size: Some([256, 1, 1]),
                    vector_width: Some(4),
                    unroll_factor: Some(4),
                    auto_vectorize: true,
                }),
                target_hints: HashMap::new(),
            },
        }
    }
}

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

    #[test]
    fn test_jit_compiler_creation() {
        let config = JitConfig::default();
        let compiler = JitCompiler::new(config);
        assert!(compiler.is_ok());
    }

    #[test]
    fn test_kernel_source_creation() {
        let source = KernelSource {
            id: "test_kernel".to_string(),
            source: "kernel void test() {}".to_string(),
            language: KernelLanguage::OpenCl,
            entry_point: "test".to_string(),
            input_types: vec![DataType::F32],
            output_types: vec![DataType::F32],
            hints: CompilationHints::default(),
        };

        assert_eq!(source.id, "test_kernel");
        assert_eq!(source.language, KernelLanguage::OpenCl);
    }

    #[test]
    fn test_dsl_arithmetic_kernel() {
        let kernel = jit_dsl::create_arithmetic_kernel("sqrt", DataType::F32, DataType::F32);
        assert_eq!(kernel.id, "arithmetic_sqrt");
        assert!(!kernel.source.is_empty());
        assert_eq!(kernel.input_types.len(), 1);
        assert_eq!(kernel.output_types.len(), 1);
    }

    #[test]
    fn test_dsl_reduction_kernel() {
        let kernel = jit_dsl::create_reduction_kernel("max", DataType::F32);
        assert_eq!(kernel.id, "reduction_max");
        assert!(!kernel.source.is_empty());
        assert!(kernel.hints.workload_size.is_some());
    }

    #[test]
    fn test_kernel_cache() {
        let mut cache = KernelCache::size(1024 * 1024); // 1MB cache

        let kernel = CompiledKernel {
            id: "test".to_string(),
            binary: vec![0; 1024],
            backend: JitBackend::Interpreter,
            target_arch: TargetArchitecture::X86_64,
            metadata: KernelMetadata {
                compiled_at: Instant::now(),
                compilation_time: Duration::from_millis(100),
                optimization_level: OptimizationLevel::O2,
                binary_size: 1024,
                register_usage: None,
                shared_memory_usage: None,
                compiler_info: "test".to_string(),
            },
            performance: KernelPerformance::default(),
        };

        cache.insert(kernel);
        assert!(cache.contains_kernel("test"));
        assert!(cache.get("test").is_some());
    }

    #[test]
    fn test_interpreter_backend() {
        let backend = InterpreterBackend::new();
        assert!(backend.is_available());

        let capabilities = backend.get_capabilities();
        assert!(!capabilities.supported_types.is_empty());
        assert!(capabilities.supports_debugging);
    }

    #[test]
    fn test_compilation_with_interpreter() {
        let config = JitConfig {
            backend: JitBackend::Interpreter,
            ..Default::default()
        };

        let compiler = JitCompiler::new(config).expect("Operation failed");

        let source = KernelSource {
            id: "test_kernel".to_string(),
            source: "void test() { /* test kernel */ }".to_string(),
            language: KernelLanguage::HighLevel,
            entry_point: "test".to_string(),
            input_types: vec![],
            output_types: vec![],
            hints: CompilationHints::default(),
        };

        let result = compiler.compile_kernel(source);
        assert!(result.is_ok());
    }

    // ── AdaptiveOptimizer tests ──────────────────────────────────────────────

    /// `update_performance_data` should record at least one entry in history.
    #[test]
    fn test_adaptive_optimizer_update_records_history() {
        let mut optimizer = AdaptiveOptimizer::new();

        let perf = KernelPerformance {
            execution_count: 5,
            totalexecution_time: Duration::from_millis(50),
            avgexecution_time: Duration::from_millis(10),
            bestexecution_time: Duration::from_millis(8),
            worstexecution_time: Duration::from_millis(15),
            throughput: 1.0e8,
            energy_efficiency: None,
        };

        optimizer.update_performance_data(&perf);

        // Should have recorded into the aggregate key.
        let history = optimizer
            .optimization_history
            .get("__perf_trends__")
            .expect("Expected history to be populated after update");
        assert_eq!(
            history.len(),
            1,
            "Exactly one record should have been added"
        );
    }

    /// Multiple `update_performance_data` calls accumulate multiple records.
    #[test]
    fn test_adaptive_optimizer_update_accumulates() {
        let mut optimizer = AdaptiveOptimizer::new();

        for i in 0..5u64 {
            let perf = KernelPerformance {
                execution_count: (i + 1) as usize,
                totalexecution_time: Duration::from_millis(10 * (i + 1)),
                avgexecution_time: Duration::from_millis(10),
                bestexecution_time: Duration::from_millis(8),
                worstexecution_time: Duration::from_millis(15),
                throughput: 2.0e9 * (i + 1) as f64, // varying throughput
                energy_efficiency: None,
            };
            optimizer.update_performance_data(&perf);
        }

        let history = optimizer
            .optimization_history
            .get("__perf_trends__")
            .expect("history should exist");
        assert_eq!(history.len(), 5);
    }

    /// `optimize_kernel` with no history should return a default directive.
    #[test]
    fn test_adaptive_optimizer_default_directive_no_history() {
        let optimizer = AdaptiveOptimizer::new();
        let config = JitConfig::default(); // O2
        let result = optimizer.optimize_kernel("unknown_kernel", &config);
        assert!(result.is_ok(), "Should succeed even with no history");
        let directive = result.expect("optimizer returned Err unexpectedly");
        assert!(!directive.is_empty(), "Directive string must not be empty");
        assert!(
            directive.contains("-O2"),
            "O2 config should produce -O2 flag, got: {directive}"
        );
    }

    /// `optimize_kernel` after recording successful improvements picks
    /// the strategy with the highest improvement.
    #[test]
    fn test_adaptive_optimizer_picks_best_strategy_from_history() {
        let mut optimizer = AdaptiveOptimizer::new();

        // Record a low-throughput run first (triggers ConstantFolding for count < 10).
        let perf_low = KernelPerformance {
            execution_count: 3,
            totalexecution_time: Duration::from_millis(300),
            avgexecution_time: Duration::from_millis(100),
            bestexecution_time: Duration::from_millis(90),
            worstexecution_time: Duration::from_millis(120),
            throughput: 5.0e7, // low
            energy_efficiency: None,
        };
        optimizer.update_performance_data(&perf_low);

        // Record a very high-throughput run (triggers Vectorization).
        let perf_high = KernelPerformance {
            execution_count: 20,
            totalexecution_time: Duration::from_millis(20),
            avgexecution_time: Duration::from_millis(1),
            bestexecution_time: Duration::from_millis(1),
            worstexecution_time: Duration::from_millis(2),
            throughput: 5.0e9, // > 1 GOp/s
            energy_efficiency: None,
        };
        optimizer.update_performance_data(&perf_high);

        let config = JitConfig {
            optimization_level: OptimizationLevel::O3,
            ..Default::default()
        };
        let result = optimizer
            .optimize_kernel("my_kernel", &config)
            .expect("optimize_kernel should succeed");

        // The vectorize record has improvement = 5.0e9 / 1e9 = 5.0 (clamped to 1.0)
        // which is higher than the constfold record, so "vectorize" should appear.
        assert!(
            result.contains("vectorize"),
            "Expected 'vectorize' in directive, got: {result}"
        );
    }
}