diff --git a/crates/ml-core/src/cuda_autograd/elementwise.rs b/crates/ml-core/src/cuda_autograd/elementwise.rs new file mode 100644 index 000000000..3d605f32f --- /dev/null +++ b/crates/ml-core/src/cuda_autograd/elementwise.rs @@ -0,0 +1,751 @@ +//! Element-wise and memory-layout CUDA kernels for `GpuTensor` operations. +//! +//! Replaces host-roundtrip workarounds with GPU-native kernels: +//! +//! - **Binary ops**: add, sub, mul (parameterized via `op` arg) +//! - **Unary ops**: powf, sqr, floor, relu, clamp (parameterized via `op` arg) +//! - **Broadcast ops**: scalar/row broadcast mul, div, expand +//! - **Memory ops**: transpose, narrow (dim-0 DtoD), index_select (gather) +//! +//! All kernels are compiled once via `compile_ptx_for_device()` and cached in +//! a `OnceLock`. Thread-safe, no global mutable state. + +#![allow(unsafe_code)] // CUDA FFI requires unsafe for kernel launches. + +use std::sync::{Arc, OnceLock}; + +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::Ptx; + +use crate::MLError; + +/// Launch config helper for element-wise kernels. +fn elem_cfg(n: usize) -> LaunchConfig { + let threads = 256_u32; + let blocks = ((n as u32) + threads - 1) / threads; + LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (threads, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Compiled element-wise kernels -- cached per CUDA context. +/// +/// Construct once via [`ElementwiseKernels::new`], then reuse for all +/// element-wise operations on tensors sharing the same context. +#[allow(missing_debug_implementations)] +pub struct ElementwiseKernels { + binary_fn: CudaFunction, + unary_fn: CudaFunction, + broadcast_binary_fn: CudaFunction, + broadcast_row_binary_fn: CudaFunction, + expand_fn: CudaFunction, + transpose_2d_fn: CudaFunction, + gather_fn: CudaFunction, + stream: Arc, +} + +impl ElementwiseKernels { + /// Compile all element-wise kernels from a single CUDA source. + pub fn new(stream: &Arc) -> Result { + let context = stream.context(); + let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ELEMENTWISE_CUDA_SRC, &context) + .map_err(|e| MLError::ModelError(format!("elementwise kernel compilation: {e}")))?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("elementwise module load: {e}")) + })?; + + let load = |name: &str| -> Result { + module.load_function(name).map_err(|e| { + MLError::ModelError(format!("elementwise kernel '{name}' load: {e}")) + }) + }; + + Ok(Self { + binary_fn: load("elementwise_binary")?, + unary_fn: load("elementwise_unary")?, + broadcast_binary_fn: load("broadcast_scalar_binary")?, + broadcast_row_binary_fn: load("broadcast_row_binary")?, + expand_fn: load("expand_broadcast")?, + transpose_2d_fn: load("transpose_2d")?, + gather_fn: load("gather_select")?, + stream: Arc::clone(stream), + }) + } + + /// Element-wise binary operation on two same-shape buffers. + /// + /// `op`: 0=add, 1=sub, 2=mul, 3=div + pub fn binary( + &self, + a: &CudaSlice, + b: &CudaSlice, + n: usize, + op: i32, + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("binary alloc: {e}")) + })?; + let n_i32 = n as i32; + let cfg = elem_cfg(n); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* a, const float* b, float* out, int n, int op) + // a and b are valid CudaSlice of at least n elements. + // out is a freshly allocated CudaSlice of n elements. + unsafe { + self.stream + .launch_builder(&self.binary_fn) + .arg(a) + .arg(b) + .arg(&out) + .arg(&n_i32) + .arg(&op) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("elementwise_binary(op={op}): {e}")))?; + } + Ok(out) + } + + /// Element-wise unary operation. + /// + /// `op`: 0=powf(param1), 1=sqr, 2=floor, 3=relu, 4=clamp(param1,param2) + pub fn unary( + &self, + x: &CudaSlice, + n: usize, + op: i32, + param1: f32, + param2: f32, + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("unary alloc: {e}")) + })?; + let n_i32 = n as i32; + let cfg = elem_cfg(n); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* x, float* out, int n, int op, float param1, float param2) + // x is a valid CudaSlice of at least n elements. + // out is a freshly allocated CudaSlice of n elements. + unsafe { + self.stream + .launch_builder(&self.unary_fn) + .arg(x) + .arg(&out) + .arg(&n_i32) + .arg(&op) + .arg(¶m1) + .arg(¶m2) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("elementwise_unary(op={op}): {e}")))?; + } + Ok(out) + } + + /// Broadcast a scalar (1-element buffer) with a tensor via binary op. + /// + /// `op`: 0=mul, 1=div, 2=add, 3=sub + /// If `scalar_is_lhs` is true: out[i] = scalar op tensor[i] + /// Otherwise: out[i] = tensor[i] op scalar + pub fn broadcast_scalar( + &self, + tensor: &CudaSlice, + scalar: &CudaSlice, + n: usize, + op: i32, + scalar_is_lhs: i32, + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("broadcast_scalar alloc: {e}")) + })?; + let n_i32 = n as i32; + let cfg = elem_cfg(n); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* tensor, const float* scalar, float* out, + // int n, int op, int scalar_is_lhs) + // tensor is at least n elements, scalar is at least 1 element. + unsafe { + self.stream + .launch_builder(&self.broadcast_binary_fn) + .arg(tensor) + .arg(scalar) + .arg(&out) + .arg(&n_i32) + .arg(&op) + .arg(&scalar_is_lhs) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("broadcast_scalar(op={op}): {e}")))?; + } + Ok(out) + } + + /// Row-broadcast binary op: [M, N] op [1, N] or [1, N] op [M, N]. + /// + /// `op`: 0=mul, 1=div + /// `row_is_lhs`: if true, out[i] = row[i%cols] op matrix[i]; else matrix[i] op row[i%cols] + pub fn broadcast_row( + &self, + matrix: &CudaSlice, + row: &CudaSlice, + total: usize, + cols: usize, + op: i32, + row_is_lhs: i32, + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(total).map_err(|e| { + MLError::ModelError(format!("broadcast_row alloc: {e}")) + })?; + let total_i32 = total as i32; + let cols_i32 = cols as i32; + let cfg = elem_cfg(total); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* matrix, const float* row, float* out, + // int total, int cols, int op, int row_is_lhs) + // matrix is at least `total` elements, row is at least `cols` elements. + unsafe { + self.stream + .launch_builder(&self.broadcast_row_binary_fn) + .arg(matrix) + .arg(row) + .arg(&out) + .arg(&total_i32) + .arg(&cols_i32) + .arg(&op) + .arg(&row_is_lhs) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("broadcast_row(op={op}): {e}")))?; + } + Ok(out) + } + + /// Expand (broadcast) from source to target using pre-computed stride tables. + /// + /// `src_strides` and `src_shape` are GPU-resident u32 buffers of length `ndim`. + /// `tgt_strides` is a GPU-resident u32 buffer of length `ndim`. + pub fn expand( + &self, + src: &CudaSlice, + src_strides_dev: &CudaSlice, + src_shape_dev: &CudaSlice, + tgt_strides_dev: &CudaSlice, + total: usize, + ndim: usize, + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(total).map_err(|e| { + MLError::ModelError(format!("expand alloc: {e}")) + })?; + let total_i32 = total as i32; + let ndim_i32 = ndim as i32; + let cfg = elem_cfg(total); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* src, const unsigned int* src_strides, + // const unsigned int* src_shape, const unsigned int* tgt_strides, + // float* out, int total, int ndim) + // All buffers are valid with sizes consistent with ndim and total. + unsafe { + self.stream + .launch_builder(&self.expand_fn) + .arg(src) + .arg(src_strides_dev) + .arg(src_shape_dev) + .arg(tgt_strides_dev) + .arg(&out) + .arg(&total_i32) + .arg(&ndim_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("expand_broadcast: {e}")))?; + } + Ok(out) + } + + /// Transpose a 2-D matrix [rows, cols] -> [cols, rows]. + pub fn transpose_2d( + &self, + src: &CudaSlice, + rows: usize, + cols: usize, + ) -> Result, MLError> { + let total = rows * cols; + let out = self.stream.alloc_zeros::(total).map_err(|e| { + MLError::ModelError(format!("transpose alloc: {e}")) + })?; + let total_i32 = total as i32; + let rows_i32 = rows as i32; + let cols_i32 = cols as i32; + let cfg = elem_cfg(total); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* src, float* dst, int total, int rows, int cols) + // src and dst are both valid CudaSlice of `total` elements. + unsafe { + self.stream + .launch_builder(&self.transpose_2d_fn) + .arg(src) + .arg(&out) + .arg(&total_i32) + .arg(&rows_i32) + .arg(&cols_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("transpose_2d: {e}")))?; + } + Ok(out) + } + + /// Gather (index_select) along a dimension. + /// + /// For dim-0: out[i, ...] = src[indices[i], ...] + /// `indices_dev` is a GPU-resident u32 buffer. + pub fn gather( + &self, + src: &CudaSlice, + indices_dev: &CudaSlice, + num_indices: usize, + inner_size: usize, + outer_count: usize, + src_dim_size: usize, + ) -> Result, MLError> { + let total = outer_count * num_indices * inner_size; + let out = self.stream.alloc_zeros::(total).map_err(|e| { + MLError::ModelError(format!("gather alloc: {e}")) + })?; + let total_i32 = total as i32; + let num_indices_i32 = num_indices as i32; + let inner_size_i32 = inner_size as i32; + let src_dim_stride_i32 = (src_dim_size * inner_size) as i32; + let cfg = elem_cfg(total); + + // SAFETY: Kernel arguments match the CUDA function signature exactly: + // (const float* src, const unsigned int* indices, float* out, + // int total, int num_indices, int inner_size, int src_dim_stride) + // src is valid with enough elements, indices has num_indices elements, + // out has total elements allocated. + unsafe { + self.stream + .launch_builder(&self.gather_fn) + .arg(src) + .arg(indices_dev) + .arg(&out) + .arg(&total_i32) + .arg(&num_indices_i32) + .arg(&inner_size_i32) + .arg(&src_dim_stride_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("gather_select: {e}")))?; + } + Ok(out) + } + + /// Reference to the underlying stream. + pub fn stream(&self) -> &Arc { + &self.stream + } +} + +/// Global singleton for compiled elementwise kernels. +/// +/// Lazily compiled on first use. Thread-safe via `OnceLock`. +static KERNELS: OnceLock> = OnceLock::new(); + +/// Get or compile the elementwise kernels for the given stream's device. +/// +/// The kernels are compiled once and cached globally. If compilation fails, +/// the error is cached and returned on every subsequent call. +pub fn get_or_compile(stream: &Arc) -> Result<&'static ElementwiseKernels, MLError> { + let result = KERNELS.get_or_init(|| { + ElementwiseKernels::new(stream).map_err(|e| format!("{e}")) + }); + match result { + Ok(k) => Ok(k), + Err(e) => Err(MLError::ModelError(format!("elementwise kernels: {e}"))), + } +} + +// ── CUDA source ────────────────────────────────────────────────────────── + +const ELEMENTWISE_CUDA_SRC: &str = r#" +// ═══════════════════════════════════════════════════════════════════════ +// Element-wise and memory-layout kernels for GpuTensor. +// +// All kernels use 1-D grid-stride pattern with 256 threads per block. +// ═══════════════════════════════════════════════════════════════════════ + +// ── Binary element-wise: add, sub, mul, div ────────────────────────── + +extern "C" __global__ +void elementwise_binary( + const float* __restrict__ a, + const float* __restrict__ b, + float* __restrict__ out, + int n, + int op +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float va = a[i]; + float vb = b[i]; + float r; + switch (op) { + case 0: r = va + vb; break; // add + case 1: r = va - vb; break; // sub + case 2: r = va * vb; break; // mul + case 3: r = va / vb; break; // div + default: r = 0.0f; break; + } + out[i] = r; + } +} + +// ── Unary element-wise: powf, sqr, floor, relu, clamp ──────────────── + +extern "C" __global__ +void elementwise_unary( + const float* __restrict__ x, + float* __restrict__ out, + int n, + int op, + float param1, + float param2 +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float v = x[i]; + float r; + switch (op) { + case 0: r = powf(v, param1); break; // powf + case 1: r = v * v; break; // sqr + case 2: r = floorf(v); break; // floor + case 3: r = fmaxf(v, 0.0f); break; // relu + case 4: r = fminf(fmaxf(v, param1), param2); break; // clamp + default: r = v; break; + } + out[i] = r; + } +} + +// ── Broadcast scalar binary: tensor op scalar or scalar op tensor ──── + +extern "C" __global__ +void broadcast_scalar_binary( + const float* __restrict__ tensor, + const float* __restrict__ scalar_ptr, + float* __restrict__ out, + int n, + int op, + int scalar_is_lhs +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float s = scalar_ptr[0]; + float t = tensor[i]; + float lhs = scalar_is_lhs ? s : t; + float rhs = scalar_is_lhs ? t : s; + float r; + switch (op) { + case 0: r = lhs * rhs; break; // mul + case 1: r = lhs / rhs; break; // div + case 2: r = lhs + rhs; break; // add + case 3: r = lhs - rhs; break; // sub + default: r = 0.0f; break; + } + out[i] = r; + } +} + +// ── Broadcast row binary: [M,N] op [1,N] or [1,N] op [M,N] ───────── + +extern "C" __global__ +void broadcast_row_binary( + const float* __restrict__ matrix, + const float* __restrict__ row, + float* __restrict__ out, + int total, + int cols, + int op, + int row_is_lhs +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < total) { + int col = i % cols; + float m = matrix[i]; + float rv = row[col]; + float lhs = row_is_lhs ? rv : m; + float rhs = row_is_lhs ? m : rv; + float r; + switch (op) { + case 0: r = lhs * rhs; break; // mul + case 1: r = lhs / rhs; break; // div + default: r = 0.0f; break; + } + out[i] = r; + } +} + +// ── Expand (general N-D broadcast) ────────────────────────────────── + +extern "C" __global__ +void expand_broadcast( + const float* __restrict__ src, + const unsigned int* __restrict__ src_strides, + const unsigned int* __restrict__ src_shape, + const unsigned int* __restrict__ tgt_strides, + float* __restrict__ out, + int total, + int ndim +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < total) { + int src_idx = 0; + int remaining = i; + for (int d = 0; d < ndim; d++) { + int tgt_stride = tgt_strides[d]; + int coord = remaining / tgt_stride; + remaining = remaining % tgt_stride; + int src_dim = src_shape[d]; + int src_coord = (src_dim == 1) ? 0 : coord; + src_idx += src_coord * src_strides[d]; + } + out[i] = src[src_idx]; + } +} + +// ── Transpose 2-D ─────────────────────────────────────────────────── + +extern "C" __global__ +void transpose_2d( + const float* __restrict__ src, + float* __restrict__ dst, + int total, + int rows, + int cols +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < total) { + int r = i / cols; + int c = i % cols; + // src[r, c] -> dst[c, r] + dst[c * rows + r] = src[i]; + } +} + +// ── Gather (index_select along a dimension) ────────────────────────── + +extern "C" __global__ +void gather_select( + const float* __restrict__ src, + const unsigned int* __restrict__ indices, + float* __restrict__ out, + int total, + int num_indices, + int inner_size, + int src_dim_stride +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < total) { + // Decompose flat output index into (outer, idx_pos, inner) + int out_idx_inner = i % inner_size; + int tmp = i / inner_size; + int idx_pos = tmp % num_indices; + int outer = tmp / num_indices; + + unsigned int src_idx = indices[idx_pos]; + int src_flat = outer * src_dim_stride + (int)src_idx * inner_size + out_idx_inner; + out[i] = src[src_flat]; + } +} +"#; + +#[cfg(test)] +mod tests { + use super::*; + + fn make_stream() -> Arc { + cudarc::driver::CudaContext::new(0) + .expect("CUDA required") + .new_stream() + .expect("fork stream") + } + + #[test] + fn test_binary_add() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let a_host = vec![1.0_f32, 2.0, 3.0, 4.0]; + let b_host = vec![10.0_f32, 20.0, 30.0, 40.0]; + let mut a = stream.alloc_zeros::(4).unwrap(); + let mut b = stream.alloc_zeros::(4).unwrap(); + stream.memcpy_htod(&a_host, &mut a).unwrap(); + stream.memcpy_htod(&b_host, &mut b).unwrap(); + + let out = kernels.binary(&a, &b, 4, 0).unwrap(); // add + let mut host = [0.0_f32; 4]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - 11.0).abs() < 1e-5); + assert!((host[1] - 22.0).abs() < 1e-5); + assert!((host[2] - 33.0).abs() < 1e-5); + assert!((host[3] - 44.0).abs() < 1e-5); + } + + #[test] + fn test_binary_sub() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let a_host = vec![10.0_f32, 20.0, 30.0]; + let b_host = vec![1.0_f32, 2.0, 3.0]; + let mut a = stream.alloc_zeros::(3).unwrap(); + let mut b = stream.alloc_zeros::(3).unwrap(); + stream.memcpy_htod(&a_host, &mut a).unwrap(); + stream.memcpy_htod(&b_host, &mut b).unwrap(); + + let out = kernels.binary(&a, &b, 3, 1).unwrap(); // sub + let mut host = [0.0_f32; 3]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - 9.0).abs() < 1e-5); + assert!((host[1] - 18.0).abs() < 1e-5); + assert!((host[2] - 27.0).abs() < 1e-5); + } + + #[test] + fn test_binary_mul() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let a_host = vec![2.0_f32, 3.0, 4.0]; + let b_host = vec![5.0_f32, 6.0, 7.0]; + let mut a = stream.alloc_zeros::(3).unwrap(); + let mut b = stream.alloc_zeros::(3).unwrap(); + stream.memcpy_htod(&a_host, &mut a).unwrap(); + stream.memcpy_htod(&b_host, &mut b).unwrap(); + + let out = kernels.binary(&a, &b, 3, 2).unwrap(); // mul + let mut host = [0.0_f32; 3]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - 10.0).abs() < 1e-5); + assert!((host[1] - 18.0).abs() < 1e-5); + assert!((host[2] - 28.0).abs() < 1e-5); + } + + #[test] + fn test_unary_relu() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let x_host = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0]; + let mut x = stream.alloc_zeros::(5).unwrap(); + stream.memcpy_htod(&x_host, &mut x).unwrap(); + + let out = kernels.unary(&x, 5, 3, 0.0, 0.0).unwrap(); // relu + let mut host = [0.0_f32; 5]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0]).abs() < 1e-5); + assert!((host[1]).abs() < 1e-5); + assert!((host[2]).abs() < 1e-5); + assert!((host[3] - 1.0).abs() < 1e-5); + assert!((host[4] - 2.0).abs() < 1e-5); + } + + #[test] + fn test_unary_sqr() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let x_host = vec![1.0_f32, 2.0, 3.0, -4.0]; + let mut x = stream.alloc_zeros::(4).unwrap(); + stream.memcpy_htod(&x_host, &mut x).unwrap(); + + let out = kernels.unary(&x, 4, 1, 0.0, 0.0).unwrap(); // sqr + let mut host = [0.0_f32; 4]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - 1.0).abs() < 1e-5); + assert!((host[1] - 4.0).abs() < 1e-5); + assert!((host[2] - 9.0).abs() < 1e-5); + assert!((host[3] - 16.0).abs() < 1e-5); + } + + #[test] + fn test_unary_clamp() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let x_host = vec![-5.0_f32, -1.0, 0.5, 2.0, 10.0]; + let mut x = stream.alloc_zeros::(5).unwrap(); + stream.memcpy_htod(&x_host, &mut x).unwrap(); + + let out = kernels.unary(&x, 5, 4, -2.0, 3.0).unwrap(); // clamp(-2, 3) + let mut host = [0.0_f32; 5]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - (-2.0)).abs() < 1e-5); + assert!((host[1] - (-1.0)).abs() < 1e-5); + assert!((host[2] - 0.5).abs() < 1e-5); + assert!((host[3] - 2.0).abs() < 1e-5); + assert!((host[4] - 3.0).abs() < 1e-5); + } + + #[test] + fn test_transpose_2d() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + // 2x3 matrix: [[1,2,3],[4,5,6]] + let x_host = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let mut x = stream.alloc_zeros::(6).unwrap(); + stream.memcpy_htod(&x_host, &mut x).unwrap(); + + let out = kernels.transpose_2d(&x, 2, 3).unwrap(); + let mut host = [0.0_f32; 6]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + // Expected: [[1,4],[2,5],[3,6]] = [1,4,2,5,3,6] + assert!((host[0] - 1.0).abs() < 1e-5); + assert!((host[1] - 4.0).abs() < 1e-5); + assert!((host[2] - 2.0).abs() < 1e-5); + assert!((host[3] - 5.0).abs() < 1e-5); + assert!((host[4] - 3.0).abs() < 1e-5); + assert!((host[5] - 6.0).abs() < 1e-5); + } + + #[test] + fn test_broadcast_scalar_mul() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + let tensor_host = vec![1.0_f32, 2.0, 3.0, 4.0]; + let scalar_host = vec![10.0_f32]; + let mut tensor = stream.alloc_zeros::(4).unwrap(); + let mut scalar = stream.alloc_zeros::(1).unwrap(); + stream.memcpy_htod(&tensor_host, &mut tensor).unwrap(); + stream.memcpy_htod(&scalar_host, &mut scalar).unwrap(); + + let out = kernels.broadcast_scalar(&tensor, &scalar, 4, 0, 0).unwrap(); // mul, scalar is rhs + let mut host = [0.0_f32; 4]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + assert!((host[0] - 10.0).abs() < 1e-5); + assert!((host[1] - 20.0).abs() < 1e-5); + assert!((host[2] - 30.0).abs() < 1e-5); + assert!((host[3] - 40.0).abs() < 1e-5); + } + + #[test] + fn test_gather_dim0() { + let stream = make_stream(); + let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); + // 4x3 src: row0=[0,1,2], row1=[3,4,5], row2=[6,7,8], row3=[9,10,11] + let src_host: Vec = (0..12).map(|i| i as f32).collect(); + let mut src = stream.alloc_zeros::(12).unwrap(); + stream.memcpy_htod(&src_host, &mut src).unwrap(); + + // Select rows 2, 0 + let idx_host = vec![2_u32, 0]; + let mut idx = stream.alloc_zeros::(2).unwrap(); + stream.memcpy_htod(&idx_host, &mut idx).unwrap(); + + // dim=0, inner_size=3, outer_count=1, src_dim_size=4 + let out = kernels.gather(&src, &idx, 2, 3, 1, 4).unwrap(); + let mut host = [0.0_f32; 6]; + stream.memcpy_dtoh(&out, &mut host).unwrap(); + // row2: [6,7,8], row0: [0,1,2] + assert!((host[0] - 6.0).abs() < 1e-5); + assert!((host[1] - 7.0).abs() < 1e-5); + assert!((host[2] - 8.0).abs() < 1e-5); + assert!((host[3] - 0.0).abs() < 1e-5); + assert!((host[4] - 1.0).abs() < 1e-5); + assert!((host[5] - 2.0).abs() < 1e-5); + } +} diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index 3539fa6d7..a60b90c44 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -4,6 +4,8 @@ //! wrapper — no autograd graph, no dtype dispatch, no device polymorphism. //! All data is F32 on a single CUDA device. +#![allow(unsafe_code)] // CUDA FFI requires unsafe for DtoD memcpy and kernel launches. + use std::mem::ManuallyDrop; use std::sync::Arc; @@ -320,20 +322,31 @@ impl GpuTensor { // ----------------------------------------------------------------------- // Tensor algebra: element-wise ops, reductions, slicing. - // Most ops download to host, compute, re-upload. Marked for future CUDA kernels. + // All ops are GPU-native — no host downloads except to_host/to_scalar. // ----------------------------------------------------------------------- /// Clone this tensor by copying GPU data to a new allocation. /// - /// Uses host-side roundtrip (same as `clone_gpu_tensor` in linear.rs). - /// TODO: CUDA kernel — use `cuMemcpyDtoDAsync` for zero-host clone. + /// Uses async device-to-device memcpy — zero host download. pub fn gpu_clone(&self, stream: &Arc) -> Result { let n = self.numel(); - let mut host = vec![0.0_f32; n]; - stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| { - MLError::ModelError(format!("gpu_clone DtoH: {e}")) + let dst = stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("gpu_clone alloc: {e}")) })?; - Self::from_host(&host, self.shape.clone(), stream) + let num_bytes = n * std::mem::size_of::(); + let src_ptr = raw_ptr(&self.data, stream); + let dst_ptr = raw_ptr(&dst, stream); + // SAFETY: src and dst are valid device allocations on the same context. + // num_bytes = n * sizeof(f32) does not exceed either allocation. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("gpu_clone DtoD: {e}")))?; + } + Ok(Self { + data: dst, + shape: self.shape.clone(), + }) } /// No-op dtype cast — `GpuTensor` is always F32. Returns a clone. @@ -384,7 +397,8 @@ impl GpuTensor { } /// Element-wise addition. Shapes must match exactly. - /// TODO: CUDA kernel + /// + /// GPU-native: uses `elementwise_binary` kernel (op=0). pub fn add(&self, other: &Self, stream: &Arc) -> Result { if self.shape != other.shape { return Err(MLError::DimensionMismatch { @@ -392,14 +406,14 @@ impl GpuTensor { actual: other.numel(), }); } - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; - let c: Vec = a.iter().zip(b.iter()).map(|(x, y)| x + y).collect(); - Self::from_host(&c, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.binary(&self.data, &other.data, self.numel(), 0)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Element-wise subtraction. Shapes must match exactly. - /// TODO: CUDA kernel + /// + /// GPU-native: uses `elementwise_binary` kernel (op=1). pub fn sub(&self, other: &Self, stream: &Arc) -> Result { if self.shape != other.shape { return Err(MLError::DimensionMismatch { @@ -407,14 +421,14 @@ impl GpuTensor { actual: other.numel(), }); } - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; - let c: Vec = a.iter().zip(b.iter()).map(|(x, y)| x - y).collect(); - Self::from_host(&c, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.binary(&self.data, &other.data, self.numel(), 1)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Element-wise multiply. Shapes must match exactly. - /// TODO: CUDA kernel + /// + /// GPU-native: uses `elementwise_binary` kernel (op=2). pub fn mul(&self, other: &Self, stream: &Arc) -> Result { if self.shape != other.shape { return Err(MLError::DimensionMismatch { @@ -422,10 +436,9 @@ impl GpuTensor { actual: other.numel(), }); } - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; - let c: Vec = a.iter().zip(b.iter()).map(|(x, y)| x * y).collect(); - Self::from_host(&c, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.binary(&self.data, &other.data, self.numel(), 2)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Element-wise multiply with broadcasting. @@ -434,28 +447,24 @@ impl GpuTensor { /// - Same shape: plain element-wise multiply. /// - Scalar (numel=1) * tensor: broadcast the scalar. /// - `[1, N]` * `[M, N]` or `[M, N]` * `[1, N]`: broadcast rows. - /// TODO: CUDA kernel + /// + /// GPU-native: all paths use CUDA kernels with zero host download. pub fn broadcast_mul(&self, other: &Self, stream: &Arc) -> Result { // Fast path: identical shapes if self.shape == other.shape { return self.mul(other, stream); } + let kernels = super::elementwise::get_or_compile(stream)?; // Scalar broadcast if other.numel() == 1 { - let scalar = other.to_host(stream)?; - let s = scalar.first().copied().unwrap_or(0.0); - let a = self.to_host(stream)?; - let c: Vec = a.iter().map(|x| x * s).collect(); - return Self::from_host(&c, self.shape.clone(), stream); + let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 0, 0)?; + return Ok(Self { data: out, shape: self.shape.clone() }); } if self.numel() == 1 { - let scalar = self.to_host(stream)?; - let s = scalar.first().copied().unwrap_or(0.0); - let b = other.to_host(stream)?; - let c: Vec = b.iter().map(|x| x * s).collect(); - return Self::from_host(&c, other.shape.clone(), stream); + let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 0, 1)?; + return Ok(Self { data: out, shape: other.shape.clone() }); } - // Row broadcast: [1, N] * [M, N] + // Row broadcast: [1, N] * [M, N] or [M, N] * [1, N] if self.ndim() == 2 && other.ndim() == 2 { let (m_a, n_a) = ( self.shape.first().copied().unwrap_or(0), @@ -466,25 +475,15 @@ impl GpuTensor { other.shape.get(1).copied().unwrap_or(0), ); if n_a == n_b { - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; if m_a == 1 { - let c: Vec = (0..m_b * n_b) - .map(|i| { - let col = i % n_b; - b.get(i).copied().unwrap_or(0.0) * a.get(col).copied().unwrap_or(0.0) - }) - .collect(); - return Self::from_host(&c, other.shape.clone(), stream); + // [1, N] * [M, N]: self is the row, other is the matrix + let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 0, 0)?; + return Ok(Self { data: out, shape: other.shape.clone() }); } if m_b == 1 { - let c: Vec = (0..m_a * n_a) - .map(|i| { - let col = i % n_a; - a.get(i).copied().unwrap_or(0.0) * b.get(col).copied().unwrap_or(0.0) - }) - .collect(); - return Self::from_host(&c, self.shape.clone(), stream); + // [M, N] * [1, N]: other is the row, self is the matrix + let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 0, 0)?; + return Ok(Self { data: out, shape: self.shape.clone() }); } } } @@ -497,29 +496,25 @@ impl GpuTensor { /// Element-wise divide with broadcasting. /// /// Supports the same broadcast patterns as `broadcast_mul`. - /// TODO: CUDA kernel + /// + /// GPU-native: all paths use CUDA kernels with zero host download. pub fn broadcast_div(&self, other: &Self, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; // Fast path: identical shapes if self.shape == other.shape { - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; - let c: Vec = a.iter().zip(b.iter()).map(|(x, y)| x / y).collect(); - return Self::from_host(&c, self.shape.clone(), stream); + let out = kernels.binary(&self.data, &other.data, self.numel(), 3)?; + return Ok(Self { data: out, shape: self.shape.clone() }); } // Scalar broadcast (divisor is scalar) if other.numel() == 1 { - let scalar = other.to_host(stream)?; - let s = scalar.first().copied().unwrap_or(1.0); - let a = self.to_host(stream)?; - let c: Vec = a.iter().map(|x| x / s).collect(); - return Self::from_host(&c, self.shape.clone(), stream); + // tensor / scalar: scalar_is_lhs=0, op=1 (div) + let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 1, 0)?; + return Ok(Self { data: out, shape: self.shape.clone() }); } if self.numel() == 1 { - let scalar = self.to_host(stream)?; - let s = scalar.first().copied().unwrap_or(0.0); - let b = other.to_host(stream)?; - let c: Vec = b.iter().map(|x| s / x).collect(); - return Self::from_host(&c, other.shape.clone(), stream); + // scalar / tensor: scalar_is_lhs=1, op=1 (div) + let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 1, 1)?; + return Ok(Self { data: out, shape: other.shape.clone() }); } // Row broadcast for 2-D if self.ndim() == 2 && other.ndim() == 2 { @@ -532,27 +527,15 @@ impl GpuTensor { other.shape.get(1).copied().unwrap_or(0), ); if n_a == n_b { - let a = self.to_host(stream)?; - let b = other.to_host(stream)?; if m_b == 1 { - // [M, N] / [1, N] - let c: Vec = (0..m_a * n_a) - .map(|i| { - let col = i % n_a; - a.get(i).copied().unwrap_or(0.0) / b.get(col).copied().unwrap_or(1.0) - }) - .collect(); - return Self::from_host(&c, self.shape.clone(), stream); + // [M, N] / [1, N]: matrix=self, row=other, op=1 (div), row_is_lhs=0 + let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 1, 0)?; + return Ok(Self { data: out, shape: self.shape.clone() }); } if m_a == 1 { - // [1, N] / [M, N] - let c: Vec = (0..m_b * n_b) - .map(|i| { - let col = i % n_b; - a.get(col).copied().unwrap_or(0.0) / b.get(i).copied().unwrap_or(1.0) - }) - .collect(); - return Self::from_host(&c, other.shape.clone(), stream); + // [1, N] / [M, N]: matrix=other, row=self, op=1 (div), row_is_lhs=1 + let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 1, 1)?; + return Ok(Self { data: out, shape: other.shape.clone() }); } } } @@ -690,8 +673,8 @@ impl GpuTensor { /// Slice along a dimension: extract `[start..start+len)` along `dim`. /// - /// Only dim=0 is currently optimized; other dims use host roundtrip. - /// TODO: CUDA kernel for general narrow. + /// For dim=0 (contiguous), uses async DtoD memcpy — zero host download. + /// For other dims, uses the gather kernel on GPU. pub fn narrow( &self, dim: usize, @@ -712,8 +695,6 @@ impl GpuTensor { ))); } - let host = self.to_host(stream)?; - // Compute strides for row-major layout let ndim = self.ndim(); let mut strides = vec![1_usize; ndim]; @@ -730,54 +711,53 @@ impl GpuTensor { *d = len; } let new_numel: usize = new_shape.iter().product(); - let mut out = Vec::with_capacity(new_numel); let stride = strides.get(dim).copied().unwrap_or(1); - // For dim=0 this is a contiguous slice + // For dim=0 this is a contiguous sub-range — use DtoD memcpy if dim == 0 { let offset = start * stride; - let end = offset + len * stride; - for i in offset..end { - out.push(host.get(i).copied().unwrap_or(0.0)); - } - } else { - // General case: iterate over all output elements - let mut idx = vec![0_usize; ndim]; - for _ in 0..new_numel { - // Compute source flat index - let mut src_flat = 0; - for d in 0..ndim { - let coord = if d == dim { - idx.get(d).copied().unwrap_or(0) + start - } else { - idx.get(d).copied().unwrap_or(0) - }; - src_flat += coord * strides.get(d).copied().unwrap_or(1); - } - out.push(host.get(src_flat).copied().unwrap_or(0.0)); - - // Increment multi-index - for d in (0..ndim).rev() { - if let Some(v) = idx.get_mut(d) { - *v += 1; - let limit = new_shape.get(d).copied().unwrap_or(1); - if *v < limit { - break; - } - *v = 0; - } - } + let src_view = self.data.slice(offset..offset + new_numel); + let dst = stream.alloc_zeros::(new_numel).map_err(|e| { + MLError::ModelError(format!("narrow alloc: {e}")) + })?; + let num_bytes = new_numel * std::mem::size_of::(); + let dst_ptr = raw_ptr(&dst, stream); + let (src_ptr, _src_sync) = src_view.device_ptr(stream); + // SAFETY: src_view and dst are valid device allocations on the same context. + // num_bytes does not exceed either allocation. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("narrow DtoD: {e}")))?; } + return Ok(Self { data: dst, shape: new_shape }); } - Self::from_host(&out, new_shape, stream) + // General case: use gather kernel with sequential indices + let indices: Vec = (start..start + len).map(|i| i as u32).collect(); + let mut indices_dev = stream.alloc_zeros::(len).map_err(|e| { + MLError::ModelError(format!("narrow indices alloc: {e}")) + })?; + stream.memcpy_htod(&indices, &mut indices_dev).map_err(|e| { + MLError::ModelError(format!("narrow indices HtoD: {e}")) + })?; + + let inner_size = stride; // elements per index along this dim + let outer_count = if dim == 0 { 1 } else { + self.shape.iter().take(dim).product::() + }; + + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.gather(&self.data, &indices_dev, len, inner_size, outer_count, dim_size)?; + Ok(Self { data: out, shape: new_shape }) } /// Argmax along a dimension. Returns indices as `Vec`. /// /// For a 2-D tensor `[M, N]` with `dim=1`, returns `M` indices in `[0..N)`. - /// TODO: CUDA kernel + /// + /// GPU-native: delegates to `ReductionKernels::argmax` / `argmax_rows`. pub fn argmax(&self, dim: usize, stream: &Arc) -> Result, MLError> { if dim >= self.ndim() { return Err(MLError::ModelError(format!( @@ -785,59 +765,30 @@ impl GpuTensor { self.ndim() ))); } - let host = self.to_host(stream)?; + let reductions = super::reductions::ReductionKernels::new(stream)?; - // For 1-D, return single argmax + // For 1-D, return single argmax via flat kernel if self.ndim() == 1 { - let mut best_idx = 0_u32; - let mut best_val = f32::NEG_INFINITY; - for (i, v) in host.iter().enumerate() { - if *v > best_val { - best_val = *v; - best_idx = i as u32; - } - } - return Ok(vec![best_idx]); + let idx = reductions.argmax(&self.data, self.numel())?; + return Ok(vec![idx]); } - // For 2-D with dim=1: argmax per row + // For 2-D with dim=1: argmax per row — use ReductionKernels::argmax_rows if self.ndim() == 2 && dim == 1 { let rows = self.shape.first().copied().unwrap_or(0); let cols = self.shape.get(1).copied().unwrap_or(0); - let mut indices = Vec::with_capacity(rows); - for r in 0..rows { - let mut best_idx = 0_u32; - let mut best_val = f32::NEG_INFINITY; - for c in 0..cols { - let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY); - if v > best_val { - best_val = v; - best_idx = c as u32; - } - } - indices.push(best_idx); - } - return Ok(indices); + return reductions.argmax_rows(&self.data, rows, cols); } - // For 2-D with dim=0: argmax per column + // For 2-D with dim=0: transpose then argmax per row of transposed if self.ndim() == 2 && dim == 0 { let rows = self.shape.first().copied().unwrap_or(0); let cols = self.shape.get(1).copied().unwrap_or(0); - let mut indices = Vec::with_capacity(cols); - for c in 0..cols { - let mut best_idx = 0_u32; - let mut best_val = f32::NEG_INFINITY; - for r in 0..rows { - let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY); - if v > best_val { - best_val = v; - best_idx = r as u32; - } - } - indices.push(best_idx); - } - return Ok(indices); + let kernels = super::elementwise::get_or_compile(stream)?; + let transposed = kernels.transpose_2d(&self.data, rows, cols)?; + // Now transposed is [cols, rows]. argmax per row of [cols, rows] gives + // per-original-column argmax along original rows (dim=0). + return reductions.argmax_rows(&transposed, cols, rows); } Err(MLError::ModelError(format!( @@ -859,15 +810,16 @@ impl GpuTensor { } /// Mean of all elements. - /// TODO: CUDA kernel + /// + /// GPU-native: delegates to `ReductionKernels::stats().mean`. pub fn mean_all(&self, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let n = host.len(); + let n = self.numel(); if n == 0 { return Err(MLError::ModelError("mean_all: empty tensor".to_string())); } - let sum: f64 = host.iter().map(|&v| v as f64).sum(); - Ok((sum / n as f64) as f32) + let reductions = super::reductions::ReductionKernels::new(stream)?; + let stats = reductions.stats(&self.data, n)?; + Ok(stats.mean) } // ----------------------------------------------------------------------- @@ -961,26 +913,31 @@ impl GpuTensor { } } - /// Element-wise power. Host-side implementation. - /// TODO: CUDA kernel + /// Element-wise power. + /// + /// GPU-native: uses `elementwise_unary` kernel (op=0, param1=exp). pub fn powf(&self, exp: f32, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.powf(exp)).collect(); - Self::from_host(&result, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.unary(&self.data, self.numel(), 0, exp, 0.0)?; + Ok(Self { data: out, shape: self.shape.clone() }) } - /// Element-wise square (`x^2`). Alias for `powf(2.0, stream)`. - /// TODO: CUDA kernel + /// Element-wise square (`x^2`). + /// + /// GPU-native: uses `elementwise_unary` kernel (op=1) for `v*v`. pub fn sqr(&self, stream: &Arc) -> Result { - self.powf(2.0, stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.unary(&self.data, self.numel(), 1, 0.0, 0.0)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Element-wise floor. - /// TODO: CUDA kernel + /// + /// GPU-native: uses `elementwise_unary` kernel (op=2). pub fn floor(&self, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.floor()).collect(); - Self::from_host(&result, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.unary(&self.data, self.numel(), 2, 0.0, 0.0)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Extract a single scalar from a 1-element tensor, converting to type `T`. @@ -1009,19 +966,25 @@ impl GpuTensor { } /// Element-wise ReLU: `max(0, x)`. - /// TODO: CUDA kernel — use ActivationKernels::relu for hot path + /// + /// GPU-native: uses `elementwise_unary` kernel (op=3). pub fn relu(&self, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.max(0.0)).collect(); - Self::from_host(&result, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.unary(&self.data, self.numel(), 3, 0.0, 0.0)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Sum of all elements, returned as a scalar tensor. - /// TODO: CUDA kernel + /// + /// GPU-native: delegates to `ReductionKernels::sum()`. pub fn sum_all(&self, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let total: f64 = host.iter().map(|&v| v as f64).sum(); - Self::scalar(total as f32, stream) + let n = self.numel(); + if n == 0 { + return Self::scalar(0.0, stream); + } + let reductions = super::reductions::ReductionKernels::new(stream)?; + let total = reductions.sum(&self.data, n)?; + Self::scalar(total, stream) } /// Flatten all dimensions into a single 1-D tensor. @@ -1031,7 +994,8 @@ impl GpuTensor { } /// Transpose a 2-D tensor (swap dims 0 and 1). - /// TODO: CUDA kernel + /// + /// GPU-native: uses `transpose_2d` CUDA kernel. pub fn transpose(&self, dim0: usize, dim1: usize, stream: &Arc) -> Result { if self.ndim() != 2 || dim0 > 1 || dim1 > 1 || dim0 == dim1 { return Err(MLError::ModelError(format!( @@ -1041,30 +1005,25 @@ impl GpuTensor { } 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) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.transpose_2d(&self.data, rows, cols)?; + Ok(Self { data: out, shape: vec![cols, rows] }) } /// Element-wise clamp to `[min, max]`. - /// TODO: CUDA kernel + /// + /// GPU-native: uses `elementwise_unary` kernel (op=4, param1=min, param2=max). pub fn clamp(&self, min: f32, max: f32, stream: &Arc) -> Result { - let host = self.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.clamp(min, max)).collect(); - Self::from_host(&result, self.shape.clone(), stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.unary(&self.data, self.numel(), 4, min, max)?; + Ok(Self { data: out, shape: self.shape.clone() }) } /// Expand (broadcast) this tensor to match a target shape. /// - /// Only supports expanding dimensions of size 1. Host-side implementation. - /// TODO: CUDA kernel + /// Only supports expanding dimensions of size 1. + /// + /// GPU-native: uploads stride tables to GPU, runs `expand_broadcast` kernel. pub fn expand(&self, target_shape: &[usize], stream: &Arc) -> Result { if self.ndim() != target_shape.len() { return Err(MLError::DimensionMismatch { @@ -1080,44 +1039,71 @@ impl GpuTensor { ))); } } - // 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]; + let total: usize = target_shape.iter().product(); + + // Compute per-element strides for source and target + let mut src_strides_host = vec![1_u32; ndim]; + let mut tgt_strides_host = vec![1_u32; ndim]; + let mut src_shape_host: Vec = self.shape.iter().map(|&s| s as u32).collect(); + // Pad to at least ndim elements (always should be correct) + while src_shape_host.len() < ndim { + src_shape_host.push(1); + } + 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) { + let sh = self.shape.get(i + 1).copied().unwrap_or(1) as u32; + let th = target_shape.get(i + 1).copied().unwrap_or(1) as u32; + let next_src = src_strides_host.get(i + 1).copied().unwrap_or(1); + let next_tgt = tgt_strides_host.get(i + 1).copied().unwrap_or(1); + if let Some(ss) = src_strides_host.get_mut(i) { *ss = next_src * sh; } - if let Some(ts) = tgt_strides.get_mut(i) { + if let Some(ts) = tgt_strides_host.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)); - } + // For expand, source strides for dims of size 1 must be 0 so the kernel + // reads the same element. But our kernel handles this via src_shape check, + // so we need per-element strides (stride within the source buffer). + // The kernel computes: src_coord = (src_dim==1) ? 0 : coord + // Then: src_idx += src_coord * src_strides[d] + // We need src_strides[d] to be the stride *within the source allocation*, + // i.e. for a source of shape [1, N], src_strides = [N, 1]. + // This is already what we computed above (product of subsequent dims). - Self::from_host(&result, target_shape.to_vec(), stream) + // Upload stride tables to GPU + let mut src_strides_dev = stream.alloc_zeros::(ndim).map_err(|e| { + MLError::ModelError(format!("expand src_strides alloc: {e}")) + })?; + let mut src_shape_dev = stream.alloc_zeros::(ndim).map_err(|e| { + MLError::ModelError(format!("expand src_shape alloc: {e}")) + })?; + let mut tgt_strides_dev = stream.alloc_zeros::(ndim).map_err(|e| { + MLError::ModelError(format!("expand tgt_strides alloc: {e}")) + })?; + stream.memcpy_htod(&src_strides_host, &mut src_strides_dev).map_err(|e| { + MLError::ModelError(format!("expand src_strides HtoD: {e}")) + })?; + stream.memcpy_htod(&src_shape_host, &mut src_shape_dev).map_err(|e| { + MLError::ModelError(format!("expand src_shape HtoD: {e}")) + })?; + stream.memcpy_htod(&tgt_strides_host, &mut tgt_strides_dev).map_err(|e| { + MLError::ModelError(format!("expand tgt_strides HtoD: {e}")) + })?; + + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.expand( + &self.data, + &src_strides_dev, + &src_shape_dev, + &tgt_strides_dev, + total, + ndim, + )?; + Ok(Self { data: out, shape: target_shape.to_vec() }) } /// Alias for `expand()` — Candle callers use `broadcast_as`. @@ -1125,8 +1111,9 @@ impl GpuTensor { self.expand(shape, stream) } - /// Select specific indices along a dimension. Host-side implementation. - /// TODO: CUDA kernel + /// Select specific indices along a dimension. + /// + /// GPU-native: uploads indices to GPU, runs `gather_select` kernel. pub fn index_select( &self, dim: usize, @@ -1138,7 +1125,6 @@ impl GpuTensor { "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 @@ -1156,27 +1142,22 @@ impl GpuTensor { 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 { + let inner_size = strides.get(dim).copied().unwrap_or(1); + let outer_count = if dim == 0 { 1 } else { self.shape.iter().take(dim).product::() }; - 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)); - } - } - } + // Upload indices to GPU + let mut indices_dev = stream.alloc_zeros::(indices.len()).map_err(|e| { + MLError::ModelError(format!("index_select indices alloc: {e}")) + })?; + stream.memcpy_htod(indices, &mut indices_dev).map_err(|e| { + MLError::ModelError(format!("index_select indices HtoD: {e}")) + })?; - Self::from_host(&out, new_shape, stream) + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.gather(&self.data, &indices_dev, indices.len(), inner_size, outer_count, dim_size)?; + Ok(Self { data: out, shape: new_shape }) } /// Reshape without copying — returns error if total elements differ. diff --git a/crates/ml-core/src/cuda_autograd/linear.rs b/crates/ml-core/src/cuda_autograd/linear.rs index 76314f34b..33533acd9 100644 --- a/crates/ml-core/src/cuda_autograd/linear.rs +++ b/crates/ml-core/src/cuda_autograd/linear.rs @@ -405,16 +405,13 @@ fn reduce_sum_axis0( } /// Clone a GpuTensor by copying its data to a new allocation. +/// +/// Uses async device-to-device memcpy — zero host download. pub fn clone_gpu_tensor( src: &GpuTensor, stream: &Arc, ) -> Result { - let n = src.numel(); - let mut host = vec![0.0_f32; n]; - stream.memcpy_dtoh(&src.data, &mut host).map_err(|e| { - MLError::ModelError(format!("clone DtoH: {e}")) - })?; - GpuTensor::from_host(&host, src.shape.clone(), stream) + src.gpu_clone(stream) } #[cfg(test)] diff --git a/crates/ml-core/src/cuda_autograd/mod.rs b/crates/ml-core/src/cuda_autograd/mod.rs index 8a8a588c8..590e027df 100644 --- a/crates/ml-core/src/cuda_autograd/mod.rs +++ b/crates/ml-core/src/cuda_autograd/mod.rs @@ -41,6 +41,7 @@ pub mod loss; pub mod dropout; pub mod layer_norm; pub mod reductions; +pub mod elementwise; pub mod stream_ops; pub use gpu_tensor::GpuTensor; @@ -60,5 +61,6 @@ pub use dropout::GpuDropout; pub use layer_norm::GpuLayerNorm; pub use reductions::ReductionKernels; pub use reductions::Stats as ReductionStats; +pub use elementwise::ElementwiseKernels; pub use stream_ops::StreamTensor; pub use stream_ops::StreamLinear;