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

feat(ml): KAN spline GPU-native forward — eliminate GPU→CPU→GPU roundtrip

Browse Source 
Replace clamped.to_vec1() CPU loop with GPU-native floor/ceil/frac
operations. Removes 3 tensor transfers per forward pass (N floats
down + 3*N up). All index computation now stays on GPU.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-02 17:41:14 +01:00
parent 3aab43614a
commit f8cec1d6f9
1 changed files with 43 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
crates/ml/src/kan/spline.rs
View File

@@ -189,8 +189,8 @@ impl BSplineBasis {
/// GPU evaluation via linear interpolation on pre-computed grid.
///
/// Pulls only N normalized index floats to CPU (not N * num_bases),
/// builds floor/ceil u32 index tensors, then does gather + lerp on GPU.
/// All index computation (floor, ceil, frac) stays on GPU — zero CPU
/// roundtrip. Gather + lerp completes the lookup entirely on device.
fn evaluate_gpu(
&self,
x: &Tensor,
@@ -217,7 +217,7 @@ impl BSplineBasis {
let res_max = (self.grid_resolution - 1) as f32;
// -- Normalize x to [0, grid_resolution-1] on GPU, then pull small vec --
// -- Normalize x to [0, grid_resolution-1] on GPU --
let grid_min_t = Tensor::new(&[self.grid_min], device)
.map_err(|e| MLError::ModelError(format!("BSpline grid_min tensor: {e}")))?;
let inv_range_scaled = Tensor::new(&[res_max / range], device)
@@ -239,44 +239,50 @@ impl BSplineBasis {
.clamp(&zero_t, &max_t)
.map_err(|e| MLError::ModelError(format!("BSpline clamp: {e}")))?;
// Pull clamped floats to CPU -- this is just N floats, not N * num_bases
let clamped_vec: Vec<f32> = clamped.to_vec1().map_err(|e| {
MLError::ModelError(format!("BSpline clamped to_vec1: {e}"))
})?;
// GPU-native floor/ceil/frac — zero CPU roundtrip.
// clamped is shape (n,), range [0, res_max - 1].
let floor_f = clamped
.floor()
.map_err(|e| MLError::ModelError(format!("KAN floor: {e}")))?;
// Build floor/ceil index arrays and fraction array on CPU
let mut floor_indices = vec![0_u32; n * self.num_bases];
let mut ceil_indices = vec![0_u32; n * self.num_bases];
let mut fractions = vec![0.0_f32; n * self.num_bases];
// frac = clamped - floor (stays F32, shape (n,))
let frac_1d = clamped
.sub(&floor_f)
.map_err(|e| MLError::ModelError(format!("KAN frac: {e}")))?;
let res_u32 = (self.grid_resolution - 1) as u32;
for (i, val) in clamped_vec.iter().enumerate() {
let f = val.floor();
let floor_u = (f as u32).min(res_u32.saturating_sub(1));
let ceil_u = (floor_u + 1).min(res_u32);
let frac = val - f;
// ceil = floor + 1, clamped to [0, res_max] so we don't exceed lookup table
let ones_1d = Tensor::ones(floor_f.shape(), DType::F32, device)
.map_err(|e| MLError::ModelError(format!("KAN ones: {e}")))?;
let ceil_max = Tensor::full(res_max, floor_f.shape(), device)
.map_err(|e| MLError::ModelError(format!("KAN ceil_max: {e}")))?;
let ceil_f = floor_f
.add(&ones_1d)
.and_then(|t| t.minimum(&ceil_max))
.map_err(|e| MLError::ModelError(format!("KAN ceil: {e}")))?;
// Broadcast to all num_bases columns
for j in 0..self.num_bases {
if let Some(slot) = floor_indices.get_mut(i * self.num_bases + j) {
*slot = floor_u;
}
if let Some(slot) = ceil_indices.get_mut(i * self.num_bases + j) {
*slot = ceil_u;
}
if let Some(slot) = fractions.get_mut(i * self.num_bases + j) {
*slot = frac;
}
}
}
// Cast to U32 for gather indices, then broadcast (n,) → (n, num_bases)
let floor_1d = floor_f
.to_dtype(DType::U32)
.map_err(|e| MLError::ModelError(format!("KAN floor u32: {e}")))?;
let ceil_1d = ceil_f
.to_dtype(DType::U32)
.map_err(|e| MLError::ModelError(format!("KAN ceil u32: {e}")))?;
// Upload index tensors and fractions to GPU
let floor_t = Tensor::from_vec(floor_indices, &[n, self.num_bases], device)
.map_err(|e| MLError::ModelError(format!("BSpline floor tensor: {e}")))?;
let ceil_t = Tensor::from_vec(ceil_indices, &[n, self.num_bases], device)
.map_err(|e| MLError::ModelError(format!("BSpline ceil tensor: {e}")))?;
let frac_t = Tensor::from_vec(fractions, &[n, self.num_bases], device)
.map_err(|e| MLError::ModelError(format!("BSpline frac tensor: {e}")))?;
let floor_t = floor_1d
.unsqueeze(1)
.and_then(|t| t.broadcast_as(&[n, self.num_bases]))
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("KAN floor broadcast: {e}")))?;
let ceil_t = ceil_1d
.unsqueeze(1)
.and_then(|t| t.broadcast_as(&[n, self.num_bases]))
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("KAN ceil broadcast: {e}")))?;
let frac_t = frac_1d
.unsqueeze(1)
.and_then(|t| t.broadcast_as(&[n, self.num_bases]))
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("KAN frac broadcast: {e}")))?;
// GPU gather: grid is (resolution, num_bases), indices are (n, num_bases)
// gather(dim=0) selects rows per-column