Add elu activation

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 }

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<f32>,
+    ) -> 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<f32>,
+        y: impl AsTensorRef<f32>,
+    ) -> 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<f32>,
+    ) -> 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<f32>,
+    ) -> 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::<glamx::Vec4>();
+
+        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<f32>,
+        y: impl AsTensorRef<f32>,
+    ) -> 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(),
+        )
+    }
+}

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;
 

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<AdamParams>` (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<AdamParams>,
+        mut theta: impl AsTensorMut<f32>,
+        grad: impl AsTensorRef<f32>,
+        mut m: impl AsTensorMut<f32>,
+        mut v: impl AsTensorMut<f32>,
+    ) -> 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(),
+            &params.buffer(),
+            &mut buf_theta,
+            &grad.buffer(),
+            &mut buf_m,
+            &mut buf_v,
+        )
+    }
+}

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]

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 };
+    }
+}