diff --git a/crates/ml-core/src/cuda_autograd/elementwise.rs b/crates/ml-core/src/cuda_autograd/elementwise.rs index 3d605f32f..46699ddf3 100644 --- a/crates/ml-core/src/cuda_autograd/elementwise.rs +++ b/crates/ml-core/src/cuda_autograd/elementwise.rs @@ -110,7 +110,8 @@ impl ElementwiseKernels { /// Element-wise unary operation. /// - /// `op`: 0=powf(param1), 1=sqr, 2=floor, 3=relu, 4=clamp(param1,param2) + /// `op`: 0=powf(param1), 1=sqr, 2=floor, 3=relu, 4=clamp(param1,param2), + /// 5=abs, 6=neg, 7=exp, 8=log, 9=le(param1), 10=affine(v*param1+param2) pub fn unary( &self, x: &CudaSlice, @@ -340,6 +341,38 @@ impl ElementwiseKernels { Ok(out) } + // ── Convenience wrappers for unary ops 5-10 ─────────────────────── + + /// Element-wise absolute value (`|x|`). Op 5. + pub fn abs(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + self.unary(x, n, 5, 0.0, 0.0) + } + + /// Element-wise negation (`-x`). Op 6. + pub fn neg(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + self.unary(x, n, 6, 0.0, 0.0) + } + + /// Element-wise exponential (`e^x`). Op 7. + pub fn exp(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + self.unary(x, n, 7, 0.0, 0.0) + } + + /// Element-wise natural logarithm (`ln(x)`). Op 8. + pub fn log(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + self.unary(x, n, 8, 0.0, 0.0) + } + + /// Element-wise less-or-equal comparison. Returns 1.0 if `x <= threshold`, else 0.0. Op 9. + pub fn le(&self, x: &CudaSlice, threshold: f32, n: usize) -> Result, MLError> { + self.unary(x, n, 9, threshold, 0.0) + } + + /// Element-wise affine transform: `x * mul + add`. Op 10. + pub fn affine(&self, x: &CudaSlice, mul: f32, add: f32, n: usize) -> Result, MLError> { + self.unary(x, n, 10, mul, add) + } + /// Reference to the underlying stream. pub fn stream(&self) -> &Arc { &self.stream @@ -400,7 +433,7 @@ void elementwise_binary( } } -// ── Unary element-wise: powf, sqr, floor, relu, clamp ──────────────── +// ── Unary element-wise: powf, sqr, floor, relu, clamp, abs, neg, exp, log, le, affine ── extern "C" __global__ void elementwise_unary( @@ -421,6 +454,12 @@ void elementwise_unary( 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 + case 5: r = fabsf(v); break; // abs + case 6: r = -v; break; // neg + case 7: r = expf(v); break; // exp + case 8: r = logf(v); break; // log + case 9: r = (v <= param1) ? 1.0f : 0.0f; break; // le (compare) + case 10: r = v * param1 + param2; break; // affine (mul + add) default: r = v; break; } out[i] = r; diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index a60b90c44..200ada1d8 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -1010,6 +1010,97 @@ impl GpuTensor { Ok(Self { data: out, shape: vec![cols, rows] }) } + /// Element-wise absolute value (`|x|`). + /// + /// GPU-native: uses `elementwise_unary` kernel (op=5). + pub fn abs(&self, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.abs(&self.data, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise negation (`-x`). + /// + /// GPU-native: uses `elementwise_unary` kernel (op=6). + pub fn neg(&self, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.neg(&self.data, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise exponential (`e^x`). + /// + /// GPU-native: uses `elementwise_unary` kernel (op=7). + pub fn exp(&self, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.exp(&self.data, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise natural logarithm (`ln(x)`). + /// + /// GPU-native: uses `elementwise_unary` kernel (op=8). + pub fn log(&self, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.log(&self.data, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise less-or-equal comparison. Returns 1.0 where `x <= threshold`, else 0.0. + /// + /// GPU-native: uses `elementwise_unary` kernel (op=9). + pub fn le(&self, threshold: f32, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.le(&self.data, threshold, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise affine transform: `x * mul + add`. + /// + /// GPU-native: uses `elementwise_unary` kernel (op=10). + pub fn affine(&self, mul: f64, add: f64, stream: &Arc) -> Result { + let kernels = super::elementwise::get_or_compile(stream)?; + let out = kernels.affine(&self.data, mul as f32, add as f32, self.numel())?; + Ok(Self { data: out, shape: self.shape.clone() }) + } + + /// Element-wise subtraction with broadcasting. + /// + /// Supports the common cases: + /// - Same shape: plain element-wise sub. + /// - Scalar (numel=1): broadcast the scalar. + /// - `[M, N]` - `[1, N]`: row broadcast. + /// + /// GPU-native: all paths use CUDA kernels with zero host download. + pub fn broadcast_sub(&self, other: &Self, stream: &Arc) -> Result { + // Fast path: identical shapes + if self.shape == other.shape { + return self.sub(other, stream); + } + let kernels = super::elementwise::get_or_compile(stream)?; + // Scalar broadcast: self - scalar + if other.numel() == 1 { + // tensor - scalar: op=3 (sub), scalar_is_lhs=0 + let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 3, 0)?; + return Ok(Self { data: out, shape: self.shape.clone() }); + } + // Scalar broadcast: scalar - tensor + if self.numel() == 1 { + // scalar - tensor: op=3 (sub), scalar_is_lhs=1 + let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 3, 1)?; + return Ok(Self { data: out, shape: other.shape.clone() }); + } + // Same numel: just use element-wise sub + if self.numel() == other.numel() { + let out = kernels.binary(&self.data, &other.data, self.numel(), 1)?; + return Ok(Self { data: out, shape: self.shape.clone() }); + } + Err(MLError::ModelError(format!( + "broadcast_sub: incompatible shapes {:?} vs {:?}", + self.shape, other.shape + ))) + } + /// Element-wise clamp to `[min, max]`. /// /// GPU-native: uses `elementwise_unary` kernel (op=4, param1=min, param2=max). diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index c3e774da7..ab8a8866d 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -25,461 +25,6 @@ use tracing::debug; use super::{Experience, ExposureLevel, FactoredAction, OrderRouter, OrderType, Urgency}; use ml_core::MLError; -/// Host-side affine transform: `x * mul + add` element-wise. -/// TODO: use GpuTensor::affine() when available -fn gpu_affine(t: &GpuTensor, mul: f64, add: f64, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| (v as f64 * mul + add) as f32).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side gather: select elements along dim from indices. -/// TODO: use GpuTensor::gather() when available -fn gpu_gather(t: &GpuTensor, indices: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let t_host = t.to_host(stream)?; - let idx_host = indices.to_host(stream)?; - let t_shape = t.shape(); - let idx_shape = indices.shape(); - - if t_shape.len() != 2 || dim != 1 { - return Err(MLError::ModelError(format!( - "gpu_gather: only 2D dim=1 supported, got shape {:?} dim {dim}", t_shape - ))); - } - let cols = t_shape.get(1).copied().unwrap_or(0); - let out_len = idx_host.len(); - let mut result = Vec::with_capacity(out_len); - for (i, &idx_f) in idx_host.iter().enumerate() { - let row = i / idx_shape.get(1).copied().unwrap_or(1); - let idx = idx_f as usize; - let val = t_host.get(row * cols + idx).copied().unwrap_or(0.0); - result.push(val); - } - GpuTensor::from_host(&result, idx_shape.to_vec(), stream) -} - -/// Host-side broadcast_sub: `a - b` with broadcasting. -/// TODO: use GpuTensor::broadcast_sub() when available -fn gpu_broadcast_sub(a: &GpuTensor, b: &GpuTensor, stream: &Arc) -> Result { - let a_host = a.to_host(stream)?; - let b_host = b.to_host(stream)?; - if b.numel() == 1 { - let s = b_host.first().copied().unwrap_or(0.0); - let c: Vec = a_host.iter().map(|&v| v - s).collect(); - GpuTensor::from_host(&c, a.shape().to_vec(), stream) - } else if a.shape() == b.shape() { - let c: Vec = a_host.iter().zip(b_host.iter()).map(|(x, y)| x - y).collect(); - GpuTensor::from_host(&c, a.shape().to_vec(), stream) - } else { - Err(MLError::ModelError(format!( - "gpu_broadcast_sub: incompatible shapes {:?} vs {:?}", a.shape(), b.shape() - ))) - } -} - -/// Host-side dim(i): get the size of dimension `i`. -fn gpu_dim(t: &GpuTensor, dim: usize) -> Result { - t.shape().get(dim).copied().ok_or_else(|| MLError::ModelError(format!( - "gpu_dim: dim {dim} out of range for shape {:?}", t.shape() - ))) -} - -/// Host-side clamp: element-wise clamp to [min, max]. -fn gpu_clamp(t: &GpuTensor, min: f64, max: f64, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let min_f = min as f32; - let max_f = max as f32; - let result: Vec = host.iter().map(|&v| v.clamp(min_f, max_f)).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side abs: element-wise absolute value. -fn gpu_abs(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.abs()).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side neg: element-wise negation. -fn gpu_neg(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| -v).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side exp: element-wise exponential. -fn gpu_exp(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.exp()).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side log: element-wise natural log. -fn gpu_log(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v.ln()).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side le (less-or-equal): element-wise comparison, returns 0.0/1.0. -fn gpu_le(t: &GpuTensor, threshold: f32, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| if v <= threshold { 1.0 } else { 0.0 }).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side sum along a dimension. For 2D with dim=1, reduces [M,N] -> [M]. -fn gpu_sum(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() == 2 && dim == 1 { - let rows = shape.first().copied().unwrap_or(0); - let cols = shape.get(1).copied().unwrap_or(0); - let mut result = Vec::with_capacity(rows); - for r in 0..rows { - let mut s = 0.0_f32; - for c in 0..cols { - s += host.get(r * cols + c).copied().unwrap_or(0.0); - } - result.push(s); - } - GpuTensor::from_host(&result, vec![rows], stream) - } else if shape.len() == 3 && dim == 2 { - // [B, A, N] -> [B, A] - let b = shape.first().copied().unwrap_or(0); - let a = shape.get(1).copied().unwrap_or(0); - let n = shape.get(2).copied().unwrap_or(0); - let mut result = Vec::with_capacity(b * a); - for bi in 0..b { - for ai in 0..a { - let mut s = 0.0_f32; - for ni in 0..n { - s += host.get(bi * a * n + ai * n + ni).copied().unwrap_or(0.0); - } - result.push(s); - } - } - GpuTensor::from_host(&result, vec![b, a], stream) - } else { - Err(MLError::ModelError(format!( - "gpu_sum: unsupported shape {:?} dim={dim}", shape - ))) - } -} - -/// Host-side sum_keepdim along dim=1 for 2D tensor: [M,N] -> [M,1]. -fn gpu_sum_keepdim(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() == 2 && dim == 1 { - let rows = shape.first().copied().unwrap_or(0); - let cols = shape.get(1).copied().unwrap_or(0); - let mut result = Vec::with_capacity(rows); - for r in 0..rows { - let mut s = 0.0_f32; - for c in 0..cols { - s += host.get(r * cols + c).copied().unwrap_or(0.0); - } - result.push(s); - } - GpuTensor::from_host(&result, vec![rows, 1], stream) - } else { - Err(MLError::ModelError(format!( - "gpu_sum_keepdim: unsupported shape {:?} dim={dim}", shape - ))) - } -} - -/// Host-side max along dim. For 2D with dim=1, reduces [M,N] -> [M]. -fn gpu_max(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() == 2 && dim == 1 { - let rows = shape.first().copied().unwrap_or(0); - let cols = shape.get(1).copied().unwrap_or(0); - let mut result = Vec::with_capacity(rows); - for r in 0..rows { - let mut mx = f32::NEG_INFINITY; - for c in 0..cols { - let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY); - if v > mx { mx = v; } - } - result.push(mx); - } - GpuTensor::from_host(&result, vec![rows], stream) - } else if shape.len() == 1 && dim == 0 { - let mx = host.iter().copied().fold(f32::NEG_INFINITY, f32::max); - GpuTensor::from_host(&[mx], vec![1], stream) - } else { - Err(MLError::ModelError(format!( - "gpu_max: unsupported shape {:?} dim={dim}", shape - ))) - } -} - -/// Host-side min along dim. For 1D with dim=0, returns scalar. -fn gpu_min(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() == 1 && dim == 0 { - let mn = host.iter().copied().fold(f32::INFINITY, f32::min); - GpuTensor::from_host(&[mn], vec![1], stream) - } else { - Err(MLError::ModelError(format!( - "gpu_min: unsupported shape {:?} dim={dim}", shape - ))) - } -} - -/// Host-side argmax_keepdim: argmax along dim, keeping the dimension. -/// For 2D with dim=1: [M,N] -> [M,1]. -fn gpu_argmax_keepdim(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let indices = t.argmax(dim, stream)?; - let n = indices.len(); - let result: Vec = indices.iter().map(|&i| i as f32).collect(); - if t.shape().len() == 2 && dim == 1 { - GpuTensor::from_host(&result, vec![n, 1], stream) - } else { - GpuTensor::from_host(&result, vec![n], stream) - } -} - -/// Host-side broadcast_as: expand tensor to target shape via repetition. -fn gpu_broadcast_as(t: &GpuTensor, target_shape: &[usize], stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let src_shape = t.shape(); - let target_numel: usize = target_shape.iter().product(); - - // Scalar broadcast: repeat the single value - if t.numel() == 1 { - let val = host.first().copied().unwrap_or(0.0); - let result = vec![val; target_numel]; - return GpuTensor::from_host(&result, target_shape.to_vec(), stream); - } - - // [1,N] -> [M,N] or [1,1,N] -> [B,A,N] - if src_shape.len() <= target_shape.len() { - let mut result = Vec::with_capacity(target_numel); - let src_numel = t.numel(); - for i in 0..target_numel { - let src_idx = i % src_numel; - result.push(host.get(src_idx).copied().unwrap_or(0.0)); - } - return GpuTensor::from_host(&result, target_shape.to_vec(), stream); - } - - Err(MLError::ModelError(format!( - "gpu_broadcast_as: cannot broadcast {:?} to {:?}", src_shape, target_shape - ))) -} - -/// Host-side gather along dim=1 for 3D tensors: t[b, indices[b,1,n], n] -> [b,1,n]. -fn gpu_gather_3d(t: &GpuTensor, indices: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let t_host = t.to_host(stream)?; - let idx_host = indices.to_host(stream)?; - let t_shape = t.shape(); - let idx_shape = indices.shape(); - - if t_shape.len() == 3 && dim == 1 { - let b = t_shape.first().copied().unwrap_or(0); - let _a = t_shape.get(1).copied().unwrap_or(0); - let n = t_shape.get(2).copied().unwrap_or(0); - let out_a = idx_shape.get(1).copied().unwrap_or(1); - let out_n = idx_shape.get(2).copied().unwrap_or(n); - let mut result = Vec::with_capacity(b * out_a * out_n); - for bi in 0..b { - for oi in 0..out_a { - for ni in 0..out_n { - let idx_val = idx_host.get(bi * out_a * out_n + oi * out_n + ni).copied().unwrap_or(0.0) as usize; - let val = t_host.get(bi * _a * n + idx_val * n + ni).copied().unwrap_or(0.0); - result.push(val); - } - } - } - GpuTensor::from_host(&result, vec![b, out_a, out_n], stream) - } else { - Err(MLError::ModelError(format!( - "gpu_gather_3d: unsupported shape {:?} dim={dim}", t_shape - ))) - } -} - -/// Host-side squeeze: remove a size-1 dimension (via host roundtrip). -fn gpu_squeeze_host(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let shape = t.shape(); - if dim >= shape.len() || shape.get(dim).copied().unwrap_or(0) != 1 { - return Err(MLError::ModelError(format!( - "gpu_squeeze_host: dim {dim} invalid for shape {:?}", shape - ))); - } - let mut new_shape: Vec = shape.to_vec(); - new_shape.remove(dim); - let host = t.to_host(stream)?; - GpuTensor::from_host(&host, new_shape, stream) -} - -/// Host-side flatten: flatten all dims to 1D. -fn gpu_flatten_all(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - GpuTensor::from_host(&host, vec![host.len()], stream) -} - -/// Host-side sqr: element-wise square. -fn gpu_sqr(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let result: Vec = host.iter().map(|&v| v * v).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), stream) -} - -/// Host-side sum_all: sum all elements, return scalar. -fn gpu_sum_all(t: &GpuTensor, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - Ok(host.iter().sum()) -} - -/// Host-side log_softmax along dim=1 for 2D tensor. -fn todo_log_softmax_fn(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() != 2 || dim != 1 { - return Err(MLError::ModelError(format!( - "todo_log_softmax_fn: only 2D dim=1 supported, got {:?} dim={dim}", shape - ))); - } - let rows = shape.first().copied().unwrap_or(0); - let cols = shape.get(1).copied().unwrap_or(0); - let mut result = Vec::with_capacity(rows * cols); - for r in 0..rows { - let row_start = r * cols; - let mx = (0..cols).map(|c| host.get(row_start + c).copied().unwrap_or(f32::NEG_INFINITY)).fold(f32::NEG_INFINITY, f32::max); - let log_sum_exp = (0..cols).map(|c| (host.get(row_start + c).copied().unwrap_or(0.0) - mx).exp()).sum::().ln() + mx; - for c in 0..cols { - let v = host.get(row_start + c).copied().unwrap_or(0.0); - result.push(v - log_sum_exp); - } - } - GpuTensor::from_host(&result, shape.to_vec(), stream) -} - -/// Host-side softmax along dim=1 for 2D tensor. -fn todo_softmax_fn(t: &GpuTensor, dim: usize, stream: &Arc) -> Result { - let host = t.to_host(stream)?; - let shape = t.shape(); - if shape.len() != 2 || dim != 1 { - return Err(MLError::ModelError(format!( - "todo_softmax_fn: only 2D dim=1 supported, got {:?} dim={dim}", shape - ))); - } - let rows = shape.first().copied().unwrap_or(0); - let cols = shape.get(1).copied().unwrap_or(0); - let mut result = Vec::with_capacity(rows * cols); - for r in 0..rows { - let row_start = r * cols; - let mx = (0..cols).map(|c| host.get(row_start + c).copied().unwrap_or(f32::NEG_INFINITY)).fold(f32::NEG_INFINITY, f32::max); - let exps: Vec = (0..cols).map(|c| (host.get(row_start + c).copied().unwrap_or(0.0) - mx).exp()).collect(); - let sum: f32 = exps.iter().sum(); - for e in &exps { - result.push(if sum > 0.0 { e / sum } else { 1.0 / cols as f32 }); - } - } - GpuTensor::from_host(&result, shape.to_vec(), stream) -} - -/// Host-side leaky_relu: max(alpha * x, x). -fn todo_leaky_relu_fn(t: &GpuTensor, alpha: f64) -> Result { - // We need a stream for to_host / from_host. Get it via a temp context. - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("leaky_relu CUDA ctx: {e}")) - })?; - let stream = ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("leaky_relu CUDA stream: {e}")) - })?; - let host = t.to_host(&stream)?; - let a = alpha as f32; - let result: Vec = host.iter().map(|&v| if v >= 0.0 { v } else { a * v }).collect(); - GpuTensor::from_host(&result, t.shape().to_vec(), &stream) -} - -/// Host-side multiply element-wise. -fn gpu_mul_elementwise(a: &GpuTensor, b: &GpuTensor, stream: &Arc) -> Result { - if a.shape() == b.shape() { - a.mul(b, stream) - } else if b.numel() == 1 { - let b_host = b.to_host(stream)?; - let s = b_host.first().copied().unwrap_or(1.0); - gpu_affine(a, s as f64, 0.0, stream) - } else if a.numel() == b.numel() { - a.mul(b, stream) - } else { - // Try broadcast - let b_bc = gpu_broadcast_as(b, a.shape(), stream)?; - a.mul(&b_bc, stream) - } -} - -/// Host-side add element-wise with broadcasting. -fn gpu_add_bc(a: &GpuTensor, b: &GpuTensor, stream: &Arc) -> Result { - if a.shape() == b.shape() { - a.add(b, stream) - } else if b.numel() == 1 { - let b_host = b.to_host(stream)?; - let s = b_host.first().copied().unwrap_or(0.0); - gpu_affine(a, 1.0, s as f64, stream) - } else if a.numel() == b.numel() { - a.add(b, stream) - } else { - let b_bc = gpu_broadcast_as(b, a.shape(), stream)?; - a.add(&b_bc, stream) - } -} - -/// Host-side sub element-wise with broadcasting. -fn gpu_sub_bc(a: &GpuTensor, b: &GpuTensor, stream: &Arc) -> Result { - if a.shape() == b.shape() { - a.sub(b, stream) - } else if b.numel() == 1 { - let b_host = b.to_host(stream)?; - let s = b_host.first().copied().unwrap_or(0.0); - gpu_affine(a, 1.0, -(s as f64), stream) - } else { - let b_bc = gpu_broadcast_as(b, a.shape(), stream)?; - a.sub(&b_bc, stream) - } -} - -/// Multiply a GpuTensor by a scalar f64. -fn gpu_mul_scalar(t: &GpuTensor, scalar: f64, stream: &Arc) -> Result { - gpu_affine(t, scalar, 0.0, stream) -} - -/// Host-side copy: deep clone of tensor data. -fn gpu_copy(t: &GpuTensor, stream: &Arc) -> Result { - t.gpu_clone(stream) -} - -/// Polyak-update for paired variable lists (NoisyLinear heads not in GpuVarStore). -/// target[i] = (1-tau) * target[i] + tau * online[i] element-wise. -fn polyak_update_var_pairs( - online: &[GpuTensor], - target: &[GpuTensor], - tau: f64, - stream: &Arc, -) -> Result<(), MLError> { - let tau_f32 = tau as f32; - let one_minus_tau = 1.0_f32 - tau_f32; - for (o, t) in online.iter().zip(target.iter()) { - let o_host = o.to_host(stream)?; - let mut t_host = t.to_host(stream)?; - for (tv, ov) in t_host.iter_mut().zip(o_host.iter()) { - *tv = one_minus_tau * *tv + tau_f32 * *ov; - } - // We can't mutate t in-place since it's &GpuTensor. - // The caller needs to handle mutability. For now, this is a no-op - // that at least compiles. The target_update module handles the real polyak. - } - Ok(()) -} - /// Configuration for the `DQN` #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DQNConfig { @@ -1943,26 +1488,37 @@ impl DQN { let v_max = self.config.v_max; let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0); - let batch_size = gpu_dim(&state, 0)?; + let batch_size = state.dim(0)?; let num_actions = self.config.num_actions; - // Host-side: build atoms tensor [batch, num_actions, num_atoms] - let atoms_1d = gpu_affine( - &GpuTensor::arange(0.0_f32, num_atoms as f32, 1.0_f32, &self.stream)?, - delta_z as f64, v_min as f64, &self.stream, - )?; + // GPU-native: build atoms tensor [batch, num_actions, num_atoms] + let atoms_1d = GpuTensor::arange(0.0_f32, num_atoms as f32, 1.0_f32, &self.stream)? + .affine(delta_z as f64, v_min as f64, &self.stream)?; let atoms_unsq = atoms_1d.unsqueeze(0, &self.stream)?.unsqueeze(0, &self.stream)?; - let atoms_tensor = gpu_broadcast_as(&atoms_unsq, &[batch_size, num_actions, num_atoms], &self.stream)?; + let atoms_tensor = atoms_unsq.broadcast_as(&[batch_size, num_actions, num_atoms], &self.stream)?; // Compute expected Q-values: Q(s,a) = Σ(z_i * p_i) - let prod = gpu_mul_elementwise(&z_probs, &atoms_tensor, &self.stream)?; - gpu_sum(&prod, 2, &self.stream)? // Sum over atoms + let prod = z_probs.broadcast_mul(&atoms_tensor, &self.stream)?; + // Sum over atoms: [B, A, N] -> [B, A] — host-side reduction for 3D + let prod_host = prod.to_host(&self.stream)?; + let (b, a, n_at) = prod.dims3()?; + let mut summed = Vec::with_capacity(b * a); + for bi in 0..b { + for ai in 0..a { + let mut s = 0.0_f32; + for ni in 0..n_at { + s += prod_host.get(bi * a * n_at + ai * n_at + ni).copied().unwrap_or(0.0); + } + summed.push(s); + } + } + GpuTensor::from_host(&summed, vec![b, a], &self.stream)? } else if let Some(ref dueling_net) = self.dueling_q_network { // Wave 11.4: Priority 2 - Dueling only { // DuelingQNetwork takes &[f32], convert from GpuTensor (dead code path — always None) let host = state.to_host(&self.stream)?; - let batch_size = gpu_dim(&state, 0)?; + let batch_size = state.dim(0)?; let result = dueling_net.forward(&host, batch_size)?; GpuTensor::from_host(&result, vec![batch_size, self.config.num_actions], &self.stream)? } @@ -1988,10 +1544,9 @@ impl DQN { // 2. Gradients flow through during backward (no gradient death) // 3. Prevents extreme Q-values from propagating through action selection let q_values = if self.config.enable_q_value_clipping { - let clamped = gpu_clamp( - &q_values, - self.config.q_value_clip_min as f64, - self.config.q_value_clip_max as f64, + let clamped = q_values.clamp( + self.config.q_value_clip_min, + self.config.q_value_clip_max, &self.stream, )?; @@ -2585,7 +2140,7 @@ impl DQN { MLError::ModelError("IQN network not initialized despite use_iqn=true".into()) })?; let state_embed = self.get_state_embedding(states)?; - let batch_size = gpu_dim(&state_embed, 0)?; + let batch_size = state_embed.dim(0)?; let taus = iqn_net.sample_uniform_quantiles(batch_size) .map_err(|e| MLError::ModelError(format!("Failed to sample taus: {}", e)))?; // Forward: [batch, num_actions, num_quantiles] @@ -2650,8 +2205,8 @@ impl DQN { // Host-side Gumbel-max: Q/T + Gumbel noise, then argmax let q_host = q_values.to_host(&self.stream)?; - let n = gpu_dim(&q_values, 0)?; - let num_actions = gpu_dim(&q_values, 1)?; + let n = q_values.dim(0)?; + let num_actions = q_values.dim(1)?; let mut rng = thread_rng(); let mut result = Vec::with_capacity(n); for i in 0..n { @@ -2740,12 +2295,14 @@ impl DQN { let states = states.to_dtype(NativeDType::F32, &self.stream) .map_err(|e| MLError::ModelError(e.to_string()))?; + let act_kernels = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?; let mut x = states; let num_layers = self.q_network.noisy_layers.len(); for (i, layer) in self.q_network.noisy_layers.iter().enumerate() { if i >= num_layers - 1 { break; } // Skip output layer x = layer.forward(&x)?; - x = todo_leaky_relu_fn(&x, self.q_network.leaky_relu_alpha)?; + let (y, _mask) = act_kernels.leaky_relu_fwd(&x, self.q_network.leaky_relu_alpha as f32, &self.stream)?; + x = y; } Ok(x) } @@ -2809,7 +2366,7 @@ impl DQN { "GPU batch required — GPU PER is mandatory for DQN training.".to_owned(), ) })?; - let batch_size = gpu_dim(&gpu.states, 0).map_err(|e| { + let batch_size = gpu.states.dim(0).map_err(|e| { MLError::TrainingError(format!("GPU batch dim error: {}", e)) })?; let states_tensor = gpu.states.to_dtype(NativeDType::F32, stream).map_err(|e| { @@ -3286,7 +2843,7 @@ impl DQN { // Get Q-values for first state in batch let first_state = states_tensor.narrow(0, 0, 1, &self.stream)?; let q_values = self.forward(&first_state)?; - let n_actions = gpu_dim(&q_values, 1).unwrap_or(0); + let n_actions = q_values.dim(1).unwrap_or(0); // Host-side stats computation let host = q_values.to_host(&self.stream)?;