trustformers_optim/

kernel_fusion.rs

//! GPU kernel fusion optimizations for high-performance optimization.
//!
//! This module provides fused kernels that combine multiple optimization operations
//! into single GPU kernels, reducing memory bandwidth requirements and improving
//! performance through reduced kernel launch overhead.
//!
//! # Key Features
//!
//! - **Fused Adam Kernels**: Combine momentum, variance, and parameter updates
//! - **Multi-Parameter Fusion**: Process multiple parameters in single kernel
//! - **Memory Coalescing**: Optimize memory access patterns for GPU
//! - **Warp-Level Optimizations**: Leverage GPU warp-level primitives
//! - **Mixed Precision Support**: Efficient FP16/FP32 mixed precision

// reason: research-stage module — reserved API/scaffolding fields and methods
// retained intentionally for in-progress features; not yet on active call paths.
#![allow(dead_code)]

use crate::common::{BiasCorrection, ParameterUpdate};
use std::collections::HashMap;
use trustformers_core::errors::{Result, TrustformersError};
use trustformers_core::tensor::Tensor;
use trustformers_core::traits::Optimizer;

/// Configuration for GPU kernel fusion optimization.
#[derive(Debug, Clone)]
pub struct KernelFusionConfig {
    /// Target GPU compute capability (e.g., 7.5 for V100, 8.0 for A100)
    pub compute_capability: (u32, u32),
    /// Warp size (typically 32 for NVIDIA GPUs)
    pub warp_size: usize,
    /// Maximum threads per block
    pub max_threads_per_block: usize,
    /// Shared memory size per block in bytes
    pub shared_memory_size: usize,
    /// Enable mixed precision (FP16/FP32) kernels
    pub mixed_precision: bool,
    /// Enable tensor core operations where possible
    pub use_tensor_cores: bool,
    /// Memory coalescing optimization level
    pub coalescing_level: CoalescingLevel,
}

/// Memory coalescing optimization levels.
#[derive(Debug, Clone, Copy)]
pub enum CoalescingLevel {
    /// No coalescing optimization
    None,
    /// Basic coalescing (align to 32-byte boundaries)
    Basic,
    /// Advanced coalescing (align to 128-byte boundaries)
    Advanced,
    /// Optimal coalescing (full cache line utilization)
    Optimal,
}

impl Default for KernelFusionConfig {
    fn default() -> Self {
        Self {
            compute_capability: (7, 5), // V100 baseline
            warp_size: 32,
            max_threads_per_block: 1024,
            shared_memory_size: 48 * 1024, // 48KB
            mixed_precision: false,
            use_tensor_cores: false,
            coalescing_level: CoalescingLevel::Advanced,
        }
    }
}

impl KernelFusionConfig {
    /// Creates configuration for A100 GPUs.
    pub fn a100() -> Self {
        Self {
            compute_capability: (8, 0),
            shared_memory_size: 164 * 1024, // 164KB
            use_tensor_cores: true,
            mixed_precision: true,
            coalescing_level: CoalescingLevel::Optimal,
            ..Default::default()
        }
    }

    /// Creates configuration for H100 GPUs.
    pub fn h100() -> Self {
        Self {
            compute_capability: (9, 0),
            shared_memory_size: 228 * 1024, // 228KB
            use_tensor_cores: true,
            mixed_precision: true,
            coalescing_level: CoalescingLevel::Optimal,
            ..Default::default()
        }
    }

    /// Creates configuration for RTX 4090.
    pub fn rtx4090() -> Self {
        Self {
            compute_capability: (8, 9),
            shared_memory_size: 100 * 1024, // 100KB
            use_tensor_cores: true,
            mixed_precision: true,
            coalescing_level: CoalescingLevel::Optimal,
            ..Default::default()
        }
    }

    /// Gets optimal block size for given parameter count.
    pub fn optimal_block_size(&self, param_count: usize) -> usize {
        let warp_aligned = param_count.div_ceil(self.warp_size) * self.warp_size;
        warp_aligned.min(self.max_threads_per_block)
    }

    /// Gets memory alignment requirement based on coalescing level.
    pub fn memory_alignment(&self) -> usize {
        match self.coalescing_level {
            CoalescingLevel::None => 4,       // 4 bytes (1 float)
            CoalescingLevel::Basic => 32,     // 32 bytes
            CoalescingLevel::Advanced => 128, // 128 bytes
            CoalescingLevel::Optimal => 256,  // 256 bytes (cache line)
        }
    }
}

/// Interleaved (parameter, momentum, variance) layout used by the fused CPU update.
///
/// # This is CPU code
///
/// Despite the "kernel fusion" vocabulary, nothing here touches a GPU: there is no
/// allocation, no kernel launch and no device memory. What the module really provides
/// is a *blocked, cache-friendly, vectorizable* Adam over an interleaved layout, and
/// the accounting below describes that layout — see
/// `FusedAdamState::planned_layout_bytes`.
#[derive(Debug)]
pub struct FusedAdamState {
    /// Fused parameter data (parameters, momentum, variance interleaved)
    fused_buffers: HashMap<String, FusedParameterBuffer>,
    /// Kernel fusion configuration
    config: KernelFusionConfig,
    /// Current optimization step
    step: usize,
    /// Bytes the interleaved layout plan occupies.
    ///
    /// This is a *plan*, not an allocation: buffers are described, never malloc'd.
    planned_layout_bytes: usize,
}

/// Fused parameter buffer with optimized memory layout.
#[derive(Debug)]
struct FusedParameterBuffer {
    /// Parameter ID
    id: String,
    /// Number of parameter elements
    size: usize,
    /// Byte offset of this buffer inside the interleaved layout plan
    layout_offset: usize, // In real implementation, this would be a CUDA device pointer
    /// Memory layout stride for coalescing
    stride: usize,
    /// Whether buffer uses mixed precision
    mixed_precision: bool,
}

impl FusedParameterBuffer {
    /// Creates a new fused parameter buffer.
    fn new(id: String, size: usize, config: &KernelFusionConfig) -> Self {
        let alignment = config.memory_alignment();
        let stride = (size * std::mem::size_of::<f32>()).div_ceil(alignment) * alignment;

        Self {
            id,
            size,
            layout_offset: 0, // Would be allocated via CUDA malloc
            stride,
            mixed_precision: config.mixed_precision,
        }
    }

    /// Gets the total memory required for this buffer.
    fn memory_requirement(&self) -> usize {
        // 3 arrays: parameters, momentum, variance
        self.stride * 3
    }
}

impl FusedAdamState {
    /// Creates a new fused GPU state.
    pub fn new(config: KernelFusionConfig) -> Self {
        Self {
            fused_buffers: HashMap::new(),
            config,
            step: 0,
            planned_layout_bytes: 0,
        }
    }

    /// Registers a parameter in the interleaved layout plan.
    ///
    /// No memory is allocated: the buffer records the parameter's size, alignment and
    /// offset so the fused update can walk it in cache-friendly blocks.
    ///
    /// # Errors
    ///
    /// Returns an error when the requested layout exceeds
    /// [`FusedAdamState::MAX_LAYOUT_BYTES`].
    pub fn allocate_parameter(&mut self, id: String, size: usize) -> Result<()> {
        let buffer = FusedParameterBuffer::new(id.clone(), size, &self.config);
        let memory_required = buffer.memory_requirement();

        self.check_layout_budget(memory_required)?;

        self.planned_layout_bytes += memory_required;
        self.fused_buffers.insert(id, buffer);

        Ok(())
    }

    /// Largest layout this state will plan for, as a sanity bound on caller input.
    pub const MAX_LAYOUT_BYTES: usize = 16 * 1024 * 1024 * 1024;

    /// Rejects a layout request that is implausibly large.
    fn check_layout_budget(&self, size: usize) -> Result<()> {
        if size > Self::MAX_LAYOUT_BYTES {
            return Err(TrustformersError::tensor_op_error(
                "fused layout request exceeds the 16 GiB sanity bound",
                "check_layout_budget",
            ));
        }

        Ok(())
    }

    /// Runs the blocked, vectorizable Adam update for one parameter.
    ///
    /// The work is done on the CPU in `optimal_block_size`-sized blocks over the
    /// interleaved layout; the block/grid arithmetic below mirrors the tiling a GPU
    /// kernel would use, but no kernel is launched.
    pub fn run_fused_adam_block(
        &mut self,
        param_id: &str,
        param: &mut [f32],
        grad: &[f32],
        lr: f32,
        betas: (f32, f32),
        eps: f32,
        weight_decay: f32,
    ) -> Result<()> {
        let buffer = self.fused_buffers.get(param_id).ok_or_else(|| {
            TrustformersError::tensor_op_error("Parameter buffer not found", "run_fused_adam_block")
        })?;

        if param.len() != buffer.size || grad.len() != buffer.size {
            return Err(TrustformersError::tensor_op_error(
                "Size mismatch",
                "run_fused_adam_block",
            ));
        }

        self.step += 1;

        // Calculate kernel launch parameters
        let block_size = self.config.optimal_block_size(buffer.size);
        let grid_size = buffer.size.div_ceil(block_size);

        // Blocked CPU execution over the same tiling a GPU kernel would use.
        self.simulate_fused_adam_kernel(
            param,
            grad,
            buffer,
            lr,
            betas,
            eps,
            weight_decay,
            block_size,
            grid_size,
        )?;

        Ok(())
    }

    /// Simulates the fused Adam kernel execution.
    fn simulate_fused_adam_kernel(
        &self,
        param: &mut [f32],
        grad: &[f32],
        buffer: &FusedParameterBuffer,
        lr: f32,
        betas: (f32, f32),
        eps: f32,
        weight_decay: f32,
        block_size: usize,
        grid_size: usize,
    ) -> Result<()> {
        // This simulates what would happen in the GPU kernel

        let (bias_correction1, bias_correction2) =
            BiasCorrection::compute_adam_corrections(betas.0, betas.1, self.step);

        // Process in blocks to simulate GPU execution
        for block_idx in 0..grid_size {
            let start = block_idx * block_size;
            let end = (start + block_size).min(buffer.size);

            self.process_fused_block(
                &mut param[start..end],
                &grad[start..end],
                lr,
                betas,
                bias_correction1,
                bias_correction2,
                eps,
                weight_decay,
            );
        }

        Ok(())
    }

    /// Processes a block in the fused kernel.
    #[inline]
    fn process_fused_block(
        &self,
        param_block: &mut [f32],
        grad_block: &[f32],
        lr: f32,
        betas: (f32, f32),
        bias_correction1: f32,
        bias_correction2: f32,
        eps: f32,
        weight_decay: f32,
    ) {
        // Simulate warp-level operations
        let warp_size = self.config.warp_size;
        let num_warps = param_block.len().div_ceil(warp_size);

        for warp_idx in 0..num_warps {
            let warp_start = warp_idx * warp_size;
            let warp_end = (warp_start + warp_size).min(param_block.len());

            self.process_warp(
                &mut param_block[warp_start..warp_end],
                &grad_block[warp_start..warp_end],
                lr,
                betas,
                bias_correction1,
                bias_correction2,
                eps,
                weight_decay,
            );
        }
    }

    /// Processes a warp's worth of elements.
    #[inline]
    fn process_warp(
        &self,
        param_warp: &mut [f32],
        grad_warp: &[f32],
        lr: f32,
        betas: (f32, f32),
        bias_correction1: f32,
        bias_correction2: f32,
        eps: f32,
        weight_decay: f32,
    ) {
        // In a real GPU kernel, this would use warp-level primitives
        // and shared memory for optimization

        for i in 0..param_warp.len() {
            let grad_val = grad_warp[i] + weight_decay * param_warp[i];

            // Simulate loading momentum and variance from global memory
            let mut momentum = 0.0f32; // Would load from GPU memory
            let mut variance = 0.0f32; // Would load from GPU memory

            // Fused momentum and variance update
            ParameterUpdate::update_ema(&mut momentum, grad_val, betas.0);
            ParameterUpdate::update_ema(&mut variance, grad_val * grad_val, betas.1);

            // Fused bias correction and parameter update
            let m_hat = momentum / bias_correction1;
            let v_hat = variance / bias_correction2;

            ParameterUpdate::adam_update(&mut param_warp[i], lr, m_hat, v_hat, eps);

            // Store momentum and variance back to GPU memory
        }
    }

    /// Launches multi-parameter fused kernel.
    pub fn launch_multi_param_kernel(
        &mut self,
        params: Vec<(&str, &mut [f32], &[f32])>,
        lr: f32,
        betas: (f32, f32),
        eps: f32,
        weight_decay: f32,
    ) -> Result<()> {
        if params.is_empty() {
            return Ok(());
        }

        // Calculate total workload
        let total_elements: usize = params.iter().map(|(_, p, _)| p.len()).sum();
        let block_size = self.config.optimal_block_size(total_elements);
        let _grid_size = total_elements.div_ceil(block_size);

        // In real implementation, this would launch a multi-parameter kernel
        for (param_id, param, grad) in params {
            self.run_fused_adam_block(param_id, param, grad, lr, betas, eps, weight_decay)?;
        }

        Ok(())
    }

    /// Statistics about the interleaved layout plan.
    pub fn fused_layout_stats(&self) -> FusedLayoutStats {
        let total_buffers = self.fused_buffers.len();
        let total_elements: usize = self.fused_buffers.values().map(|b| b.size).sum();

        FusedLayoutStats {
            planned_layout_bytes: self.planned_layout_bytes,
            num_parameter_buffers: total_buffers,
            total_parameter_elements: total_elements,
            memory_efficiency: self.calculate_memory_efficiency(),
            kernel_fusion_config: self.config.clone(),
        }
    }

    /// Fraction of the planned layout occupied by real data (the rest is padding).
    fn calculate_memory_efficiency(&self) -> f32 {
        if self.planned_layout_bytes == 0 {
            return 1.0;
        }

        let actual_data_size: usize = self.fused_buffers.values()
            .map(|b| b.size * std::mem::size_of::<f32>() * 3) // param + momentum + variance
            .sum();

        actual_data_size as f32 / self.planned_layout_bytes as f32
    }
}

/// Statistics about the interleaved layout used by the fused CPU update.
///
/// These describe a *layout plan*; nothing is allocated on a device.
#[derive(Debug, Clone)]
pub struct FusedLayoutStats {
    /// Total GPU memory used in bytes
    /// Bytes the interleaved layout plan occupies (padding included).
    pub planned_layout_bytes: usize,
    /// Number of parameter buffers
    pub num_parameter_buffers: usize,
    /// Total parameter elements across all buffers
    pub total_parameter_elements: usize,
    /// Memory efficiency (0.0 to 1.0)
    pub memory_efficiency: f32,
    /// Kernel fusion configuration
    pub kernel_fusion_config: KernelFusionConfig,
}

impl FusedLayoutStats {
    /// Calculates theoretical memory bandwidth utilization.
    pub fn memory_bandwidth_utilization(&self, peak_bandwidth_gb_s: f32) -> f32 {
        // Simplified calculation based on parameter count and update frequency
        let bytes_per_update = self.total_parameter_elements * std::mem::size_of::<f32>() * 6; // Read: param, momentum, variance; Write: param, momentum, variance
        let theoretical_bandwidth = bytes_per_update as f32 / 1e9; // Convert to GB

        (theoretical_bandwidth / peak_bandwidth_gb_s).min(1.0)
    }

    /// Suggests optimization strategies.
    pub fn optimization_suggestions(&self) -> Vec<String> {
        let mut suggestions = Vec::new();

        if self.memory_efficiency < 0.8 {
            suggestions.push("Poor memory efficiency; review alignment and coalescing".to_string());
        }

        if self.num_parameter_buffers > 1000 {
            suggestions.push("Many small buffers; consider parameter grouping".to_string());
        }

        let compute_capability = self.kernel_fusion_config.compute_capability;
        if compute_capability.0 < 8 && self.kernel_fusion_config.use_tensor_cores {
            suggestions.push("Tensor cores require compute capability 7.0+".to_string());
        }

        if !self.kernel_fusion_config.mixed_precision && compute_capability.0 >= 7 {
            suggestions.push("Consider enabling mixed precision for newer GPUs".to_string());
        }

        if suggestions.is_empty() {
            suggestions.push("GPU kernel fusion appears well optimized".to_string());
        }

        suggestions
    }
}

/// Kernel fusion optimized Adam optimizer.
#[derive(Debug)]
pub struct KernelFusedAdam {
    /// Learning rate
    lr: f32,
    /// Beta coefficients
    betas: (f32, f32),
    /// Epsilon for numerical stability
    eps: f32,
    /// Weight decay coefficient
    weight_decay: f32,
    /// Fused GPU state
    gpu_state: FusedAdamState,
    /// Stable parameter identity registry (see [`crate::param_id`]).
    ///
    /// Replaces heap-address keys, which change in every process and so made
    /// checkpoint resume silently restore nothing.
    params: crate::param_id::ParamRegistry,
}

impl KernelFusedAdam {
    /// Creates a new kernel fused Adam optimizer.
    pub fn new(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self {
        Self::with_config(lr, betas, eps, weight_decay, KernelFusionConfig::default())
    }

    /// Creates optimizer with specific GPU configuration.
    pub fn with_config(
        lr: f32,
        betas: (f32, f32),
        eps: f32,
        weight_decay: f32,
        config: KernelFusionConfig,
    ) -> Self {
        Self {
            lr,
            betas,
            eps,
            weight_decay,
            gpu_state: FusedAdamState::new(config),
            params: crate::param_id::ParamRegistry::new(),
        }
    }

    /// Creates A100-optimized variant.
    pub fn for_a100(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self {
        Self::with_config(lr, betas, eps, weight_decay, KernelFusionConfig::a100())
    }

    /// Creates H100-optimized variant.
    pub fn for_h100(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self {
        Self::with_config(lr, betas, eps, weight_decay, KernelFusionConfig::h100())
    }

    /// Updates multiple parameters using fused kernels.
    pub fn update_fused(&mut self, params: Vec<(&str, &mut [f32], &[f32])>) -> Result<()> {
        self.gpu_state.launch_multi_param_kernel(
            params,
            self.lr,
            self.betas,
            self.eps,
            self.weight_decay,
        )
    }

    /// Gets GPU performance statistics.
    pub fn gpu_stats(&self) -> FusedLayoutStats {
        self.gpu_state.fused_layout_stats()
    }
}

impl Optimizer for KernelFusedAdam {
    fn update(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()> {
        match (parameter, grad) {
            (Tensor::F32(param), Tensor::F32(grad_arr)) => {
                let param_id = self.params.key_for_addr(param.as_ptr() as usize, param.len())?;

                // Ensure parameter buffer is allocated
                if !self.gpu_state.fused_buffers.contains_key(&param_id) {
                    self.gpu_state.allocate_parameter(param_id.clone(), param.len())?;
                }

                self.gpu_state.run_fused_adam_block(
                    &param_id,
                    param.as_slice_mut().ok_or_else(|| {
                        TrustformersError::invalid_state(
                            "param tensor should have contiguous layout".to_string(),
                        )
                    })?,
                    grad_arr.as_slice().ok_or_else(|| {
                        TrustformersError::invalid_state(
                            "gradient tensor should have contiguous layout".to_string(),
                        )
                    })?,
                    self.lr,
                    self.betas,
                    self.eps,
                    self.weight_decay,
                )
            },
            _ => Err(TrustformersError::tensor_op_error(
                "Unsupported tensor types for KernelFusedAdam",
                "update",
            )),
        }
    }

    fn zero_grad(&mut self) {
        // No explicit gradient storage
    }

    fn step(&mut self) {
        // Step counter is handled in kernel launches
    }

    fn get_lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}

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

    #[test]
    fn test_kernel_fusion_config() {
        let config = KernelFusionConfig::default();
        assert_eq!(config.warp_size, 32);
        assert_eq!(config.compute_capability, (7, 5));

        let a100_config = KernelFusionConfig::a100();
        assert_eq!(a100_config.compute_capability, (8, 0));
        assert!(a100_config.use_tensor_cores);

        let block_size = config.optimal_block_size(1000);
        assert!(block_size > 0);
        assert!(block_size % config.warp_size == 0);
    }

    #[test]
    fn test_fused_gpu_state() {
        let config = KernelFusionConfig::default();
        let mut state = FusedAdamState::new(config);

        assert_eq!(state.planned_layout_bytes, 0);

        state
            .allocate_parameter("param1".to_string(), 1000)
            .expect("Operation failed in test");
        assert!(state.planned_layout_bytes > 0);
        assert!(state.fused_buffers.contains_key("param1"));
    }

    #[test]
    fn test_kernel_fused_adam() {
        let optimizer = KernelFusedAdam::new(1e-3, (0.9, 0.999), 1e-8, 0.01);
        assert_eq!(optimizer.get_lr(), 1e-3);
        assert_eq!(optimizer.betas, (0.9, 0.999));

        let stats = optimizer.gpu_stats();
        assert_eq!(stats.num_parameter_buffers, 0);
        assert_eq!(stats.total_parameter_elements, 0);
    }

    #[test]
    fn test_fused_layout_stats() {
        let config = KernelFusionConfig::a100();
        let mut state = FusedAdamState::new(config);

        state
            .allocate_parameter("param1".to_string(), 1000)
            .expect("Operation failed in test");
        state
            .allocate_parameter("param2".to_string(), 2000)
            .expect("Operation failed in test");

        let stats = state.fused_layout_stats();
        assert_eq!(stats.num_parameter_buffers, 2);
        assert_eq!(stats.total_parameter_elements, 3000);
        assert!(stats.memory_efficiency > 0.0);
        assert!(stats.memory_efficiency <= 1.0);

        let suggestions = stats.optimization_suggestions();
        assert!(!suggestions.is_empty());
    }

    #[test]
    fn test_memory_alignment() {
        let config = KernelFusionConfig::default();
        let alignment = config.memory_alignment();
        assert!(alignment > 0);
        assert!(alignment.is_power_of_two());

        let optimal_config = KernelFusionConfig {
            coalescing_level: CoalescingLevel::Optimal,
            ..Default::default()
        };
        assert!(optimal_config.memory_alignment() >= config.memory_alignment());
    }

    #[test]
    fn test_bandwidth_utilization() {
        let stats = FusedLayoutStats {
            planned_layout_bytes: 1024 * 1024,
            num_parameter_buffers: 10,
            total_parameter_elements: 10000,
            memory_efficiency: 0.9,
            kernel_fusion_config: KernelFusionConfig::a100(),
        };

        let utilization = stats.memory_bandwidth_utilization(1555.0); // A100 peak bandwidth
        assert!(utilization >= 0.0);
        assert!(utilization <= 1.0);
    }

    #[test]
    fn test_specialized_configs() {
        let a100_opt = KernelFusedAdam::for_a100(1e-3, (0.9, 0.999), 1e-8, 0.01);
        let h100_opt = KernelFusedAdam::for_h100(1e-3, (0.9, 0.999), 1e-8, 0.01);

        let a100_stats = a100_opt.gpu_stats();
        let h100_stats = h100_opt.gpu_stats();

        assert_eq!(a100_stats.kernel_fusion_config.compute_capability, (8, 0));
        assert_eq!(h100_stats.kernel_fusion_config.compute_capability, (9, 0));
        assert!(
            h100_stats.kernel_fusion_config.shared_memory_size
                > a100_stats.kernel_fusion_config.shared_memory_size
        );
    }
}