diff --git a/crates/ml/build.rs b/crates/ml/build.rs index c2bce4fe2..24a9666fc 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -1,3 +1,117 @@ +use std::path::{Path, PathBuf}; +use std::process::Command; + fn main() { println!("cargo:rerun-if-changed=build.rs"); + + // Only compile CUDA kernels when the cuda feature is enabled + if std::env::var("CARGO_FEATURE_CUDA").is_err() { + return; + } + + let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); + let kernel_dir = Path::new("src/cuda_pipeline"); + let common_header = kernel_dir.join("common_device_functions.cuh"); + + // Detect GPU architecture from env or default to sm_80 + let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string()); + let arch = format!("sm_{cuda_compute_cap}"); + + // Check if nvcc is available -- gracefully skip if not (non-CUDA builds) + let nvcc_path = find_nvcc(); + let nvcc = match nvcc_path { + Some(p) => p, + None => { + eprintln!(" warning: nvcc not found, skipping CUDA kernel precompilation"); + eprintln!(" Install CUDA toolkit or set CUDA_HOME for GPU builds"); + return; + } + }; + + // Proof-of-concept: compile epsilon_greedy_kernel.cu (no #define dependencies) + let simple_kernels = [ + "epsilon_greedy_kernel.cu", + ]; + + // Read common header once + let common_src = std::fs::read_to_string(&common_header) + .unwrap_or_else(|e| panic!("Failed to read {}: {e}", common_header.display())); + + println!("cargo:rerun-if-changed={}", common_header.display()); + + for kernel_name in &simple_kernels { + let kernel_path = kernel_dir.join(kernel_name); + let cubin_name = kernel_name.replace(".cu", ".cubin"); + let cubin_path = out_dir.join(&cubin_name); + + println!("cargo:rerun-if-changed={}", kernel_path.display()); + + let kernel_src = std::fs::read_to_string(&kernel_path) + .unwrap_or_else(|e| panic!("Failed to read {}: {e}", kernel_path.display())); + + // Compose source: common header + kernel + // The common header requires STATE_DIM, MARKET_DIM, PORTFOLIO_DIM to be defined. + // For kernels that DON'T use these symbols, we inject dummy defines so the + // header compiles without errors. These kernels don't reference the values. + let defines = "\ + #define STATE_DIM 72\n\ + #define MARKET_DIM 42\n\ + #define PORTFOLIO_DIM 8\n"; + let full_source = format!("{defines}{common_src}\n{kernel_src}"); + + // Write composed source to temp file + let tmp_src = out_dir.join(format!("_{kernel_name}")); + std::fs::write(&tmp_src, &full_source).unwrap(); + + // Compile with nvcc + let status = Command::new(&nvcc) + .args([ + "-cubin", + &format!("-arch={arch}"), + "-O3", + "--use_fast_math", + "--ftz=true", + "--fmad=true", + "-o", cubin_path.to_str().unwrap(), + tmp_src.to_str().unwrap(), + ]) + .status(); + + match status { + Ok(s) if s.success() => { + eprintln!(" Compiled {kernel_name} -> {cubin_name} ({arch})"); + } + Ok(s) => { + panic!("nvcc failed to compile {kernel_name} (exit={})", s.code().unwrap_or(-1)); + } + Err(e) => { + panic!("Failed to execute nvcc: {e}"); + } + } + } +} + +/// Find nvcc: prefer $CUDA_HOME/bin/nvcc, then check PATH +fn find_nvcc() -> Option { + // Try $CUDA_HOME/bin/nvcc first + if let Ok(home) = std::env::var("CUDA_HOME") { + let nvcc = PathBuf::from(home).join("bin/nvcc"); + if nvcc.exists() { + return Some(nvcc); + } + } + + // Try common CUDA paths + for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] { + let p = PathBuf::from(path); + if p.exists() { + return Some(p); + } + } + + // Check if nvcc is in PATH + match Command::new("nvcc").arg("--version").output() { + Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")), + _ => None, + } } diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 268f01969..7da74b516 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -6,26 +6,17 @@ //! indices. The selector stores an `Arc`. //! //! Pure cudarc implementation with no Candle dependency. +//! Kernel loaded from build-time precompiled cubin (zero runtime nvcc). -use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use tracing::info; use crate::MLError; -static EPSILON_GREEDY_PTX: OnceLock> = OnceLock::new(); - -fn compile_kernel_ptx(context: &CudaContext) -> Result { - let defines = "\ - #define STATE_DIM 72\n\ - #define MARKET_DIM 42\n\ - #define PORTFOLIO_DIM 8\n"; - let common_src = include_str!("common_device_functions.cuh"); - let kernel_src = include_str!("epsilon_greedy_kernel.cu"); - let full_source = format!("{defines}{common_src}\n{kernel_src}"); - crate::cuda_pipeline::compile_ptx_for_device(&full_source, context) -} +/// Precompiled epsilon_greedy_kernel cubin, embedded at compile time by build.rs. +static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin")); /// GPU-fused epsilon-greedy action selector -- pure cudarc API. pub struct GpuActionSelector { @@ -52,9 +43,8 @@ pub struct GpuActionSelector { impl GpuActionSelector { pub fn new(stream: Arc, max_batch_size: usize, seed: u64) -> Result { let context = stream.context(); - let ptx_result = EPSILON_GREEDY_PTX.get_or_init(|| compile_kernel_ptx(&context)); - let ptx = ptx_result.as_ref().map_err(|e| MLError::ModelError(format!("epsilon_greedy PTX: {e}")))?; - let module = context.load_module(ptx.clone()).map_err(|e| MLError::ModelError(format!("epsilon_greedy module load: {e}")))?; + let ptx = Ptx::from_binary(EPSILON_GREEDY_CUBIN.to_vec()); + let module = context.load_module(ptx).map_err(|e| MLError::ModelError(format!("load epsilon_greedy cubin: {e}")))?; let kernel_func = module.load_function("epsilon_greedy_select").map_err(|e| MLError::ModelError(format!("epsilon_greedy function load: {e}")))?; let routed_kernel_func = module.load_function("epsilon_greedy_routed").map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed function load: {e}")))?; let branching_kernel_func = module.load_function("branching_action_select").map_err(|e| MLError::ModelError(format!("branching_action_select function load: {e}")))?; @@ -68,7 +58,7 @@ impl GpuActionSelector { } let mut rng_states = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?; stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| MLError::ModelError(format!("upload rng_states: {e}")))?; - info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, kernel compiled"); + info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, precompiled cubin loaded"); Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, actions_buf, fill_mask_buf, max_batch_size, stream, q_gap_threshold: 0.0, bonus_exposure_ptr: 0, bonus_order_ptr: 0, bonus_urgency_ptr: 0, bonus_exposure_buf: None, bonus_order_buf: None, bonus_urgency_buf: None, @@ -227,11 +217,15 @@ mod tests { } #[test] - fn test_ptx_compilation() { + fn test_cubin_loads() { let stream = cuda_stream(); let context = stream.context(); - let result = compile_kernel_ptx(&context); - if let Err(ref e) = result { panic!("PTX compilation failed: {e}"); } + let ptx = Ptx::from_binary(EPSILON_GREEDY_CUBIN.to_vec()); + let module = context.load_module(ptx).expect("load precompiled cubin"); + module.load_function("epsilon_greedy_select").expect("load epsilon_greedy_select"); + module.load_function("epsilon_greedy_routed").expect("load epsilon_greedy_routed"); + module.load_function("branching_action_select").expect("load branching_action_select"); + module.load_function("batch_route_exposure_to_factored").expect("load batch_route_exposure_to_factored"); } #[test]