rustytorch/crates/training/rtx-flash-metal-attention/src/kernels.rs

//! Metal kernel compilation and pipeline management

use crate::config::FlashAttentionConfig;
use crate::error::{FlashError, FlashResult};
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_foundation::NSString;
use objc2_metal::{MTLComputePipelineState, MTLCreateSystemDefaultDevice, MTLDevice, MTLLibrary};
use tracing::{debug, info};

/// MSL source code for Flash Attention forward kernel (f32)
const FORWARD_KERNEL_SOURCE: &str = include_str!("../metal/flash_attention_forward.metal");

/// MSL source code for Flash Attention backward dQ kernel (f32)
const BACKWARD_DQ_KERNEL_SOURCE: &str = include_str!("../metal/flash_attention_backward_dq.metal");

/// MSL source code for Flash Attention backward dKV kernel (f32)
const BACKWARD_DKV_KERNEL_SOURCE: &str =
    include_str!("../metal/flash_attention_backward_dkv.metal");

/// MSL source code for Flash Attention forward kernel (f16)
const FORWARD_KERNEL_SOURCE_F16: &str = include_str!("../metal/flash_attention_forward_f16.metal");

/// MSL source code for Flash Attention backward dQ kernel (f16)
const BACKWARD_DQ_KERNEL_SOURCE_F16: &str =
    include_str!("../metal/flash_attention_backward_dq_f16.metal");

/// MSL source code for Flash Attention backward dKV kernel (f16)
const BACKWARD_DKV_KERNEL_SOURCE_F16: &str =
    include_str!("../metal/flash_attention_backward_dkv_f16.metal");

/// Get the default Metal device
pub fn get_default_device() -> FlashResult<Retained<ProtocolObject<dyn MTLDevice>>> {
    let device = MTLCreateSystemDefaultDevice();
    device.ok_or_else(|| FlashError::device("No Metal device available"))
}

/// Preprocess MSL source to inject configuration values
fn preprocess_source(source: &str, config: &FlashAttentionConfig) -> String {
    // Replace preprocessor macros with actual values
    // This is a simple string replacement approach that works without
    // needing the complex NSMutableDictionary API
    source
        .replace(
            "#define BLOCK_Q 64",
            &format!("#define BLOCK_Q {}", config.block_q),
        )
        .replace(
            "#define BLOCK_KV 64",
            &format!("#define BLOCK_KV {}", config.block_kv),
        )
        .replace(
            "#define HEAD_DIM 128",
            &format!("#define HEAD_DIM {}", config.max_head_dim),
        )
}

/// Compile MSL source code into a library
fn compile_library(
    device: &ProtocolObject<dyn MTLDevice>,
    source: &str,
    config: &FlashAttentionConfig,
) -> FlashResult<Retained<ProtocolObject<dyn MTLLibrary>>> {
    // Preprocess source with configuration values
    let processed_source = preprocess_source(source, config);

    // Compile the source
    let source_ns = NSString::from_str(&processed_source);
    let library = device
        .newLibraryWithSource_options_error(&source_ns, None)
        .map_err(|e| FlashError::shader(format!("Compilation failed: {:?}", e)))?;

    debug!("Successfully compiled Metal library");
    Ok(library)
}

/// Create a compute pipeline for a kernel function
fn create_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    library: &ProtocolObject<dyn MTLLibrary>,
    function_name: &str,
) -> FlashResult<Retained<ProtocolObject<dyn MTLComputePipelineState>>> {
    let func_name = NSString::from_str(function_name);
    let function = library
        .newFunctionWithName(&func_name)
        .ok_or_else(|| FlashError::pipeline(format!("Function '{}' not found", function_name)))?;

    let pipeline = device
        .newComputePipelineStateWithFunction_error(&function)
        .map_err(|e| FlashError::pipeline(format!("Pipeline creation failed: {:?}", e)))?;

    debug!("Created pipeline for function: {}", function_name);
    Ok(pipeline)
}

/// Compile all Flash Attention pipelines
///
/// Returns (forward_pipeline, backward_dq_pipeline, backward_dkv_pipeline)
pub fn compile_pipelines(
    device: &ProtocolObject<dyn MTLDevice>,
    config: &FlashAttentionConfig,
) -> FlashResult<(
    Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    Retained<ProtocolObject<dyn MTLComputePipelineState>>,
)> {
    let precision = if config.use_f16 { "f16" } else { "f32" };
    info!(
        "Compiling Flash Attention kernels ({}) with block_q={}, block_kv={}, max_head_dim={}",
        precision, config.block_q, config.block_kv, config.max_head_dim
    );

    // Select shader sources based on precision
    let (forward_source, forward_name) = if config.use_f16 {
        (FORWARD_KERNEL_SOURCE_F16, "flash_attention_forward_f16")
    } else {
        (FORWARD_KERNEL_SOURCE, "flash_attention_forward")
    };

    let (backward_dq_source, backward_dq_name) = if config.use_f16 {
        (
            BACKWARD_DQ_KERNEL_SOURCE_F16,
            "flash_attention_backward_dq_f16",
        )
    } else {
        (BACKWARD_DQ_KERNEL_SOURCE, "flash_attention_backward_dq")
    };

    let (backward_dkv_source, backward_dkv_name) = if config.use_f16 {
        (
            BACKWARD_DKV_KERNEL_SOURCE_F16,
            "flash_attention_backward_dkv_f16",
        )
    } else {
        (BACKWARD_DKV_KERNEL_SOURCE, "flash_attention_backward_dkv")
    };

    // Compile forward kernel
    let forward_lib = compile_library(device, forward_source, config)?;
    let forward_pipeline = create_pipeline(device, &forward_lib, forward_name)?;

    // Compile backward dQ kernel
    let backward_dq_lib = compile_library(device, backward_dq_source, config)?;
    let backward_dq_pipeline = create_pipeline(device, &backward_dq_lib, backward_dq_name)?;

    // Compile backward dKV kernel
    let backward_dkv_lib = compile_library(device, backward_dkv_source, config)?;
    let backward_dkv_pipeline = create_pipeline(device, &backward_dkv_lib, backward_dkv_name)?;

    info!(
        "Successfully compiled all Flash Attention pipelines ({})",
        precision
    );

    Ok((
        forward_pipeline,
        backward_dq_pipeline,
        backward_dkv_pipeline,
    ))
}

/// Query Metal device capabilities relevant to Flash Attention
pub fn query_device_capabilities(device: &ProtocolObject<dyn MTLDevice>) -> DeviceCapabilities {
    let max_threads_per_threadgroup = device.maxThreadsPerThreadgroup();

    let max_threadgroup_memory = device.maxThreadgroupMemoryLength();

    let unified_memory = device.hasUnifiedMemory();

    DeviceCapabilities {
        max_threads_per_threadgroup: max_threads_per_threadgroup.width,
        max_threadgroup_memory,
        unified_memory,
    }
}

/// Device capabilities for Flash Attention optimization
#[derive(Debug, Clone)]
pub struct DeviceCapabilities {
    /// Maximum threads per threadgroup
    pub max_threads_per_threadgroup: usize,
    /// Maximum threadgroup memory in bytes
    pub max_threadgroup_memory: usize,
    /// Whether device has unified memory (all Apple Silicon does)
    pub unified_memory: bool,
}

impl DeviceCapabilities {
    /// Compute optimal block sizes for the given head dimension
    pub fn optimal_block_sizes(&self, head_dim: usize) -> (usize, usize) {
        // Each thread needs space for Q row, and we share K/V tiles
        // Threadgroup memory usage: BLOCK_Q * head_dim + 2 * BLOCK_KV * head_dim (floats)
        let bytes_per_float = 4;

        // Start with max block size and reduce if needed
        let mut block_q = 64.min(self.max_threads_per_threadgroup);
        let mut block_kv = 64;

        loop {
            let threadgroup_mem = (block_q * head_dim + 2 * block_kv * head_dim) * bytes_per_float;
            if threadgroup_mem <= self.max_threadgroup_memory {
                break;
            }
            // Reduce block sizes
            if block_q > 32 {
                block_q /= 2;
            } else if block_kv > 32 {
                block_kv /= 2;
            } else {
                break; // Can't reduce further
            }
        }

        (block_q, block_kv)
    }
}