rustytorch/crates/core/rtx-backend-cuda/src/kernels.rs

//! CUDA kernel management for rtx-backend-cuda.
//!
//! This module provides kernel loading, caching, and launch utilities
//! for GPU-accelerated tensor operations.

#[cfg(feature = "cuda")]
use cudarc::driver::{CudaFunction, CudaModule, LaunchConfig};
#[cfg(feature = "cuda")]
use cudarc::nvrtc::{CompileOptions, Ptx, compile_ptx_with_opts};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;

use crate::{CudaDevice, CudaError, CudaResult};

/// CUDA source for element-wise kernels (embedded for runtime compilation).
pub const ELEMENT_WISE_CUDA_SRC: &str =
    include_str!("../../rtx-tensor/src/cuda_kernels/element_wise.cu");

/// CUDA source preamble with missing defines for NVRTC.
const CUDA_PREAMBLE: &str = r#"
#ifndef INFINITY
#define INFINITY __int_as_float(0x7f800000)
#endif
#ifndef NAN
#define NAN __int_as_float(0x7fffffff)
#endif
"#;

#[cfg(feature = "cuda")]
/// Compile CUDA source to PTX at runtime using NVRTC.
/// This ensures PTX compatibility with the current driver.
fn compile_cuda_to_ptx(cuda_src: &str) -> CudaResult<Ptx> {
    // Add preamble with missing definitions
    let full_source = format!("{}\n{}", CUDA_PREAMBLE, cuda_src);

    // Use NVRTC with architecture targeting current GPU
    // Add CUDA include paths for standard headers
    let opts = CompileOptions {
        arch: Some("sm_86"), // Target RTX 3050/3090 (Ampere)
        include_paths: vec![
            "/usr/local/cuda/include".to_string(),
            "/usr/local/cuda-13.1/include".to_string(),
        ],
        ..Default::default()
    };

    compile_ptx_with_opts(&full_source, opts)
        .map_err(|e| CudaError::KernelCompilation(format!("NVRTC compilation failed: {:?}", e)))
}

#[cfg(feature = "cuda")]
/// Lazily compiled PTX for element-wise kernels.
static COMPILED_ELEMENT_WISE_PTX: once_cell::sync::Lazy<Result<Ptx, String>> =
    once_cell::sync::Lazy::new(|| {
        compile_cuda_to_ptx(ELEMENT_WISE_CUDA_SRC).map_err(|e| format!("{:?}", e))
    });

#[cfg(feature = "cuda")]
/// Get the element-wise PTX (compiled at runtime).
pub fn get_element_wise_ptx() -> CudaResult<&'static Ptx> {
    COMPILED_ELEMENT_WISE_PTX
        .as_ref()
        .map_err(|e| CudaError::KernelCompilation(e.clone()))
}

#[cfg(feature = "cuda")]
/// Global module cache: (device_index, module_id) -> CudaModule
static MODULE_CACHE: once_cell::sync::Lazy<
    RwLock<HashMap<(usize, &'static str), Arc<CudaModule>>>,
> = once_cell::sync::Lazy::new(|| RwLock::new(HashMap::new()));

#[cfg(feature = "cuda")]
/// Get or load the element-wise CUDA module (uses runtime-compiled PTX).
fn get_or_load_element_wise_module(device: &CudaDevice) -> CudaResult<Arc<CudaModule>> {
    let cache_key = (device.index, "element_wise");

    // Check cache first (read lock)
    {
        let cache = MODULE_CACHE.read();
        if let Some(module) = cache.get(&cache_key) {
            return Ok(module.clone());
        }
    }

    // Load and cache module (write lock)
    let mut cache = MODULE_CACHE.write();

    // Double-check after acquiring write lock
    if let Some(module) = cache.get(&cache_key) {
        return Ok(module.clone());
    }

    // Get runtime-compiled PTX
    let ptx = get_element_wise_ptx()?;

    // Load PTX module
    let module = device
        .context
        .load_module(ptx.clone())
        .map_err(|e| CudaError::KernelCompilation(format!("Failed to load PTX module: {:?}", e)))?;

    cache.insert(cache_key, module.clone());
    Ok(module)
}

#[cfg(feature = "cuda")]
/// Get a CUDA kernel function from the element-wise module.
///
/// This function caches modules per device to avoid redundant PTX compilation.
/// The PTX is compiled at runtime using NVRTC for driver compatibility.
pub fn get_kernel(
    device: &CudaDevice,
    kernel_name: &str,
    _ptx: &str, // Ignored - we use runtime-compiled PTX
) -> CudaResult<CudaFunction> {
    let module = get_or_load_element_wise_module(device)?;

    module.load_function(kernel_name).map_err(|e| {
        CudaError::KernelCompilation(format!(
            "Failed to load kernel function '{}': {:?}",
            kernel_name, e
        ))
    })
}

/// Constant for backwards compatibility (not used, we compile at runtime).
pub const ELEMENT_WISE_PTX: &str = "";

// Stub implementations for non-CUDA builds
#[cfg(not(feature = "cuda"))]
/// Stub launch config type.
#[derive(Clone, Debug)]
pub struct LaunchConfig {
    pub grid_dim: (u32, u32, u32),
    pub block_dim: (u32, u32, u32),
    pub shared_mem_bytes: u32,
}

#[cfg(not(feature = "cuda"))]
/// Stub kernel function type.
pub struct CudaFunction;

#[cfg(not(feature = "cuda"))]
pub fn get_kernel(
    _device: &crate::CudaDevice,
    _kernel_name: &str,
    _ptx: &str,
) -> crate::CudaResult<CudaFunction> {
    Err(crate::CudaError::KernelCompilation(
        "CUDA support not compiled".to_string(),
    ))
}

#[cfg(not(feature = "cuda"))]
pub fn element_wise_config(_numel: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (0, 0, 0),
        block_dim: (0, 0, 0),
        shared_mem_bytes: 0,
    }
}

#[cfg(not(feature = "cuda"))]
pub fn reduction_config(_numel: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (0, 0, 0),
        block_dim: (0, 0, 0),
        shared_mem_bytes: 0,
    }
}

#[cfg(not(feature = "cuda"))]
pub fn row_wise_config(_num_rows: usize, _row_size: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (0, 0, 0),
        block_dim: (0, 0, 0),
        shared_mem_bytes: 0,
    }
}

#[cfg(not(feature = "cuda"))]
pub fn transpose_2d_config(_rows: usize, _cols: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (0, 0, 0),
        block_dim: (0, 0, 0),
        shared_mem_bytes: 0,
    }
}

#[cfg(not(feature = "cuda"))]
pub fn dim_reduction_config(
    _outer_size: usize,
    _inner_size: usize,
    _reduce_size: usize,
) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (0, 0, 0),
        block_dim: (0, 0, 0),
        shared_mem_bytes: 0,
    }
}

#[cfg(feature = "cuda")]
/// Standard launch configuration for element-wise operations.
///
/// Uses 256 threads per block (optimal for memory coalescing on modern GPUs).
pub fn element_wise_config(numel: usize) -> LaunchConfig {
    let block_size = 256u32;
    let grid_size = ((numel as u32) + block_size - 1) / block_size;
    LaunchConfig {
        grid_dim: (grid_size, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 0,
    }
}

#[cfg(feature = "cuda")]
/// Launch configuration for reduction operations with shared memory.
pub fn reduction_config(numel: usize) -> LaunchConfig {
    let block_size = 256u32;
    let grid_size = ((numel as u32) + block_size - 1) / block_size;
    // Limit grid size for reduction (we'll do multi-pass for large tensors)
    let grid_size = grid_size.min(1024);
    LaunchConfig {
        grid_dim: (grid_size, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: block_size * std::mem::size_of::<f32>() as u32,
    }
}

#[cfg(feature = "cuda")]
/// Launch configuration for row-wise operations (softmax, layer_norm, rms_norm).
///
/// Each row is processed by one block using shared memory for reductions.
/// - `num_rows`: Number of rows to process (grid dimension)
/// - `row_size`: Size of each row (used for shared memory allocation)
pub fn row_wise_config(num_rows: usize, row_size: usize) -> LaunchConfig {
    // Block size must be a power of 2 for the parallel reduction to work correctly.
    // Use the next power of 2 >= row_size, capped at 256.
    let block_size = next_power_of_2(row_size as u32).min(256);
    // Shared memory: floats for parallel reductions
    let shared_mem = block_size * std::mem::size_of::<f32>() as u32;
    LaunchConfig {
        grid_dim: (num_rows as u32, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: shared_mem,
    }
}

#[cfg(feature = "cuda")]
/// Round up to the next power of 2.
fn next_power_of_2(n: u32) -> u32 {
    if n == 0 {
        return 1;
    }
    let mut v = n - 1;
    v |= v >> 1;
    v |= v >> 2;
    v |= v >> 4;
    v |= v >> 8;
    v |= v >> 16;
    v + 1
}

#[cfg(feature = "cuda")]
/// Launch configuration for 2D matrix transpose using 32x32 tiles.
///
/// Uses shared memory for coalesced access pattern.
/// - `rows`: Number of rows in input matrix
/// - `cols`: Number of columns in input matrix
pub fn transpose_2d_config(rows: usize, cols: usize) -> LaunchConfig {
    let tile_size = 32u32;
    let grid_x = ((cols as u32) + tile_size - 1) / tile_size;
    let grid_y = ((rows as u32) + tile_size - 1) / tile_size;
    LaunchConfig {
        grid_dim: (grid_x, grid_y, 1),
        block_dim: (tile_size, tile_size, 1),
        shared_mem_bytes: 0, // Shared memory is statically allocated in kernel
    }
}

#[cfg(feature = "cuda")]
/// Launch configuration for dimensional reduction (sum_dim, max_dim, min_dim).
///
/// Each output element is computed by one block.
/// - `outer_size`: Product of dimensions before the reduction dimension
/// - `inner_size`: Product of dimensions after the reduction dimension
/// - `reduce_size`: Size of the dimension being reduced
pub fn dim_reduction_config(
    outer_size: usize,
    inner_size: usize,
    reduce_size: usize,
) -> LaunchConfig {
    // Block size must be a power of 2 for the parallel reduction to work correctly.
    let block_size = next_power_of_2(reduce_size as u32).min(256).max(32);
    let shared_mem = block_size * std::mem::size_of::<f32>() as u32;
    LaunchConfig {
        grid_dim: (outer_size as u32, inner_size as u32, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: shared_mem,
    }
}

#[cfg(all(test, feature = "cuda"))]
mod tests {
    use super::*;

    #[test]
    fn test_element_wise_config() {
        let config = element_wise_config(1000);
        assert_eq!(config.block_dim.0, 256);
        assert_eq!(config.grid_dim.0, 4); // ceil(1000/256) = 4
    }

    #[test]
    fn test_reduction_config() {
        let config = reduction_config(1_000_000);
        assert_eq!(config.block_dim.0, 256);
        assert!(config.grid_dim.0 <= 1024); // Capped at 1024
        assert!(config.shared_mem_bytes > 0);
    }
}