diff --git a/crates/ml-explainability/build.rs b/crates/ml-explainability/build.rs new file mode 100644 index 000000000..51c1765d0 --- /dev/null +++ b/crates/ml-explainability/build.rs @@ -0,0 +1,102 @@ +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"); + + // Detect GPU architecture from env or default to sm_80 (matches crates/ml/build.rs). + 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; + } + }; + + let kernel_name = "ig_kernels.cu"; + if !compile_kernel(&nvcc, kernel_dir, kernel_name, &arch, &out_dir) { + panic!("nvcc failed to compile {kernel_name}"); + } + eprintln!(" Precompiled 1/1 CUDA kernels ({arch}) — f32"); +} + +/// Compile a single .cu kernel file to a .cubin via nvcc. +fn compile_kernel( + nvcc: &Path, + kernel_dir: &Path, + kernel_name: &str, + arch: &str, + out_dir: &Path, +) -> bool { + 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()); + + // Compile with nvcc. + let status = Command::new(nvcc) + .args([ + "-cubin", + &format!("-arch={arch}"), + "-O3", + "--ftz=true", + "--fmad=true", + "--prec-div=true", + "--prec-sqrt=true", + "-o", cubin_path.to_str().unwrap(), + kernel_path.to_str().unwrap(), + ]) + .status(); + + match status { + Ok(s) if s.success() => { + eprintln!(" Compiled {kernel_name} -> {cubin_name} ({arch})"); + true + } + Ok(s) => { + eprintln!(" FAILED: {kernel_name} (exit={})", s.code().unwrap_or(-1)); + false + } + Err(e) => { + eprintln!(" FAILED: {kernel_name} (nvcc error: {e})"); + false + } + } +} + +/// Find nvcc: prefer $CUDA_HOME/bin/nvcc, then check common paths / $PATH. +fn find_nvcc() -> Option { + if let Ok(home) = std::env::var("CUDA_HOME") { + let nvcc = PathBuf::from(home).join("bin/nvcc"); + if nvcc.exists() { + return Some(nvcc); + } + } + + for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] { + let p = PathBuf::from(path); + if p.exists() { + return Some(p); + } + } + + match Command::new("nvcc").arg("--version").output() { + Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")), + _ => None, + } +} diff --git a/crates/ml-explainability/src/ig_kernels.cu b/crates/ml-explainability/src/ig_kernels.cu new file mode 100644 index 000000000..c13b8c4e2 --- /dev/null +++ b/crates/ml-explainability/src/ig_kernels.cu @@ -0,0 +1,64 @@ +// ═══════════════════════════════════════════════════════════════════════════ +// CUDA kernels for GPU-resident Integrated Gradients (IG). +// +// Sundararajan et al. (2017). Attributions are obtained by integrating the +// model's gradient along the straight-line path baseline -> input: +// +// attribution_i = (x_i - baseline_i) * (1/M) * sum_{k=0..M-1} d F/d x_i +// evaluated at the k-th +// interpolation point. +// +// These kernels handle the two per-step GPU primitives: +// +// 1. `interpolate_input` : builds x(alpha) = baseline + alpha * (input - baseline) +// 2. `perturb_dimension` : copies a buffer and perturbs a single coordinate +// (used to build x_plus / x_minus for central +// finite differences). +// +// All kernels use a 1-D grid-stride layout with 256 threads per block. +// Launch config: blocks = ceil(n / 256), block_dim = 256, shared_mem = 0. +// Deterministic — no atomicAdd, no warp-level reductions, no randomness. +// ═══════════════════════════════════════════════════════════════════════════ + +// Build: output[i] = baseline[i] + alpha * diff[i] (element-wise) +// +// Launch: grid_dim = (ceil(n / 256), 1, 1), block_dim = (256, 1, 1). +extern "C" __global__ +void interpolate_input( + const float* __restrict__ baseline, + const float* __restrict__ diff, + float* __restrict__ output, + float alpha, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + output[i] = baseline[i] + alpha * diff[i]; + } +} + +// Copy src to dst, then apply dst[dim] += delta. +// +// The full copy is performed on every thread (rather than a targeted single +// write) so dst is always a complete, coherent snapshot of src with exactly +// one perturbed coordinate. This avoids needing a separate cudaMemcpy before +// the perturbation and keeps the operation a single kernel launch per probe. +// +// Launch: grid_dim = (ceil(n / 256), 1, 1), block_dim = (256, 1, 1). +extern "C" __global__ +void perturb_dimension( + const float* __restrict__ src, + float* __restrict__ dst, + int dim, + float delta, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float val = src[i]; + if (i == dim) { + val += delta; + } + dst[i] = val; + } +} diff --git a/crates/ml-explainability/src/integrated_gradients.rs b/crates/ml-explainability/src/integrated_gradients.rs index 3df4b469f..3093c78eb 100644 --- a/crates/ml-explainability/src/integrated_gradients.rs +++ b/crates/ml-explainability/src/integrated_gradients.rs @@ -17,10 +17,12 @@ //! # GPU Implementation //! //! On CUDA builds, all computation stays GPU-resident: -//! - Interpolated inputs are constructed on GPU via a fused CUDA kernel. +//! - Interpolated inputs are constructed on GPU via the `interpolate_input` kernel. //! - Forward passes operate on `GpuTensor` (GPU-in, GPU-out). -//! - Finite-difference gradient approximation perturbs tensors on GPU. -//! - Only the final attribution scores are read back to CPU. +//! - Finite-difference gradient approximation perturbs tensors on GPU via the +//! `perturb_dimension` kernel. +//! - Only the scalar outputs of the forward function and the final attribution +//! vector are read back to CPU. //! //! On non-CUDA builds, a CPU fallback with `&[f32]` is provided. @@ -28,9 +30,70 @@ use std::collections::HashMap; use ml_core::MLError; +#[cfg(feature = "cuda")] +use std::sync::Arc; + +#[cfg(feature = "cuda")] +use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; + #[cfg(feature = "cuda")] use ml_core::cuda_autograd::GpuTensor; +/// Precompiled cubin containing the IG kernels (`interpolate_input`, +/// `perturb_dimension`). Built by `build.rs` from `src/ig_kernels.cu`. +#[cfg(feature = "cuda")] +static IG_KERNELS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ig_kernels.cubin")); + +/// Number of threads per block for the 1-D IG kernels. +#[cfg(feature = "cuda")] +const IG_BLOCK_DIM: u32 = 256; + +/// Epsilon used for central finite-difference gradient estimation. +#[cfg(feature = "cuda")] +const IG_FD_EPS: f32 = 1e-4_f32; + +/// Same epsilon for the CPU path. +const IG_FD_EPS_CPU: f32 = 1e-4_f32; + +/// Build a 1-D launch configuration covering `n` elements with 256-thread blocks. +#[cfg(feature = "cuda")] +fn elem_cfg(n: usize) -> LaunchConfig { + let threads = IG_BLOCK_DIM; + let blocks = ((n as u32) + threads - 1) / threads; + LaunchConfig { + grid_dim: (blocks.max(1), 1, 1), + block_dim: (threads, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Bundle of compiled IG kernel handles, loaded once per `compute_gpu` call. +#[cfg(feature = "cuda")] +struct IgKernels { + interpolate_fn: CudaFunction, + perturb_fn: CudaFunction, +} + +#[cfg(feature = "cuda")] +impl IgKernels { + fn load(stream: &Arc) -> Result { + let module = stream + .context() + .load_cubin(IG_KERNELS_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("IG cubin load: {e}")))?; + let interpolate_fn = module + .load_function("interpolate_input") + .map_err(|e| MLError::ModelError(format!("interpolate_input: {e}")))?; + let perturb_fn = module + .load_function("perturb_dimension") + .map_err(|e| MLError::ModelError(format!("perturb_dimension: {e}")))?; + Ok(Self { + interpolate_fn, + perturb_fn, + }) + } +} + /// Integrated Gradients explainability method. /// /// Computes per-feature attribution scores by integrating model gradients @@ -59,7 +122,8 @@ impl IntegratedGradients { /// The forward function operates entirely on GPU tensors (`GpuTensor` in, /// `GpuTensor` out). Interpolated inputs are built on GPU, gradient /// approximation (central finite differences) runs on GPU, and only the - /// final attribution vector is read back to CPU. + /// scalar forward outputs plus the final attribution vector cross the + /// device boundary. /// /// # Arguments /// * `forward_fn` - GPU-resident forward pass: takes a 1-D `GpuTensor` of @@ -73,14 +137,14 @@ impl IntegratedGradients { /// # Errors /// Returns `MLError` if dimensions mismatch or if any GPU operation fails. #[cfg(feature = "cuda")] - #[allow(unsafe_code, unused_variables)] // Stub: forward_fn unused until cubins wired. + #[allow(unsafe_code)] // Raw CUDA kernel launches are inherently unsafe. pub fn compute_gpu( &self, forward_fn: &dyn Fn(&GpuTensor) -> Result, input: &[f32], baseline: Option<&[f32]>, feature_names: &[String], - stream: &std::sync::Arc, + stream: &Arc, ) -> Result, MLError> { let num_elements = input.len(); @@ -101,24 +165,124 @@ impl IntegratedGradients { }); } - // diff = input - baseline (computed on host, uploaded once) + // diff = input - baseline (computed on host, uploaded once). let diff: Vec = input .iter() .zip(baseline_data.iter()) .map(|(a, b)| a - b) .collect(); - // Upload baseline and diff to GPU once - // TODO: implement GPU IG kernel (interpolate + perturb cubins via build.rs) - let _d_baseline = GpuTensor::from_host(baseline_data, vec![num_elements], stream)?; - let _d_diff = GpuTensor::from_host(&diff, vec![num_elements], stream)?; + // Load kernels once per compute call. + let kernels = IgKernels::load(stream)?; - // TODO: implement GPU IG kernel — compile interpolate + perturb cubins - // via build.rs, load them here, and run the IG loop on device. - // For now, return an error so callers fall back to `compute()` (CPU). - Err(MLError::ModelError( - "IG GPU kernels not available: cubins not yet wired".to_owned(), - )) + // Upload baseline and diff to GPU once. + let d_baseline = GpuTensor::from_host(baseline_data, vec![num_elements], stream)?; + let d_diff = GpuTensor::from_host(&diff, vec![num_elements], stream)?; + + // Allocate reusable GPU scratch buffers (freed when this fn returns). + let mut d_interpolated = GpuTensor::zeros(&[num_elements], stream)?; + let mut d_x_plus = GpuTensor::zeros(&[num_elements], stream)?; + let mut d_x_minus = GpuTensor::zeros(&[num_elements], stream)?; + + let n_i32 = i32::try_from(num_elements).map_err(|_| { + MLError::ModelError(format!( + "IG compute_gpu: num_elements={num_elements} exceeds i32 range" + )) + })?; + let cfg = elem_cfg(num_elements); + + let divisor = if self.num_steps > 1 { + (self.num_steps - 1) as f64 + } else { + 1.0 + }; + + let mut accumulated_grads = vec![0.0_f64; num_elements]; + + for step in 0..self.num_steps { + let alpha_f32 = (step as f64 / divisor) as f32; + + // interpolated = baseline + alpha * diff + // SAFETY: Kernel arguments match `interpolate_input`'s signature: + // (const float*, const float*, float*, float, int). + // All buffers are valid CudaSlice of `num_elements` elements. + unsafe { + stream + .launch_builder(&kernels.interpolate_fn) + .arg(d_baseline.cuda_data()) + .arg(d_diff.cuda_data()) + .arg(d_interpolated.data_mut()) + .arg(&alpha_f32) + .arg(&n_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("interpolate_input: {e}")))?; + } + + for feat_idx in 0..num_elements { + let dim_i32 = i32::try_from(feat_idx).map_err(|_| { + MLError::ModelError(format!( + "IG compute_gpu: feature index {feat_idx} exceeds i32 range" + )) + })?; + let pos_eps = IG_FD_EPS; + let neg_eps = -IG_FD_EPS; + + // x_plus = interpolated, then x_plus[feat_idx] += eps + // SAFETY: Kernel arguments match `perturb_dimension`'s signature: + // (const float*, float*, int, float, int). + unsafe { + stream + .launch_builder(&kernels.perturb_fn) + .arg(d_interpolated.cuda_data()) + .arg(d_x_plus.data_mut()) + .arg(&dim_i32) + .arg(&pos_eps) + .arg(&n_i32) + .launch(cfg) + .map_err(|e| { + MLError::ModelError(format!("perturb_dimension(+eps): {e}")) + })?; + } + + // x_minus = interpolated, then x_minus[feat_idx] -= eps + // SAFETY: same signature as above. + unsafe { + stream + .launch_builder(&kernels.perturb_fn) + .arg(d_interpolated.cuda_data()) + .arg(d_x_minus.data_mut()) + .arg(&dim_i32) + .arg(&neg_eps) + .arg(&n_i32) + .launch(cfg) + .map_err(|e| { + MLError::ModelError(format!("perturb_dimension(-eps): {e}")) + })?; + } + + // GPU-resident forward passes; returns scalar GpuTensor of shape [1]. + let y_plus = forward_fn(&d_x_plus)?; + let y_minus = forward_fn(&d_x_minus)?; + let f_plus = y_plus.to_scalar(stream)?; + let f_minus = y_minus.to_scalar(stream)?; + + let grad = f64::from(f_plus - f_minus) / (2.0 * f64::from(IG_FD_EPS)); + if let Some(acc) = accumulated_grads.get_mut(feat_idx) { + *acc += grad; + } + } + } + + let n_steps = self.num_steps as f64; + let mut result = HashMap::with_capacity(feature_names.len()); + for (i, name) in feature_names.iter().enumerate() { + let avg_grad = accumulated_grads.get(i).copied().unwrap_or(0.0) / n_steps; + let d = diff.get(i).copied().unwrap_or(0.0_f32); + let attribution = avg_grad * f64::from(d); + result.insert(name.clone(), attribution); + } + + Ok(result) } /// CPU fallback for feature attribution computation. @@ -163,7 +327,7 @@ impl IntegratedGradients { 1.0 }; - let eps = 1e-4_f32; + let eps = IG_FD_EPS_CPU; let mut accumulated_grads = vec![0.0_f64; num_elements]; for step in 0..self.num_steps { @@ -207,46 +371,6 @@ impl IntegratedGradients { } } -// ── CUDA kernels for GPU-resident integrated gradients ────────────────── -// TODO: implement GPU IG kernel — move to .cu file, compile via build.rs -#[cfg(feature = "cuda")] -#[allow(dead_code)] -const _IG_CUDA_SRC: &str = r#" -// Compute interpolated = baseline + alpha * diff (element-wise) -extern "C" __global__ -void interpolate_input( - const float* __restrict__ baseline, - const float* __restrict__ diff, - float* __restrict__ output, - float alpha, - int n) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - output[i] = baseline[i] + alpha * diff[i]; - } -} - -// Copy src to dst, then perturb dst[dim] += delta -extern "C" __global__ -void perturb_dimension( - const float* __restrict__ src, - float* __restrict__ dst, - int dim, - float delta, - int n) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - float val = src[i]; - if (i == dim) { - val += delta; - } - dst[i] = val; - } -} -"#; - #[cfg(test)] mod tests { use super::*; @@ -350,4 +474,101 @@ mod tests { let ig = IntegratedGradients::new(50); assert_eq!(ig.num_steps, 50); } + + // ── GPU smoke test ──────────────────────────────────────────────────── + // + // Exercises the full `compute_gpu` pipeline on a tiny linear model whose + // forward pass is expressible with GpuTensor primitives (element-wise + // multiply + mean reduction gives a scalar output). Verifies: + // 1. Cubins load and both kernels launch without error. + // 2. Completeness axiom holds within 1% (same tolerance as CPU). + // 3. GPU attributions match the CPU path within 5%. + // + // `#[ignore]` because it requires a CUDA runtime; run with + // `cargo test -p ml-explainability --all-features -- --ignored --nocapture`. + #[cfg(feature = "cuda")] + #[test] + #[ignore = "requires CUDA runtime"] + fn test_ig_compute_gpu_linear_model() { + use ml_core::cuda_autograd::GpuTensor; + + let stream = cudarc::driver::CudaContext::new(0) + .expect("CUDA context") + .new_stream() + .expect("fork stream"); + + // Linear weights on-device: f(x) = sum_i (w_i * x_i) = numel * mean(w * x). + // For numel=4 and weights w, mean(w*x) * 4 == w . x. + // We construct f as a closure that returns a [1]-tensor whose value is + // numel * mean_all(w * x). mean_all returns an f32 scalar. + let weights_host = [2.0_f32, -1.0, 0.5, 3.0]; + let num_features = weights_host.len(); + let d_weights = GpuTensor::from_host( + &weights_host, + vec![num_features], + &stream, + ) + .expect("upload weights"); + + let stream_ref = stream.clone(); + let weights_ref = d_weights.clone(); + let num_features_f32 = num_features as f32; + + let gpu_forward = move |x: &GpuTensor| -> Result { + let wx = weights_ref.mul(x, &stream_ref)?; + let mean = wx.mean_all(&stream_ref)?; + let dot = mean * num_features_f32; + // Return a [1] GpuTensor holding the scalar. + GpuTensor::from_host(&[dot], vec![1], &stream_ref) + }; + + let input = [1.0_f32, -0.5, 0.3, 2.0]; + let feature_names: Vec = + (0..num_features).map(|i| format!("feature_{}", i)).collect(); + + let ig = IntegratedGradients::new(50); + let gpu_attributions = ig + .compute_gpu(&gpu_forward, &input, None, &feature_names, &stream) + .expect("GPU IG should succeed"); + + assert_eq!(gpu_attributions.len(), num_features); + + // Completeness: sum == F(input) - F(baseline). + let f_input: f32 = weights_host + .iter() + .zip(input.iter()) + .map(|(w, x)| w * x) + .sum(); + let expected_diff = f64::from(f_input); + let sum_attr: f64 = gpu_attributions.values().sum(); + let rel_err = if expected_diff.abs() > 1e-10 { + (sum_attr - expected_diff).abs() / expected_diff.abs() + } else { + (sum_attr - expected_diff).abs() + }; + assert!( + rel_err < 0.01, + "GPU completeness violated: sum={:.6}, F(x)-F(b)={:.6}, rel_err={:.4}", + sum_attr, expected_diff, rel_err, + ); + + // Cross-check against CPU path. + let cpu_forward = |x: &[f32]| -> Result { + Ok(weights_host.iter().zip(x.iter()).map(|(w, v)| w * v).sum()) + }; + let cpu_attributions = ig + .compute(&cpu_forward, &input, None, &feature_names) + .expect("CPU IG should succeed"); + + for name in &feature_names { + let gpu_v = gpu_attributions.get(name).copied().unwrap_or(0.0); + let cpu_v = cpu_attributions.get(name).copied().unwrap_or(0.0); + let denom = cpu_v.abs().max(1e-6); + let rel = (gpu_v - cpu_v).abs() / denom; + assert!( + rel < 0.05, + "{name}: gpu={gpu_v:.6} vs cpu={cpu_v:.6} rel_err={rel:.4}", + ); + } + } }