feat(explainability): GPU-resident Integrated Gradients kernel

Implements the CUDA kernels (`interpolate_input`, `perturb_dimension`)
and wires `compute_gpu` in `integrated_gradients.rs` to run the full
IG algorithm on-device. Replaces the stub that previously returned
`MLError::ModelError("IG GPU kernels not available: cubins not yet
wired")` and fell back to CPU.

Algorithm: for each of `num_steps` interpolation points along the
baseline->input path, compute central-finite-difference gradients for
all `num_features` dimensions via two GPU forward passes per feature.
Single cubin (`ig_kernels.cubin`) compiled via build.rs (following the
crates/ml/build.rs pattern — nvcc, `-arch=sm_\${CUDA_COMPUTE_CAP}`, O3,
f32) and embedded with `include_bytes!`. The `forward_fn` consumers
operate on `GpuTensor`, so no dtoh round-trips inside the inner loop —
only the scalar `[1]` tensor output of each forward pass is pulled to
host (via `to_scalar`) per gradient sample. The three scratch buffers
(`interpolated`, `x_plus`, `x_minus`) are allocated once outside the
step loop and reused across all steps and features. No atomicAdd —
the kernels are trivial 1-D element-wise writes.

Tests: existing CPU tests pass unchanged. Added GPU smoke test
`test_ig_compute_gpu_linear_model` (gated `#[cfg(feature = \"cuda\")]`
+ `#[ignore]`) that builds a linear model as a `GpuTensor`-native
forward (elementwise mul + mean), verifies the completeness axiom
within 1%, and cross-checks GPU attributions against the CPU path
within 5%. Passes locally on RTX 3050 Ti (sm_86).

Removes 3 TODO markers and the embedded `_IG_CUDA_SRC` const that
were awaiting this work, along with the `#[allow(unused_variables)]`
stub attribute on `compute_gpu`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 08:32:09 +02:00
parent d9fee6ef8d
commit 952302149e
3 changed files with 445 additions and 58 deletions

View File

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

View File

@@ -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;
}
}

View File

@@ -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<CudaStream>) -> Result<Self, MLError> {
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<GpuTensor, MLError>,
input: &[f32],
baseline: Option<&[f32]>,
feature_names: &[String],
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
stream: &Arc<CudaStream>,
) -> Result<HashMap<String, f64>, 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<f32> = 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<f32> 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<GpuTensor, MLError> {
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<String> =
(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<f32, MLError> {
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}",
);
}
}
}