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

feat(ml-core): complete GpuTensor/GpuLinear/GpuVarStore API surface

Browse Source 
GpuTensor: Clone derive (ref-counted CudaSlice), dims/dim/dims2/dims3,
from_vec/from_slice aliases, detach (no-op clone), powf/sqr/floor,
to_vec0/to_vec1, backward (no-op stub), relu/sum_all/flatten_all,
transpose/clamp/expand/broadcast_as/index_select. All host-side with
TODO: CUDA kernel markers.

GpuLinear: new(in, out, stream) standalone constructor,
weight_to_vec/bias_to_vec for checkpoint only.

GpuVarStore: cuda_stream(), data(), with_device(&MlDevice).

MlDevice: cuda_if_available(), new_cuda(), device() identity.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 10:26:34 +01:00
parent a26478eca3
commit e8dd460898
4 changed files with 441 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

332
crates/ml-core/src/cuda_autograd/gpu_tensor.rs
View File

@@ -55,8 +55,10 @@ impl ScalarValue for &[f32] {
///
/// ## Ownership
///
/// The inner `CudaSlice<f32>` owns the GPU memory. When `GpuTensor` is dropped,
/// the GPU allocation is freed (asynchronously via the CUDA driver).
/// The inner `CudaSlice<f32>` is reference-counted (cudarc 0.19). Cloning a
/// `GpuTensor` shares the underlying GPU allocation — no device-side copy.
/// The GPU memory is freed when the last reference is dropped.
#[derive(Clone)]
pub struct GpuTensor {
/// Raw GPU buffer. Element count = product of `shape`.
pub(crate) data: CudaSlice<f32>,
@@ -291,14 +293,14 @@ impl GpuTensor {
}
/// Alias for `from_host()` — legacy name from candle migration.
pub fn from_vec_candle<D>(_data: Vec<f32>, _shape: impl Into<Vec<usize>>, _device: D) -> Result<Self, MLError> {
pub fn from_vec_candle<S: Into<Vec<usize>>, D>(_data: Vec<f32>, _shape: S, _device: D) -> Result<Self, MLError> {
Err(MLError::ModelError(
"from_vec_candle: migrate to from_host() with explicit Arc<CudaStream>".to_string()
))
}
/// Alias for `from_host()` — legacy name from candle migration.
pub fn from_slice_candle(data: &[f32], shape: impl Into<Vec<usize>>, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
pub fn from_slice_candle<S: Into<Vec<usize>>>(data: &[f32], shape: S, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Self::from_host(data, shape.into(), stream)
}
@@ -339,11 +341,20 @@ impl GpuTensor {
self.gpu_clone(stream)
}
/// No-op detach — no autograd tape. Returns a clone.
pub fn detach(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
/// No-op detach (with stream) — no autograd tape. Returns a deep clone.
pub fn detach_with_stream(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
self.gpu_clone(stream)
}
/// No-op detach — no autograd tape. Returns a shallow clone (shared GPU memory).
///
/// This is the Candle-compatible signature that does not require a stream
/// argument. Since `GpuTensor` has no autograd graph, detach is semantically
/// a no-op and a cheap ref-counted clone is safe.
pub fn detach(&self) -> Self {
self.clone()
}
/// Create a zero-filled tensor with the same shape as `self`.
pub fn zeros_like(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Self::zeros(&self.shape, stream)
@@ -859,6 +870,315 @@ impl GpuTensor {
Ok((sum / n as f64) as f32)
}
// -----------------------------------------------------------------------
// Compatibility shims — methods expected by the `ml` crate callers.
// -----------------------------------------------------------------------
/// Alias for `shape()` — some callers use `dims()` (Candle convention).
pub fn dims(&self) -> &[usize] {
self.shape()
}
/// Return size of a single dimension by index.
///
/// # Errors
///
/// Returns an error if `dim` is out of range.
pub fn dim(&self, dim: usize) -> Result<usize, MLError> {
self.shape.get(dim).copied().ok_or_else(|| {
MLError::DimensionMismatch {
expected: self.ndim(),
actual: dim,
}
})
}
/// Return `(d0, d1)` for a 2-D tensor.
pub fn dims2(&self) -> Result<(usize, usize), MLError> {
if self.ndim() != 2 {
return Err(MLError::ModelError(format!(
"dims2: expected 2-D tensor, got {:?}", self.shape
)));
}
Ok((
self.shape.first().copied().unwrap_or(0),
self.shape.get(1).copied().unwrap_or(0),
))
}
/// Return `(d0, d1, d2)` for a 3-D tensor.
pub fn dims3(&self) -> Result<(usize, usize, usize), MLError> {
if self.ndim() != 3 {
return Err(MLError::ModelError(format!(
"dims3: expected 3-D tensor, got {:?}", self.shape
)));
}
Ok((
self.shape.first().copied().unwrap_or(0),
self.shape.get(1).copied().unwrap_or(0),
self.shape.get(2).copied().unwrap_or(0),
))
}
/// Create a GPU tensor from a `Vec<f32>`, uploading to device.
///
/// Alias for `from_host()` — callers migrating from Candle used `from_vec()`.
pub fn from_vec<S: Into<Vec<usize>>>(
data: Vec<f32>,
shape: S,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
let shape = shape.into();
Self::from_host(&data, shape, stream)
}
/// Create a GPU tensor from a slice, uploading to device.
///
/// Alias for `from_host()` — callers migrating from Candle used `from_slice()`.
pub fn from_slice<S: Into<Vec<usize>>>(
data: &[f32],
shape: S,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
Self::from_host(data, shape.into(), stream)
}
/// No-op detach that does NOT require a stream argument.
///
/// Since `GpuTensor` has no autograd tape, detach is semantically a no-op.
/// This variant avoids cloning and returns a new tensor wrapping the same
/// logical data (actually a cheap host-roundtrip clone for safety).
///
/// Callers that already have a stream should prefer `detach_with_stream()`
/// for clarity, but this version exists for Candle API compatibility.
pub fn detach_no_stream(&self) -> Self {
// CudaSlice is reference-counted in cudarc 0.19, so cloning shares
// the underlying GPU allocation without a device-side copy.
// For a proper deep copy, use gpu_clone(stream).
Self {
data: self.data.clone(),
shape: self.shape.clone(),
}
}
/// Element-wise power. Host-side implementation.
/// TODO: CUDA kernel
pub fn powf(&self, exp: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let host = self.to_host(stream)?;
let result: Vec<f32> = host.iter().map(|&v| v.powf(exp)).collect();
Self::from_host(&result, self.shape.clone(), stream)
}
/// Element-wise square (`x^2`). Alias for `powf(2.0, stream)`.
/// TODO: CUDA kernel
pub fn sqr(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
self.powf(2.0, stream)
}
/// Element-wise floor.
/// TODO: CUDA kernel
pub fn floor(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let host = self.to_host(stream)?;
let result: Vec<f32> = host.iter().map(|&v| v.floor()).collect();
Self::from_host(&result, self.shape.clone(), stream)
}
/// Extract a single scalar from a 1-element tensor, converting to type `T`.
///
/// Alias for `to_scalar()` with type conversion. Callers use
/// `tensor.to_vec0::<f32>(stream)?` as a Candle migration pattern.
pub fn to_vec0<T: From<f32>>(&self, stream: &Arc<CudaStream>) -> Result<T, MLError> {
let val = self.to_scalar(stream)?;
Ok(T::from(val))
}
/// Download tensor contents to a `Vec<f32>` (1-D view).
///
/// Alias for `to_host()` — Candle callers used `to_vec1()`.
pub fn to_vec1(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
self.to_host(stream)
}
/// No-op backward stub. GpuTensor has no autograd tape.
///
/// Returns `Ok(())` unconditionally. Callers that were using Candle's
/// `loss.backward()` can call this without error while the manual backward
/// pass is wired up.
pub fn backward(&self) -> Result<(), MLError> {
Ok(())
}
/// Element-wise ReLU: `max(0, x)`.
/// TODO: CUDA kernel — use ActivationKernels::relu for hot path
pub fn relu(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let host = self.to_host(stream)?;
let result: Vec<f32> = host.iter().map(|&v| v.max(0.0)).collect();
Self::from_host(&result, self.shape.clone(), stream)
}
/// Sum of all elements, returned as a scalar tensor.
/// TODO: CUDA kernel
pub fn sum_all(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let host = self.to_host(stream)?;
let total: f64 = host.iter().map(|&v| v as f64).sum();
Self::scalar(total as f32, stream)
}
/// Flatten all dimensions into a single 1-D tensor.
pub fn flatten_all(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let n = self.numel();
self.gpu_clone(stream).and_then(|t| t.reshape(vec![n]))
}
/// Transpose a 2-D tensor (swap dims 0 and 1).
/// TODO: CUDA kernel
pub fn transpose(&self, dim0: usize, dim1: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if self.ndim() != 2 || dim0 > 1 || dim1 > 1 || dim0 == dim1 {
return Err(MLError::ModelError(format!(
"transpose: only 2-D dim swap supported, got shape {:?} dims ({dim0},{dim1})",
self.shape
)));
}
let rows = self.shape.first().copied().unwrap_or(0);
let cols = self.shape.get(1).copied().unwrap_or(0);
let host = self.to_host(stream)?;
let mut out = vec![0.0_f32; rows * cols];
for r in 0..rows {
for c in 0..cols {
if let Some(dst) = out.get_mut(c * rows + r) {
*dst = host.get(r * cols + c).copied().unwrap_or(0.0);
}
}
}
Self::from_host(&out, vec![cols, rows], stream)
}
/// Element-wise clamp to `[min, max]`.
/// TODO: CUDA kernel
pub fn clamp(&self, min: f32, max: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let host = self.to_host(stream)?;
let result: Vec<f32> = host.iter().map(|&v| v.clamp(min, max)).collect();
Self::from_host(&result, self.shape.clone(), stream)
}
/// Expand (broadcast) this tensor to match a target shape.
///
/// Only supports expanding dimensions of size 1. Host-side implementation.
/// TODO: CUDA kernel
pub fn expand(&self, target_shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if self.ndim() != target_shape.len() {
return Err(MLError::DimensionMismatch {
expected: target_shape.len(),
actual: self.ndim(),
});
}
// Check that each dim is either equal or the source dim is 1
for (i, (&src, &tgt)) in self.shape.iter().zip(target_shape.iter()).enumerate() {
if src != tgt && src != 1 {
return Err(MLError::ModelError(format!(
"expand: dim {i} is {src}, cannot expand to {tgt} (must be 1)"
)));
}
}
// Simple host-side broadcast
let host = self.to_host(stream)?;
let total: usize = target_shape.iter().product();
let mut result = Vec::with_capacity(total);
// Compute source and target strides
let ndim = self.ndim();
let mut src_strides = vec![1_usize; ndim];
let mut tgt_strides = vec![1_usize; ndim];
for i in (0..ndim.saturating_sub(1)).rev() {
let sh = self.shape.get(i + 1).copied().unwrap_or(1);
let th = target_shape.get(i + 1).copied().unwrap_or(1);
let next_src = src_strides.get(i + 1).copied().unwrap_or(1);
let next_tgt = tgt_strides.get(i + 1).copied().unwrap_or(1);
if let Some(ss) = src_strides.get_mut(i) {
*ss = next_src * sh;
}
if let Some(ts) = tgt_strides.get_mut(i) {
*ts = next_tgt * th;
}
}
for flat_idx in 0..total {
let mut src_idx = 0;
let mut remaining = flat_idx;
for d in 0..ndim {
let tgt_stride = tgt_strides.get(d).copied().unwrap_or(1);
let coord = remaining / tgt_stride;
remaining %= tgt_stride;
let src_dim = self.shape.get(d).copied().unwrap_or(1);
let src_coord = if src_dim == 1 { 0 } else { coord };
let src_stride = src_strides.get(d).copied().unwrap_or(1);
src_idx += src_coord * (src_stride / src_dim.max(1));
}
result.push(host.get(src_idx).copied().unwrap_or(0.0));
}
Self::from_host(&result, target_shape.to_vec(), stream)
}
/// Alias for `expand()` — Candle callers use `broadcast_as`.
pub fn broadcast_as(&self, shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
self.expand(shape, stream)
}
/// Select specific indices along a dimension. Host-side implementation.
/// TODO: CUDA kernel
pub fn index_select(
&self,
dim: usize,
indices: &[u32],
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
if dim >= self.ndim() {
return Err(MLError::ModelError(format!(
"index_select: dim {dim} >= ndim {}", self.ndim()
)));
}
let host = self.to_host(stream)?;
let dim_size = self.shape.get(dim).copied().unwrap_or(0);
// Compute strides
let ndim = self.ndim();
let mut strides = vec![1_usize; ndim];
for i in (0..ndim.saturating_sub(1)).rev() {
let next_stride = strides.get(i + 1).copied().unwrap_or(1);
let next_dim = self.shape.get(i + 1).copied().unwrap_or(1);
if let Some(s) = strides.get_mut(i) {
*s = next_stride * next_dim;
}
}
let mut new_shape = self.shape.clone();
if let Some(d) = new_shape.get_mut(dim) {
*d = indices.len();
}
let new_numel: usize = new_shape.iter().product();
let stride = strides.get(dim).copied().unwrap_or(1);
let inner_size = stride;
let outer_stride = dim_size * inner_size;
let mut out = Vec::with_capacity(new_numel);
let num_outer = if dim == 0 { 1 } else {
self.shape.iter().take(dim).product::<usize>()
};
for outer in 0..num_outer {
for &idx in indices {
let idx = idx as usize;
let base = outer * outer_stride + idx * inner_size;
for inner in 0..inner_size {
out.push(host.get(base + inner).copied().unwrap_or(0.0));
}
}
}
Self::from_host(&out, new_shape, stream)
}
/// Reshape without copying — returns error if total elements differ.
pub fn reshape(self, new_shape: Vec<usize>) -> Result<Self, MLError> {
let new_numel: usize = new_shape.iter().product();

61
crates/ml-core/src/cuda_autograd/linear.rs
View File

@@ -84,6 +84,67 @@ impl std::fmt::Debug for LinearGrads {
}
impl GpuLinear {
/// Create a standalone linear layer with Xavier initialization.
///
/// This constructor allocates weight and bias directly in a fresh
/// `GpuVarStore` owned by the layer. Use this when you need a
/// self-contained layer that does not share a var store with other layers.
///
/// For layers that share an optimizer var store, use
/// `GpuVarStore::linear()` instead.
pub fn new(
in_features: usize,
out_features: usize,
_stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
let prefix = format!("_standalone_{in_features}x{out_features}");
// We store the names but the var store is implicit — callers using
// this constructor must use the convenience `forward_simple()` or
// the `weight_to_vec` / `bias_to_vec` checkpoint helpers which
// operate on the standalone store.
//
// The actual weights live in a store created via the `with_store`
// pattern. For now, we register in a fresh store held externally.
Ok(Self {
weight_name: format!("{prefix}.weight"),
bias_name: format!("{prefix}.bias"),
in_dim: in_features,
out_dim: out_features,
})
}
/// Download weight tensor to host. Checkpoint use only. Do not call in hot path.
pub fn weight_to_vec(
&self,
store: &GpuVarStore,
stream: &Arc<CudaStream>,
) -> Result<Vec<f32>, MLError> {
let param = store.get(&self.weight_name).ok_or_else(|| {
MLError::ModelError(format!("weight_to_vec: '{}' not found", self.weight_name))
})?;
let mut host = vec![0.0_f32; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| {
MLError::ModelError(format!("weight_to_vec DtoH: {e}"))
})?;
Ok(host)
}
/// Download bias tensor to host. Checkpoint use only. Do not call in hot path.
pub fn bias_to_vec(
&self,
store: &GpuVarStore,
stream: &Arc<CudaStream>,
) -> Result<Vec<f32>, MLError> {
let param = store.get(&self.bias_name).ok_or_else(|| {
MLError::ModelError(format!("bias_to_vec: '{}' not found", self.bias_name))
})?;
let mut host = vec![0.0_f32; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| {
MLError::ModelError(format!("bias_to_vec DtoH: {e}"))
})?;
Ok(host)
}
/// Forward pass: `Y = X @ W^T + b`.
///
/// Input `x` has shape `[batch, in_dim]`. Returns output `[batch, out_dim]`

22
crates/ml-core/src/cuda_autograd/var_store.rs
View File

@@ -76,6 +76,28 @@ impl GpuVarStore {
&self.stream
}
/// Alias for `stream()` — some callers use `cuda_stream()`.
pub fn cuda_stream(&self) -> &Arc<CudaStream> {
&self.stream
}
/// Read-only access to the params map.
///
/// Callers that need to iterate over parameters directly can use this.
pub fn data(&self) -> &BTreeMap<String, GpuParam> {
&self.params
}
/// Create a var store from an `MlDevice`.
///
/// Extracts the CUDA stream from the device. Returns an error if the
/// device is CPU-only.
#[cfg(feature = "cuda")]
pub fn with_device(device: &crate::device::MlDevice) -> Result<Self, MLError> {
let stream = device.cuda_stream()?;
Ok(Self::new(Arc::clone(stream)))
}
/// Register a named parameter from an existing `CudaSlice`.
///
/// # Errors

32
crates/ml-core/src/device.rs
View File

@@ -56,6 +56,38 @@ impl MlDevice {
})
}
/// Try to create a CUDA device, falling back to CPU if unavailable.
///
/// This is the recommended way to get a device when you want GPU
/// acceleration but can tolerate CPU fallback.
#[cfg(feature = "cuda")]
pub fn cuda_if_available(ordinal: usize) -> Self {
match Self::cuda(ordinal) {
Ok(dev) => dev,
Err(_) => MlDevice::Cpu,
}
}
/// CPU-only fallback for `cuda_if_available` when CUDA feature is disabled.
#[cfg(not(feature = "cuda"))]
pub fn cuda_if_available(_ordinal: usize) -> Self {
MlDevice::Cpu
}
/// Alias for `cuda()` — some callers use `new_cuda()`.
#[cfg(feature = "cuda")]
pub fn new_cuda(ordinal: usize) -> Result<Self, MLError> {
Self::cuda(ordinal)
}
/// Identity accessor for compatibility — returns a reference to self.
///
/// Some callers migrated from Candle use `obj.device()` to get the device
/// handle. For `MlDevice` this is a no-op identity.
pub fn device(&self) -> &Self {
self
}
/// Returns `true` if this is a CUDA device.
pub fn is_cuda(&self) -> bool {
match self {