diff --git a/Cargo.toml b/Cargo.toml index c888def..c7b3aac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ khal = { version = "0.1", features = ["derive"]} [dependencies] bytemuck = "1" +glamx = { version = "0.2", default-features = false, features = ["bytemuck"] } include_dir = "0.7" nalgebra = "0.34" khal = { workspace = true } diff --git a/src/linalg/activation.rs b/src/linalg/activation.rs new file mode 100644 index 0000000..19bd446 --- /dev/null +++ b/src/linalg/activation.rs @@ -0,0 +1,153 @@ +//! Element-wise activation functions (host dispatch). +//! +//! Added for zealot's MLP policy — vortx upstream has no activations. + +use crate::shaders::linalg::{GpuElu, GpuEluBackward, GpuEluVec4, GpuTanh, GpuTanhBackward}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; + +/// Element-wise activation kernels. +#[derive(Shader)] +pub struct Activation { + /// In-place tanh. + pub tanh: GpuTanh, + /// In-place tanh backward (`g *= 1 - y^2`). + pub tanh_backward: GpuTanhBackward, + /// In-place ELU. + pub elu: GpuElu, + /// In-place ELU backward (`g *= 1 if y > 0 else y + 1`). + pub elu_backward: GpuEluBackward, + /// In-place ELU, vec4 (4 contiguous f32 per thread). + pub elu_vec4: GpuEluVec4, +} + +impl Activation { + /// In-place tanh: `a = tanh(a)`. + pub fn tanh( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut a: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut a = a.as_tensor_mut(); + let shape_a = a.layout().canonicalize(); + let num_threads = a.len() as u32; + + shapes.insert(backend, shape_a)?; + let shape_a_buf = shapes.get(shape_a).unwrap(); + let mut buf_a = a.buffer_mut(); + + self.tanh + .call(pass, num_threads, &shape_a_buf.as_slice(), &mut buf_a) + } + + /// In-place tanh backward: `g *= 1 - y^2`, where `y = tanh(x)` is the forward output. + /// `g` and `y` must have the same shape. + pub fn tanh_backward( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut g: impl AsTensorMut, + y: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut g = g.as_tensor_mut(); + let y = y.as_tensor_ref(); + let shape_g = g.layout().canonicalize(); + let shape_y = y.layout().canonicalize(); + let num_threads = g.len() as u32; + + shapes.insert(backend, shape_g)?; + shapes.insert(backend, shape_y)?; + let shape_g_buf = shapes.get(shape_g).unwrap(); + let shape_y_buf = shapes.get(shape_y).unwrap(); + let mut buf_g = g.buffer_mut(); + + self.tanh_backward.call( + pass, + num_threads, + &shape_g_buf.as_slice(), + &shape_y_buf.as_slice(), + &mut buf_g, + &y.buffer(), + ) + } + + /// In-place ELU (alpha = 1): `a = a if a > 0 else exp(a) - 1`. + pub fn elu( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut a: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut a = a.as_tensor_mut(); + let shape_a = a.layout().canonicalize(); + let num_threads = a.len() as u32; + + shapes.insert(backend, shape_a)?; + let shape_a_buf = shapes.get(shape_a).unwrap(); + let mut buf_a = a.buffer_mut(); + + self.elu + .call(pass, num_threads, &shape_a_buf.as_slice(), &mut buf_a) + } + + /// In-place ELU, vec4: 4 contiguous f32 per thread (128-bit transactions). + /// Buffer length must be a multiple of 4 and contiguous (dense activations). + pub fn elu_vec4( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut a: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut a = a.as_tensor_mut(); + let shape_a = a.layout().canonicalize(); + let num_threads = (a.len() / 4) as u32; + + shapes.insert(backend, shape_a)?; + let shape_a_buf = shapes.get(shape_a).unwrap(); + let buf_a = a.buffer_mut(); + // Same bytes, viewed as vec4 (4 f32 -> 1 Vec4) for 128-bit transactions. + let mut buf_v4 = buf_a.reinterpret::(); + + self.elu_vec4 + .call(pass, num_threads, &shape_a_buf.as_slice(), &mut buf_v4) + } + + /// In-place ELU backward: `g *= 1 if y > 0 else y + 1`, where `y = elu(x)` is + /// the cached forward output. `g` and `y` must have the same shape. + pub fn elu_backward( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut g: impl AsTensorMut, + y: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut g = g.as_tensor_mut(); + let y = y.as_tensor_ref(); + let shape_g = g.layout().canonicalize(); + let shape_y = y.layout().canonicalize(); + let num_threads = g.len() as u32; + + shapes.insert(backend, shape_g)?; + shapes.insert(backend, shape_y)?; + let shape_g_buf = shapes.get(shape_g).unwrap(); + let shape_y_buf = shapes.get(shape_y).unwrap(); + let mut buf_g = g.buffer_mut(); + + self.elu_backward.call( + pass, + num_threads, + &shape_g_buf.as_slice(), + &shape_y_buf.as_slice(), + &mut buf_g, + &y.buffer(), + ) + } +} diff --git a/src/linalg/mod.rs b/src/linalg/mod.rs index 7a65987..c13ed33 100644 --- a/src/linalg/mod.rs +++ b/src/linalg/mod.rs @@ -1,14 +1,18 @@ //! Fundamental linear-algebra matrix/vector operations. +mod activation; mod contiguous; mod gemm; mod op_assign; +mod optim; mod reduce; mod repeat; +pub use activation::Activation; pub use contiguous::Contiguous; pub use gemm::{Gemm, MatrixMode, N, T}; pub use op_assign::{BinOpOffsets, OpAssign, OpAssignVariant}; +pub use optim::{Adam, AdamParams}; pub use reduce::{Reduce, ReduceVariant}; pub use repeat::Repeat; diff --git a/src/linalg/optim.rs b/src/linalg/optim.rs new file mode 100644 index 0000000..3c9b98e --- /dev/null +++ b/src/linalg/optim.rs @@ -0,0 +1,61 @@ +//! Optimizer host dispatch (Adam). Added for zealot. + +use crate::shaders::linalg::GpuAdam; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; + +// Re-export the params struct from the shader crate. +pub use vortx_shaders::linalg::optim::AdamParams; + +/// The Adam optimizer kernel. +#[derive(Shader)] +pub struct Adam { + /// One in-place Adam update step. + pub adam: GpuAdam, +} + +impl Adam { + /// Performs one in-place Adam step: updates `theta`, `m`, `v` from `grad`. + /// + /// `params` is a scalar `Tensor` (UNIFORM usage); `theta`, `grad`, + /// `m`, `v` all share the same shape. + pub fn step( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + params: impl AsTensorRef, + mut theta: impl AsTensorMut, + grad: impl AsTensorRef, + mut m: impl AsTensorMut, + mut v: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let params = params.as_tensor_ref(); + let mut theta = theta.as_tensor_mut(); + let grad = grad.as_tensor_ref(); + let mut m = m.as_tensor_mut(); + let mut v = v.as_tensor_mut(); + + let shape = theta.layout().canonicalize(); + let num_threads = theta.len() as u32; + + shapes.insert(backend, shape)?; + let shape_buf = shapes.get(shape).unwrap(); + let mut buf_theta = theta.buffer_mut(); + let mut buf_m = m.buffer_mut(); + let mut buf_v = v.buffer_mut(); + + self.adam.call( + pass, + num_threads, + &shape_buf.as_slice(), + ¶ms.buffer(), + &mut buf_theta, + &grad.buffer(), + &mut buf_m, + &mut buf_v, + ) + } +} diff --git a/vortx-shaders/Cargo.toml b/vortx-shaders/Cargo.toml index 998d75f..ebbf2f1 100644 --- a/vortx-shaders/Cargo.toml +++ b/vortx-shaders/Cargo.toml @@ -27,6 +27,10 @@ cuda = ["khal-std/cuda", "khal/cuda"] khal-std = { workspace = true } # glamx provides UVec3 and other glam types (no_std compatible, used on all targets). glamx = { version = "0.2", default-features = false, features = ["nostd-libm", "bytemuck"] } +# Cap glam < 0.33 for the shader build: spirv-std 0.10.0-alpha.1 declares glam ">=0.30.8" +# (open-ended), but glam 0.33 dropped `UVec4` under default-features=false, which breaks +# spirv-std's compile (112 errors). 0.32.1 is the newest pre-0.33 version that still works. +glam = { version = "=0.32.1", default-features = false } # Host-only dependencies (excluded on GPU targets: spirv and nvptx64). [target.'cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))'.dependencies] diff --git a/vortx-shaders/src/linalg/activation.rs b/vortx-shaders/src/linalg/activation.rs new file mode 100644 index 0000000..3967f70 --- /dev/null +++ b/vortx-shaders/src/linalg/activation.rs @@ -0,0 +1,123 @@ +//! Element-wise activation functions (tanh forward/backward). +//! +//! vortx upstream has no activations; these were added for zealot's MLP policy. +//! Uniform-shape bindings only (no push_constants variant), matching the default build. + +use super::shape::Shape; +use crate::utils::limits::MAX_NUM_WORKGROUPS; +use crate::utils::trig::stable_tanh; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use glamx::{UVec3, Vec4}; +use khal_std::{ + index::MaybeIndexUnchecked, + macros::{spirv, spirv_bindgen}, +}; + +const WORKGROUP_SIZE: u32 = 256; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Element-wise tanh, in place: `a = tanh(a)`. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_tanh( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_a: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] a: &mut [f32], +) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_a.decompose(thread_id); + let ia = shape_a.it_vec(id) as usize; + let slot = a.at_mut(ia); + *slot = stable_tanh(*slot); + } +} + +/// Backward of tanh, in place: `g *= 1 - y*y`, where `y = tanh(x)` is the forward output. +/// +/// `g` and `y` are expected to have the same shape (the per-element local derivative +/// of tanh is `1 - tanh(x)^2`, expressed in terms of the cached output `y`). +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_tanh_backward( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_g: &Shape, + #[spirv(uniform, descriptor_set = 0, binding = 1)] shape_y: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] g: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] y: &[f32], +) { + for thread_id in (invocation_id.x..shape_g.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_g.decompose(thread_id); + let ig = shape_g.it_vec(id) as usize; + let iy = shape_y.it_vec(id) as usize; + let yi = y.read(iy); + *g.at_mut(ig) *= 1.0 - yi * yi; + } +} + +/// Element-wise ELU (alpha = 1), in place: `a = a if a > 0 else exp(a) - 1`. +/// +/// Mirrors `zealot-rl`'s CPU `elu`. Hidden layers of the AGILE/rsl_rl actor/critic +/// stacks use ELU; the output layer stays linear (so this is only applied to the +/// hidden pre-activations). +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_elu( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_a: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] a: &mut [f32], +) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_a.decompose(thread_id); + let ia = shape_a.it_vec(id) as usize; + let slot = a.at_mut(ia); + let x = *slot; + *slot = if x > 0.0 { x } else { x.exp() - 1.0 }; + } +} + +/// Element-wise ELU, **vec4** in place: processes 4 contiguous f32 per thread via +/// 128-bit loads/stores. Assumes a contiguous buffer whose length is a multiple +/// of 4 (true for the dense activation buffers). The buffer is the same bytes as +/// the scalar version — only the binding type differs — so it's a drop-in for +/// contiguous tensors. Memory-bound elementwise kernels win big from the wider +/// transactions. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_elu_vec4( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_a: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] a: &mut [Vec4], +) { + let n4 = shape_a.len() / 4; + for thread_id in (invocation_id.x..n4).step_by(MAX_NUM_THREADS as usize) { + let i = thread_id as usize; + let v = a.read(i); + let e = |x: f32| if x > 0.0 { x } else { x.exp() - 1.0 }; + *a.at_mut(i) = Vec4::new(e(v.x), e(v.y), e(v.z), e(v.w)); + } +} + +/// Backward of ELU (alpha = 1), in place: `g *= 1 if y > 0 else y + 1`, where +/// `y = elu(x)` is the cached forward output. +/// +/// Valid because `elu'(x) = 1` for `x > 0` and `exp(x) = elu(x) + 1` for `x <= 0`, +/// and `y > 0 <=> x > 0`. Same cached-output formulation as `gpu_tanh_backward`, +/// matching `zealot-rl`'s `elu_grad_from_act`. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_elu_backward( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_g: &Shape, + #[spirv(uniform, descriptor_set = 0, binding = 1)] shape_y: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] g: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] y: &[f32], +) { + for thread_id in (invocation_id.x..shape_g.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_g.decompose(thread_id); + let ig = shape_g.it_vec(id) as usize; + let iy = shape_y.it_vec(id) as usize; + let yi = y.read(iy); + *g.at_mut(ig) *= if yi > 0.0 { 1.0 } else { yi + 1.0 }; + } +} diff --git a/vortx-shaders/src/linalg/mod.rs b/vortx-shaders/src/linalg/mod.rs index dd0b36b..2f9d73b 100644 --- a/vortx-shaders/src/linalg/mod.rs +++ b/vortx-shaders/src/linalg/mod.rs @@ -1,9 +1,11 @@ //! Linear algebra modules for shaders. +pub mod activation; pub mod contiguous; pub mod gemm; pub mod inv; pub mod op_assign; +pub mod optim; pub mod reduce; pub mod repeat; pub mod shape; @@ -14,12 +16,16 @@ pub use shape::{Shapes1, Shapes2, Shapes3}; // Re-export generated ShaderArgs structs (only available on host) #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] +pub use activation::{GpuElu, GpuEluBackward, GpuEluVec4, GpuTanh, GpuTanhBackward}; +#[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use contiguous::{Contiguous, ContiguousWithOffset}; #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use gemm::{GemmNaive, GemmTiled}; #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use op_assign::{GpuAdd, GpuCopy, GpuCopyWithOffsets, GpuDiv, GpuMul, GpuSub}; #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] +pub use optim::GpuAdam; +#[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use reduce::{ReduceAdd, ReduceMax, ReduceMin, ReduceMul, ReduceSqNorm}; #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use repeat::Repeat; diff --git a/vortx-shaders/src/linalg/optim.rs b/vortx-shaders/src/linalg/optim.rs new file mode 100644 index 0000000..da6b6df --- /dev/null +++ b/vortx-shaders/src/linalg/optim.rs @@ -0,0 +1,63 @@ +//! Optimizer kernels (Adam). Added for zealot; vortx upstream has no optimizers. + +use super::shape::Shape; +use crate::utils::limits::MAX_NUM_WORKGROUPS; +use glamx::UVec3; +use khal_std::{ + index::MaybeIndexUnchecked, + macros::{spirv, spirv_bindgen}, +}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; + +const WORKGROUP_SIZE: u32 = 256; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Scalar parameters for one Adam step (uniform buffer; padded to 32 bytes). +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct AdamParams { + pub lr: f32, + pub beta1: f32, + pub beta2: f32, + pub eps: f32, + /// `1 - beta1^t` (bias correction for the first moment). + pub bias_correction1: f32, + /// `1 - beta2^t` (bias correction for the second moment). + pub bias_correction2: f32, + pub pad0: f32, + pub pad1: f32, +} + +/// One in-place Adam step: updates first/second moments `m`, `v` and parameters +/// `theta` from the gradient `grad`. All buffers share `theta`'s shape. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_adam( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape: &Shape, + #[spirv(uniform, descriptor_set = 0, binding = 1)] params: &AdamParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] theta: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] grad: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] m: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] v: &mut [f32], +) { + for thread_id in (invocation_id.x..shape.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape.decompose(thread_id); + let i = shape.it_vec(id) as usize; + let g = grad.read(i); + let m_old = *m.at_mut(i); + let v_old = *v.at_mut(i); + let mi = params.beta1 * m_old + (1.0 - params.beta1) * g; + let vi = params.beta2 * v_old + (1.0 - params.beta2) * g * g; + *m.at_mut(i) = mi; + *v.at_mut(i) = vi; + let mhat = mi / params.bias_correction1; + let vhat = vi / params.bias_correction2; + *theta.at_mut(i) -= params.lr * mhat / (vhat.sqrt() + params.eps); + } +}