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:
102
crates/ml-explainability/build.rs
Normal file
102
crates/ml-explainability/build.rs
Normal 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,
|
||||
}
|
||||
}
|
||||
64
crates/ml-explainability/src/ig_kernels.cu
Normal file
64
crates/ml-explainability/src/ig_kernels.cu
Normal 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;
|
||||
}
|
||||
}
|
||||
@@ -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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user