feat: import and extend ML kernels from inferi

.github/workflows/ci.yaml

@@ -116,4 +116,5 @@ jobs:
 
       - name: Run Cargo Tests
         run: |
-          LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose -p vortx --features cpu
+          LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose -p vortx --features cpu
+          LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose -p vortx --features cpu,ml,rand

Cargo.toml

@@ -24,6 +24,8 @@ metal = ["khal/metal"]
 push_constants = ["khal/push_constants", "vortx-shaders/push_constants"]
 subgroup_ops = ["khal/subgroup_ops", "vortx-shaders/subgroup_ops"]
 unsafe_remove_boundchecks = []
+# Enables machine-learning (llm inference, reinforcement learning, etc.) operators.
+ml = ["vortx-shaders/ml"]
 
 [workspace.package]
 version = "0.4.0"
@@ -40,9 +42,11 @@ khal = { workspace = true }
 khal-std = { workspace = true }
 # Shader crate provides both GPU shader code and generated ShaderArgs via spirv_bindgen
 vortx-shaders = { version = "0.4", path = "vortx-shaders" }
+rand = { version = "0.10", optional = true }
 
 [dev-dependencies]
 nalgebra = { version = "0.35", features = ["rand"] }
+half = "2"
 futures-test = "0.3"
 serial_test = "3"
 approx = "0.5"

build.rs

@@ -18,5 +18,9 @@ fn main() {
     {
         builder = builder.feature("push_constants");
     }
+    #[cfg(feature = "ml")]
+    {
+        builder = builder.feature("ml");
+    }
     builder.build(&output_dir);
 }

src/lib.rs

@@ -14,5 +14,7 @@ pub static SPIRV_DIR: Dir<'static> = include_dir!("$OUT_DIR/shaders-spirv");
 pub use linalg::*;
 
 pub mod linalg;
+#[cfg(feature = "ml")]
+pub mod ml;
 pub mod shapes;
 pub mod tensor;

src/linalg/contiguous.rs

@@ -51,7 +51,8 @@ impl Contiguous {
         tensor_shape = tensor_shape.canonicalize();
         // println!("Tensor shape (canon): {:?}", tensor_shape);
 
-        let num_threads = tensor_shape.len() as u32;
+        let num_threads =
+            (tensor_shape.len() as u32).min(vortx_shaders::linalg::contiguous::MAX_NUM_THREADS);
 
         if let Some(offset) = offset {
             #[cfg(not(feature = "push_constants"))]

src/linalg/op_assign.rs

@@ -69,7 +69,7 @@ impl OpAssign {
         shape_a = shape_a.canonicalize();
         shape_b = shape_b.canonicalize();
 
-        let num_threads = a.len() as u32;
+        let num_threads = (a.len() as u32).min(vortx_shaders::linalg::op_assign::MAX_NUM_THREADS);
 
         #[cfg(not(feature = "push_constants"))]
         {
@@ -169,7 +169,7 @@ impl OpAssign {
         shape_a = shape_a.canonicalize();
         shape_b = shape_b.canonicalize();
 
-        let num_threads = a.len() as u32;
+        let num_threads = (a.len() as u32).min(vortx_shaders::linalg::op_assign::MAX_NUM_THREADS);
 
         // copy_with_offsets doesn't use push_constants for shapes
         shapes.insert(backend, shape_a)?;

src/ml/activation.rs

@@ -0,0 +1,85 @@
+//! Backward passes of the element-wise activations (host dispatch).
+//!
+//! The forward directions are provided by [`crate::ml::Unary`]
+//! (`UnaryOp::Tanh`, `UnaryOp::Elu`); only the gradients live here.
+
+use crate::shaders::ml::{GpuEluBackward, GpuTanhBackward};
+use crate::shapes::TensorLayoutBuffers;
+use crate::tensor::{AsTensorMut, AsTensorRef};
+use khal::Shader;
+use khal::backend::{GpuBackend, GpuBackendError, GpuPass};
+
+/// Element-wise activation gradient kernels.
+#[derive(Shader)]
+pub struct ActivationBackward {
+    /// In-place tanh backward (`g *= 1 - y^2`).
+    pub tanh_backward: GpuTanhBackward,
+    /// In-place ELU backward (`g *= 1 if y > 0 else y + 1`).
+    pub elu_backward: GpuEluBackward,
+}
+
+impl ActivationBackward {
+    /// 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).min(vortx_shaders::ml::activation::MAX_NUM_THREADS);
+
+        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 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).min(vortx_shaders::ml::activation::MAX_NUM_THREADS);
+
+        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(),
+        )
+    }
+}