gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: build.rs precompiles epsilon_greedy_kernel.cu — proof of concept for NVRTC elimination

Browse Source 
Add build.rs that compiles .cu kernels to cubins at cargo build time via nvcc.
Replace OnceLock<Ptx> + compile_ptx_for_device() in gpu_action_selector.rs with
include_bytes! + Ptx::from_binary() loading from the precompiled cubin.

- build.rs gracefully skips when nvcc is not available (non-CUDA builds)
- CUDA_COMPUTE_CAP env var controls target GPU arch (86=RTX 3050, 90=H100)
- common_device_functions.cuh prepended to kernel sources at build time
- Zero runtime compiler invocations for epsilon_greedy kernel

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-26 00:07:16 +01:00
parent 5f2e8930d3
commit a2bf5210ec
2 changed files with 129 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
crates/ml/build.rs
View File

@@ -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<PathBuf> {
// 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,
}
}

36
crates/ml/src/cuda_pipeline/gpu_action_selector.rs
View File

@@ -6,26 +6,17 @@
//! indices. The selector stores an `Arc<CudaStream>`.
//!
//! 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<Result<Ptx, String>> = OnceLock::new();
fn compile_kernel_ptx(context: &CudaContext) -> Result<Ptx, String> {
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<CudaStream>, max_batch_size: usize, seed: u64) -> Result<Self, MLError> {
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::<u32>(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]