952302149e
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>
65 lines
2.7 KiB
Plaintext
65 lines
2.7 KiB
Plaintext
// ═══════════════════════════════════════════════════════════════════════════
|
|
// 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;
|
|
}
|
|
}
|