#![allow(unsafe_code)] // Required for CUDA driver API calls (cuCtxSetLimit, cuDeviceGetAttribute)

//! L2 cache persistence hints for H100/Hopper GPUs.
//!
//! Pins frequently-accessed GPU memory (model weights) in L2 cache
//! for ~3.5x effective bandwidth improvement over HBM reads.
//! No-op on non-Hopper GPUs or when L2 cache is too small.
//!
//! # How it works
//!
//! H100 SXM5 has 50 MB of L2 cache. DQN model weights are typically ~23 MB
//! in BF16 -- they fit entirely in L2. Without persistence hints, frequently
//! accessed weight data competes with transient data for L2 residency and
//! may be evicted between kernel launches.
//!
//! `cuCtxSetLimit(CU_LIMIT_PERSISTING_L2_CACHE_SIZE, bytes)` tells the CUDA
//! runtime to reserve up to `bytes` of L2 for persistent data. Combined with
//! CUDA's access pattern tracking, this keeps model weights hot in L2 across
//! successive forward passes.
//!
//! # Supported GPUs
//!
//! - **H100 (SXM5/PCIe)**: 50 MB L2 cache
//! - **H200**: 50 MB L2 cache
//! - **GH200**: 50+ MB L2 cache
//! - Other GPUs: no-op (L2 too small or persistence not beneficial)

/// Minimum L2 cache size (in bytes) for persistence to be worthwhile.
///
/// Below this threshold the L2 is too small to hold meaningful weight data
/// alongside transient activations and the persistence hint would cause
/// more harm than good (evicting activation data that could benefit from L2).
const MIN_L2_FOR_PERSISTENCE_BYTES: usize = 40 * 1024 * 1024; // 40 MB

/// Returns `true` if the GPU supports L2 cache persistence with a large
/// enough L2 to benefit model weight pinning.
///
/// Currently returns `true` for Hopper-class GPUs (H100, H200, GH200)
/// which have 50+ MB L2 cache. Ampere GPUs (A100: 40 MB L2) are excluded
/// because their L2 is borderline after accounting for driver and
/// activation overhead.
pub fn supports_l2_persistence(gpu_name: &str) -> bool {
    let name = gpu_name.to_uppercase();
    // Hopper-class GPUs with 50+ MB L2
    name.contains("H100")
        || name.contains("H200")
        || name.contains("GH200")
}

/// Query the maximum persisting L2 cache size (in bytes) from the CUDA driver.
///
/// Returns `Ok(0)` if the GPU does not support L2 persistence.
/// Returns `Err` if the CUDA driver call fails.
#[cfg(feature = "cuda")]
fn query_max_persisting_l2_bytes() -> Result<usize, String> {
    use candle_core::cuda_backend::cudarc::driver::sys::{
        CUdevice_attribute, CUresult,
    };

    // Query via the capabilities module -- it already has the device name cached.
    // But we also need the raw CUDA attribute for the exact L2 cache size.
    let caps = super::capabilities::cached_capabilities();
    if !caps.is_cuda {
        return Ok(0);
    }

    // Use nvidia-smi-detected device name to gate the feature.
    // Only proceed if the GPU actually supports persistence.
    if !supports_l2_persistence(&caps.device_name) {
        return Ok(0);
    }

    // Query the hardware limit via CUDA driver API.
    // CU_DEVICE_ATTRIBUTE_MAX_PERSISTING_L2_CACHE_SIZE returns bytes.
    let mut max_bytes: i32 = 0;
    let result = unsafe {
        candle_core::cuda_backend::cudarc::driver::sys::cuDeviceGetAttribute(
            &mut max_bytes,
            CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_PERSISTING_L2_CACHE_SIZE,
            0, // device ordinal 0
        )
    };
    if result != CUresult::CUDA_SUCCESS {
        return Err(format!(
            "cuDeviceGetAttribute(MAX_PERSISTING_L2_CACHE_SIZE) failed: {result:?}"
        ));
    }

    Ok(max_bytes.max(0) as usize)
}

/// Set the persisting L2 cache size for the current CUDA context.
///
/// `num_bytes` is clamped to the hardware maximum. The CUDA driver will
/// reserve up to this many bytes in L2 for persistent data, allowing
/// frequently-accessed memory (model weights) to remain cache-resident
/// across kernel launches.
///
/// Returns the actual number of bytes set (may be clamped by hardware).
///
/// # No-op conditions
///
/// Returns `Ok(0)` without calling the driver when:
/// - No CUDA GPU is detected
/// - The GPU does not support L2 persistence (non-Hopper)
/// - The max persisting L2 size is below [`MIN_L2_FOR_PERSISTENCE_BYTES`]
///
/// # Errors
///
/// Returns `Err` if the CUDA driver call fails.
#[cfg(feature = "cuda")]
pub fn set_persisting_l2_cache_size(num_bytes: usize) -> Result<usize, String> {
    use candle_core::cuda_backend::cudarc::driver::sys::{
        cuCtxSetLimit, CUlimit, CUresult,
    };

    let max_bytes = query_max_persisting_l2_bytes()?;
    if max_bytes < MIN_L2_FOR_PERSISTENCE_BYTES {
        tracing::debug!(
            max_bytes,
            min_required = MIN_L2_FOR_PERSISTENCE_BYTES,
            "L2 cache persistence: GPU L2 too small or not supported, skipping"
        );
        return Ok(0);
    }

    // Clamp to hardware maximum.
    let target = num_bytes.min(max_bytes);

    let result = unsafe {
        cuCtxSetLimit(CUlimit::CU_LIMIT_PERSISTING_L2_CACHE_SIZE, target)
    };
    if result != CUresult::CUDA_SUCCESS {
        return Err(format!(
            "cuCtxSetLimit(PERSISTING_L2_CACHE_SIZE, {target}) failed: {result:?}"
        ));
    }

    tracing::info!(
        target_bytes = target,
        max_hw_bytes = max_bytes,
        target_mb = target / (1024 * 1024),
        max_hw_mb = max_bytes / (1024 * 1024),
        "L2 cache persistence: reserved {} MB of {} MB max",
        target / (1024 * 1024),
        max_bytes / (1024 * 1024),
    );

    Ok(target)
}

/// CPU-only stub: always returns `Ok(0)`.
#[cfg(not(feature = "cuda"))]
pub fn set_persisting_l2_cache_size(_num_bytes: usize) -> Result<usize, String> {
    Ok(0)
}

/// Reset the persisting L2 cache, releasing all persistent cache lines.
///
/// Call this during shutdown or when switching between models of very
/// different sizes to avoid stale persistence reservations.
///
/// No-op on non-CUDA builds or when L2 persistence is not supported.
#[cfg(feature = "cuda")]
pub fn reset_persisting_l2_cache() -> Result<(), String> {
    use candle_core::cuda_backend::cudarc::driver::sys::{
        cuCtxResetPersistingL2Cache, CUresult,
    };

    let result = unsafe { cuCtxResetPersistingL2Cache() };
    if result != CUresult::CUDA_SUCCESS {
        return Err(format!(
            "cuCtxResetPersistingL2Cache() failed: {result:?}"
        ));
    }

    tracing::debug!("L2 cache persistence: reset all persistent cache lines");
    Ok(())
}

/// CPU-only stub: always returns `Ok(())`.
#[cfg(not(feature = "cuda"))]
pub fn reset_persisting_l2_cache() -> Result<(), String> {
    Ok(())
}

/// Convenience: configure L2 persistence for model weights if the GPU supports it.
///
/// Checks whether the detected GPU name supports L2 persistence (Hopper-class),
/// and if so sets the persisting L2 cache size to accommodate `weight_bytes`.
///
/// Returns the number of bytes actually reserved (0 if not supported or not applicable).
///
/// This is the primary entry point for callers. Typical usage:
/// ```ignore
/// let reserved = pin_weights_in_l2("NVIDIA H100 80GB HBM3", 23_000_000)?;
/// if reserved > 0 {
///     tracing::info!("L2 pinning active: {} MB reserved", reserved / 1_048_576);
/// }
/// ```
pub fn pin_weights_in_l2(gpu_name: &str, weight_bytes: usize) -> Result<usize, String> {
    if !supports_l2_persistence(gpu_name) {
        tracing::debug!(
            gpu = gpu_name,
            "L2 cache persistence: GPU not supported, skipping"
        );
        return Ok(0);
    }

    set_persisting_l2_cache_size(weight_bytes)
}

/// Convenience: unpin L2 cache (alias for [`reset_persisting_l2_cache`]).
pub fn unpin_l2() -> Result<(), String> {
    reset_persisting_l2_cache()
}

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

    #[test]
    fn test_supports_l2_persistence_hopper() {
        assert!(supports_l2_persistence("NVIDIA H100 80GB HBM3"));
        assert!(supports_l2_persistence("NVIDIA H100 SXM5"));
        assert!(supports_l2_persistence("NVIDIA H100 PCIe"));
        assert!(supports_l2_persistence("NVIDIA H200 141GB HBM3e"));
        assert!(supports_l2_persistence("NVIDIA GH200 Grace Hopper"));
    }

    #[test]
    fn test_supports_l2_persistence_non_hopper() {
        assert!(!supports_l2_persistence("NVIDIA A100-SXM4-80GB"));
        assert!(!supports_l2_persistence("NVIDIA L40S"));
        assert!(!supports_l2_persistence("NVIDIA GeForce RTX 4090"));
        assert!(!supports_l2_persistence("NVIDIA GeForce RTX 3050 Ti"));
        assert!(!supports_l2_persistence("Tesla V100"));
        assert!(!supports_l2_persistence("CPU"));
        assert!(!supports_l2_persistence(""));
    }

    #[test]
    fn test_supports_l2_persistence_case_insensitive() {
        assert!(supports_l2_persistence("nvidia h100"));
        assert!(supports_l2_persistence("h200 something"));
        assert!(supports_l2_persistence("gh200"));
    }

    #[test]
    fn test_pin_weights_non_hopper_is_noop() {
        let result = pin_weights_in_l2("NVIDIA A100-SXM4-80GB", 23_000_000);
        assert!(result.is_ok());
        assert_eq!(result.ok(), Some(0));
    }

    #[test]
    fn test_pin_weights_cpu_is_noop() {
        let result = pin_weights_in_l2("CPU", 23_000_000);
        assert!(result.is_ok());
        assert_eq!(result.ok(), Some(0));
    }

    #[test]
    fn test_unpin_l2_does_not_panic() {
        // On CPU builds or without CUDA context this should not panic.
        // It may return an error if there's no CUDA context, which is fine.
        let _ = unpin_l2();
    }

    #[test]
    fn test_set_persisting_l2_noop_on_cpu() {
        // Without a CUDA device, this should return 0 or an error gracefully.
        let result = set_persisting_l2_cache_size(50 * 1024 * 1024);
        // On CPU test machines this will either return Ok(0) or an error
        // about missing CUDA context -- both are acceptable.
        match result {
            Ok(0) => {} // Expected on CPU
            Err(_) => {} // Also acceptable -- CUDA driver not available
            Ok(n) => {
                // If we somehow have a GPU, the value should be reasonable
                assert!(n <= 50 * 1024 * 1024);
            }
        }
    }
}