From 6aecbf94bf38c923bab05992acc290c104876314 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Wed, 26 Aug 2026 20:17:00 +0200 Subject: [PATCH 01/13] =?UTF-8?q?feat:=E2=80=AFmove=20inferi=20shaders=20i?= =?UTF-8?q?nto=20an=20optional=20ml=20module=20for=20vortx?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/ci.yaml | 3 +- Cargo.toml | 3 + build.rs | 4 + src/lib.rs | 2 + src/ml/batched_multiquery_attention.rs | 347 ++++ src/ml/concat.rs | 77 + src/ml/conv2d_nchw.rs | 58 + src/ml/conv_transpose_2d.rs | 154 ++ src/ml/gather.rs | 79 + src/ml/gemv_quant.rs | 1161 ++++++++++++ src/ml/get_rel_pos.rs | 87 + src/ml/im2col.rs | 94 + src/ml/layernorm.rs | 221 +++ src/ml/mod.rs | 47 + src/ml/pool2d.rs | 172 ++ src/ml/quantization.rs | 433 +++++ src/ml/quantized_matrix.rs | 118 ++ src/ml/reduce_axis.rs | 269 +++ src/ml/rms_norm.rs | 174 ++ src/ml/rope.rs | 246 +++ src/ml/select.rs | 147 ++ src/ml/silu.rs | 141 ++ src/ml/softmax.rs | 214 +++ src/ml/unary.rs | 1060 +++++++++++ src/ml/win_part.rs | 70 + vortx-shaders/Cargo.toml | 2 + vortx-shaders/src/lib.rs | 2 + .../src/ml/batched_multiquery_attention.rs | 57 + vortx-shaders/src/ml/concat.rs | 63 + vortx-shaders/src/ml/conv2d.rs | 127 ++ vortx-shaders/src/ml/conv_transpose_2d.rs | 275 +++ vortx-shaders/src/ml/fused_attention.rs | 617 ++++++ vortx-shaders/src/ml/gather.rs | 77 + vortx-shaders/src/ml/gemv_quant_q4_0x2.rs | 164 ++ vortx-shaders/src/ml/gemv_quant_q4_1x2.rs | 108 ++ vortx-shaders/src/ml/gemv_quant_q4_k.rs | 182 ++ vortx-shaders/src/ml/gemv_quant_q5_0x2.rs | 121 ++ vortx-shaders/src/ml/gemv_quant_q5_1x2.rs | 121 ++ vortx-shaders/src/ml/gemv_quant_q5_k.rs | 196 ++ vortx-shaders/src/ml/gemv_quant_q6_kx2.rs | 227 +++ vortx-shaders/src/ml/gemv_quant_q8_0x2.rs | 175 ++ vortx-shaders/src/ml/gemv_quant_q8_k.rs | 93 + vortx-shaders/src/ml/get_rel_pos.rs | 125 ++ vortx-shaders/src/ml/im2col.rs | 118 ++ vortx-shaders/src/ml/layernorm.rs | 240 +++ vortx-shaders/src/ml/mod.rs | 55 + vortx-shaders/src/ml/pool2d.rs | 283 +++ vortx-shaders/src/ml/reduce_axis.rs | 296 +++ vortx-shaders/src/ml/rms_norm.rs | 121 ++ vortx-shaders/src/ml/rope.rs | 141 ++ vortx-shaders/src/ml/select.rs | 56 + vortx-shaders/src/ml/silu.rs | 53 + vortx-shaders/src/ml/softmax.rs | 264 +++ vortx-shaders/src/ml/unary.rs | 1671 +++++++++++++++++ vortx-shaders/src/ml/win_part.rs | 109 ++ vortx-shaders/src/utils/half.rs | 32 + vortx-shaders/src/utils/mod.rs | 1 + 57 files changed, 11552 insertions(+), 1 deletion(-) create mode 100644 src/ml/batched_multiquery_attention.rs create mode 100644 src/ml/concat.rs create mode 100644 src/ml/conv2d_nchw.rs create mode 100644 src/ml/conv_transpose_2d.rs create mode 100644 src/ml/gather.rs create mode 100644 src/ml/gemv_quant.rs create mode 100644 src/ml/get_rel_pos.rs create mode 100644 src/ml/im2col.rs create mode 100644 src/ml/layernorm.rs create mode 100644 src/ml/mod.rs create mode 100644 src/ml/pool2d.rs create mode 100644 src/ml/quantization.rs create mode 100644 src/ml/quantized_matrix.rs create mode 100644 src/ml/reduce_axis.rs create mode 100644 src/ml/rms_norm.rs create mode 100644 src/ml/rope.rs create mode 100644 src/ml/select.rs create mode 100644 src/ml/silu.rs create mode 100644 src/ml/softmax.rs create mode 100644 src/ml/unary.rs create mode 100644 src/ml/win_part.rs create mode 100644 vortx-shaders/src/ml/batched_multiquery_attention.rs create mode 100644 vortx-shaders/src/ml/concat.rs create mode 100644 vortx-shaders/src/ml/conv2d.rs create mode 100644 vortx-shaders/src/ml/conv_transpose_2d.rs create mode 100644 vortx-shaders/src/ml/fused_attention.rs create mode 100644 vortx-shaders/src/ml/gather.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q4_0x2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q4_1x2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q4_k.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q5_0x2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q5_1x2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q5_k.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q6_kx2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q8_0x2.rs create mode 100644 vortx-shaders/src/ml/gemv_quant_q8_k.rs create mode 100644 vortx-shaders/src/ml/get_rel_pos.rs create mode 100644 vortx-shaders/src/ml/im2col.rs create mode 100644 vortx-shaders/src/ml/layernorm.rs create mode 100644 vortx-shaders/src/ml/mod.rs create mode 100644 vortx-shaders/src/ml/pool2d.rs create mode 100644 vortx-shaders/src/ml/reduce_axis.rs create mode 100644 vortx-shaders/src/ml/rms_norm.rs create mode 100644 vortx-shaders/src/ml/rope.rs create mode 100644 vortx-shaders/src/ml/select.rs create mode 100644 vortx-shaders/src/ml/silu.rs create mode 100644 vortx-shaders/src/ml/softmax.rs create mode 100644 vortx-shaders/src/ml/unary.rs create mode 100644 vortx-shaders/src/ml/win_part.rs create mode 100644 vortx-shaders/src/utils/half.rs diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 84d82bc..b4d5bb8 100644 --- a/.github/workflows/ci.yaml +++ b/.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 \ No newline at end of file + LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose -p vortx --features cpu + LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose -p vortx --features cpu,ml,rand diff --git a/Cargo.toml b/Cargo.toml index be3deda..2caeccb 100644 --- a/Cargo.toml +++ b/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,6 +42,7 @@ 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"] } diff --git a/build.rs b/build.rs index 3329cf3..108a260 100644 --- a/build.rs +++ b/build.rs @@ -18,5 +18,9 @@ fn main() { { builder = builder.feature("push_constants"); } + #[cfg(feature = "ml")] + { + builder = builder.feature("ml"); + } builder.build(&output_dir); } diff --git a/src/lib.rs b/src/lib.rs index 1b69520..9957bc2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,3 +16,5 @@ pub use linalg::*; pub mod linalg; pub mod shapes; pub mod tensor; +#[cfg(feature = "ml")] +pub mod ml; \ No newline at end of file diff --git a/src/ml/batched_multiquery_attention.rs b/src/ml/batched_multiquery_attention.rs new file mode 100644 index 0000000..053ee75 --- /dev/null +++ b/src/ml/batched_multiquery_attention.rs @@ -0,0 +1,347 @@ +use vortx_shaders::ml::AttentionParams; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::{DMatrix, DVector}; +use crate::tensor::{AsTensorMut, AsTensorRef}; +use crate::ml::SoftMax; + +#[derive(Shader)] +/// Fused attention shader - combines Q*K^T, scale, mask, softmax, and *V into one kernel. +pub struct FusedAttention { + pub fused_attention: vortx_shaders::ml::FusedAttention, + pub fused_attention_online: vortx_shaders::ml::FusedAttentionOnline, + pub flash_attention: vortx_shaders::ml::FlashAttention, +} + +impl FusedAttention { + /// Launch the fused attention kernel. + /// + /// This replaces the 4-dispatch attention (matmul -> mask -> softmax -> matmul) + /// with a single fused kernel dispatch. + pub fn launch( + &self, + _backend: &GpuBackend, + pass: &mut GpuPass, + params: &AttentionParams, + params_gpu: impl AsTensorRef, + q: impl AsTensorRef, + key_cache: impl AsTensorRef, + value_cache: impl AsTensorRef, + mut xb: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + const WORKGROUP_SIZE: u32 = 128; + const MAX_SEQ_LEN: u32 = 2048; + + let params_gpu = params_gpu.as_tensor_ref(); + let q = q.as_tensor_ref(); + let key_cache = key_cache.as_tensor_ref(); + let value_cache = value_cache.as_tensor_ref(); + let mut xb = xb.as_tensor_mut(); + + let n_heads = params.n_heads; + let seq_len = params.pos + 1; + + // Choose kernel based on sequence length: + // - fused_attention: stores all scores in shared memory (fast, limited to 2048 tokens) + // - flash_attention: block-wise processing with online softmax (efficient for long sequences) + if seq_len <= MAX_SEQ_LEN { + let mut buf_out = xb.buffer_mut(); + self.fused_attention.call( + pass, + [n_heads * WORKGROUP_SIZE, 1, 1], + ¶ms_gpu.buffer(), + &q.buffer(), + &key_cache.buffer(), + &value_cache.buffer(), + &mut buf_out, + ) + } else { + let mut buf_out = xb.buffer_mut(); + self.flash_attention.call( + pass, + [n_heads * WORKGROUP_SIZE, 1, 1], + ¶ms_gpu.buffer(), + &q.buffer(), + &key_cache.buffer(), + &value_cache.buffer(), + &mut buf_out, + ) + } + } + + pub fn run_cpu( + params: &AttentionParams, + q: &DVector, + key_cache: &DMatrix, + value_cache: &DMatrix, + attn: &mut DMatrix, + xb: &mut DVector, + ) { + // The number of embedding vector elements associated to each query head. + let head_size = params.head_size as usize; + // The number of query head associated to one key/value head. + let kv_mul = params.kv_mul as usize; + + // Multihead attention. Iterate over all head. + // TODO: in llama2.c, each head is iterated on in parallel. + for h in 0..params.n_heads as usize { + // Get the query vector for this head. + let q = q.rows(h * head_size, head_size); + // Attention scores for this head. + let mut att = attn.column_mut(h); + + // Iterate over all timesteps (tokens in the sequence), including the current one, but + // not past the current one due to causality. + // See the KV cache explanation there: https://youtu.be/Mn_9W1nCFLo?si=3n4GH9f2OzMb5Np0&t=2940 + // -> This is iterating through all the green columns (from K^t) that are the rotated + // (by RoPE). The values set in this loop into the `att` variable here (attention + // scores) are the elements in the pink row (at the bottom of the QK^t matrix) divide + // by sqrt(params.head_size) (in other words, this is what's given to softmax afterward. + for t in 0..=params.pos as usize { + // Get the key vector for this head and at this timestep. + let k = key_cache.column(t); // TODO: does key_cache have the right dim? + let k_head = k.rows((h / kv_mul) * head_size, head_size); + + // Calculate the attention score as the dot product of q and k. + let mut score = q.dot(&k_head); + score /= (head_size as f32).sqrt(); + // Save the score to the attention buffer. + att[t] = score; + } + + // Softmax the scores to get attention weights from 0..=pos inclusively. + SoftMax::run_cpu(&mut att.rows_mut(0, params.pos as usize + 1)); + + // Weighted sum of the values, store back into xb. + // /!\ xb is now changing semantic, storing the weighted sums for all the heads. + // Now xb contains the "Attention 4" row from https://youtu.be/Mn_9W1nCFLo?si=550ar5aUg1I1k60l&t=2940. + let mut xb = xb.rows_mut(h * head_size, head_size); + xb.fill(0.0); + for t in 0..=params.pos as usize { + let v = value_cache.column(t); + let v_head = v.rows((h / kv_mul) * head_size, head_size); + xb.axpy(att[t], &v_head, 1.0); + } + } + } +} + +/* +#[cfg(test)] +mod test { + use crate::ml::ops::{AttentionParams, SoftMax}; + use nalgebra::{DMatrix, DVector}; + use khal::gpu::GpuInstance; + use khal::kernel::CommandEncoderExt; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::{Tensor, Tensor, Tensor}; + use khal::Shader; + use crate::Gemv; + use wgpu::BufferUsages; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_attention() { + let gpu = GpuInstance::new().await.unwrap(); + let batched_multihead_attention = + super::BatchedMultiqueryAttention::from_backend(gpu.backend()).unwrap(); + let mut encoder = gpu.backend().create_command_encoder(&Default::default()); + + // let mut params = AttentionParams { seq_len: 131072, kv_dim: 256, kv_mul: 6, n_heads: 12, head_size: 128, pos: 9 }; + let params = AttentionParams { + seq_len: 1024, + kv_dim: 768, + kv_mul: 1, + n_heads: 12, + head_size: 64, + pos: 6, + }; + + let q = DVector::new_random((params.n_heads * params.head_size) as usize); + let key_cache = DMatrix::new_random(params.kv_dim as usize, params.seq_len as usize); + let value_cache = DMatrix::new_random(params.kv_dim as usize, params.seq_len as usize); + let mut attn = DMatrix::zeros(params.seq_len as usize, params.n_heads as usize); + let mut xb = DVector::zeros((params.n_heads * params.head_size) as usize); + + let gpu_params = Tensor::scalar(gpu.backend(), params, BufferUsages::UNIFORM); + let gpu_q = Tensor::vector(gpu.backend(), q.as_slice(), BufferUsages::STORAGE); + let gpu_key_cache = Tensor::matrix(gpu.backend(), &key_cache, BufferUsages::STORAGE); + let gpu_value_cache = Tensor::matrix(gpu.backend(), &value_cache, BufferUsages::STORAGE); + let gpu_attn = Tensor::matrix( + gpu.backend(), + &attn, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ); + let gpu_xb = Tensor::vector( + gpu.backend(), + xb.as_slice(), + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ); + + let gpu_staging_xb = Tensor::vector_uninit( + gpu.backend(), + xb.len() as u32, + BufferUsages::MAP_READ | BufferUsages::COPY_DST, + ); + let gpu_staging_attn = Tensor::matrix_uninit( + gpu.backend(), + attn.nrows() as u32, + attn.ncols() as u32, + BufferUsages::MAP_READ | BufferUsages::COPY_DST, + ); + + let mut pass = encoder.compute_pass("test", None); + batched_multihead_attention.launch( + gpu.backend(), + &mut pass, + params.n_heads, + &gpu_params, + &gpu_q, + &gpu_key_cache, + &gpu_value_cache, + &gpu_attn, + &gpu_xb, + ); + drop(pass); + + gpu_staging_xb.copy_from(&mut encoder, &gpu_xb); + gpu_staging_attn.copy_from(&mut encoder, &gpu_attn); + + gpu.queue().submit(Some(encoder.finish())); + + super::FudedAttention::run_cpu( + ¶ms, + &q, + &key_cache, + &value_cache, + &mut attn, + &mut xb, + ); + + approx::assert_relative_eq!( + DVector::from(gpu_staging_xb.read(gpu.backend()).await.unwrap()), + xb, + epsilon = 1.0e-5 + ); + + approx::assert_relative_eq!( + DMatrix::from_vec( + attn.nrows(), + attn.ncols(), + gpu_staging_attn.read(gpu.backend()).await.unwrap() + ), + attn, + epsilon = 1.0e-5 + ); + } + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_attention_multi() { + let gpu = GpuInstance::new().await.unwrap(); + let batched_multihead_attention = + super::BatchedMultiqueryAttention::from_backend(gpu.backend()).unwrap(); + let shapes = TensorLayoutBuffers::new(); + let matmul = Gemv::from_backend(gpu.backend()).unwrap(); + let softmax = SoftMax::from_backend(gpu.backend()).unwrap(); + + // let mut params = AttentionParams { seq_len: 131072, kv_dim: 256, kv_mul: 6, n_heads: 12, head_size: 128, pos: 0 }; + let mut params = AttentionParams { + seq_len: 1024, + kv_dim: 768, + kv_mul: 1, + n_heads: 12, + head_size: 64, + pos: 0, + }; + + let q = DVector::new_random((params.n_heads * params.head_size) as usize); + let key_cache = DMatrix::new_random(params.kv_dim as usize, params.seq_len as usize); + let value_cache = DMatrix::new_random(params.kv_dim as usize, params.seq_len as usize); + let mut attn = DMatrix::zeros(params.seq_len as usize, params.n_heads as usize); + let mut xb = DVector::zeros((params.n_heads * params.head_size) as usize); + + let gpu_q = Tensor::vector(gpu.backend(), q.as_slice(), BufferUsages::STORAGE); + let gpu_key_cache = Tensor::matrix(gpu.backend(), &key_cache, BufferUsages::STORAGE); + let gpu_value_cache = Tensor::matrix(gpu.backend(), &value_cache, BufferUsages::STORAGE); + let gpu_attn = Tensor::matrix( + gpu.backend(), + &attn, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ); + let gpu_xb = Tensor::vector( + gpu.backend(), + xb.as_slice(), + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ); + + let gpu_staging_xb = Tensor::vector_uninit( + gpu.backend(), + xb.len() as u32, + BufferUsages::MAP_READ | BufferUsages::COPY_DST, + ); + let gpu_staging_attn = Tensor::matrix_uninit( + gpu.backend(), + attn.nrows() as u32, + attn.ncols() as u32, + BufferUsages::MAP_READ | BufferUsages::COPY_DST, + ); + + for pos in 0..9 { + let mut encoder = gpu.backend().create_command_encoder(&Default::default()); + params.pos = pos; + + let gpu_params = Tensor::scalar(gpu.backend(), params, BufferUsages::UNIFORM); + + let mut pass = encoder.compute_pass("test", None); + batched_multihead_attention.launch( + gpu.backend(), + &shapes, + gpu.queue(), + &mut pass, + &matmul, + &softmax, + ¶ms, + &gpu_params, + &gpu_q, + &gpu_key_cache, + &gpu_value_cache, + &gpu_attn, + &gpu_xb, + ); + drop(pass); + + gpu_staging_xb.copy_from(&mut encoder, &gpu_xb); + gpu_staging_attn.copy_from(&mut encoder, &gpu_attn); + + gpu.queue().submit(Some(encoder.finish())); + + super::BatchedMultiqueryAttention::run_cpu( + ¶ms, + &q, + &key_cache, + &value_cache, + &mut attn, + &mut xb, + ); + + // NOTE: we can't compare attn since they don't have the same layout. + // approx::assert_relative_eq!( + // DMatrix::from_vec( + // attn.nrows(), + // attn.ncols(), + // gpu_staging_attn.read(gpu.backend()).await.unwrap() + // ), + // attn, + // epsilon = 1.0e-5 + // ); + + approx::assert_relative_eq!( + DVector::from(gpu_staging_xb.read(gpu.backend()).await.unwrap()), + xb, + epsilon = 1.0e-5 + ); + } + } +} +*/ diff --git a/src/ml/concat.rs b/src/ml/concat.rs new file mode 100644 index 0000000..da4049e --- /dev/null +++ b/src/ml/concat.rs @@ -0,0 +1,77 @@ +//! Concat operation: concatenates tensors along a given axis. + +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; + +/// Shader for the Concat operation. +#[derive(Shader)] +pub struct Concat { + pub concat_copy: vortx_shaders::ml::ConcatCopy, +} + +impl Concat { + /// Copy a source tensor into a slice of the destination tensor along a given axis. + /// + /// This is used to implement concat by calling it once per input tensor. + /// `offset` indicates where this tensor's data should be placed along the axis. + pub fn launch_copy( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut dest: impl AsTensorMut, + src: impl AsTensorRef, + axis: u32, + offset: u32, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src = src.as_tensor_ref(); + let len = src.len() as u32; + let max_threads = 65535u32; + + // Upload params [axis, offset] + let params_buf = TensorBuilder::scalar(BufferUsages::STORAGE | BufferUsages::COPY_DST) + .build_init(backend, &[axis, offset])?; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src.layout())?; + + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src = shapes.get(src.layout()).unwrap(); + let mut buf_dest = dest.buffer_mut(); + + self.concat_copy.call( + pass, + [len.min(max_threads), 1, 1], + &shape_dest.as_slice(), + &shape_src.as_slice(), + &mut buf_dest, + &src.buffer(), + ¶ms_buf.buffer().as_slice(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = crate::shaders::linalg::Shapes2 { + shape_a: dest.layout().into(), + shape_b: src.layout().into(), + }; + let mut buf_dest = dest.buffer_mut(); + + self.concat_copy.call( + pass, + [len.min(max_threads), 1, 1], + &mut buf_dest, + &src.buffer(), + ¶ms_buf.buffer().as_slice(), + shapes_val, + ) + } + } +} diff --git a/src/ml/conv2d_nchw.rs b/src/ml/conv2d_nchw.rs new file mode 100644 index 0000000..b2d7767 --- /dev/null +++ b/src/ml/conv2d_nchw.rs @@ -0,0 +1,58 @@ +//! 2D Convolution operation (NCHW format). +//! +//! This implementation works with ONNX tensor format directly. + +use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; +use khal::Shader; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +#[derive(Shader)] +pub struct Conv2dNchw { + pub conv_2d_nchw: vortx_shaders::ml::Conv2dNchw, +} + +impl Conv2dNchw { + /// Launch Conv2d operation. + /// + /// Input: [N, C_in, H, W] + /// Weight: [C_out, C_in, K_H, K_W] + /// Output: [N, C_out, H_out, W_out] + pub fn launch( + &self, + pass: &mut GpuPass, + params: &GpuBuffer, + input: impl AsTensorRef, + weight: impl AsTensorRef, + mut output: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut output = output.as_tensor_mut(); + let input = input.as_tensor_ref(); + let weight = weight.as_tensor_ref(); + + let output_len = output.len() as u32; + let mut buf_output = output.buffer_mut(); + + self.conv_2d_nchw.call( + pass, + [output_len, 1, 1], + &mut buf_output, + &input.buffer(), + &weight.buffer(), + ¶ms.as_slice(), + )?; + + Ok(()) + } +} + +/// Compute output dimensions for convolution. +pub fn conv_output_size( + input_size: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, +) -> u32 { + let effective_kernel = dilation * (kernel_size - 1) + 1; + (input_size + 2 * padding - effective_kernel) / stride + 1 +} diff --git a/src/ml/conv_transpose_2d.rs b/src/ml/conv_transpose_2d.rs new file mode 100644 index 0000000..e6a91e1 --- /dev/null +++ b/src/ml/conv_transpose_2d.rs @@ -0,0 +1,154 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; + +#[derive(Shader)] +pub struct ConvTranspose2d { + pub init_dest: vortx_shaders::ml::conv_transpose_2d::InitDest, + pub init_wdata: vortx_shaders::ml::conv_transpose_2d::InitWdata, + pub init_src_a: vortx_shaders::ml::conv_transpose_2d::InitSrcA, + pub init_src_b: vortx_shaders::ml::conv_transpose_2d::InitSrcB, + pub conv_transpose_2d_ref: vortx_shaders::ml::conv_transpose_2d::ConvTranspose2dRef, + pub conv_transpose_2d: vortx_shaders::ml::conv_transpose_2d::ConvTranspose2d, +} + +impl ConvTranspose2d { + pub fn launch_ref( + &self, + backend: &GpuBackend, + pass: &mut GpuPass, + shapes: &mut TensorLayoutBuffers, + stride: &Tensor, + mut dest: impl AsTensorMut, + src0: impl AsTensorRef, + src1: impl AsTensorRef, + mut wdata: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src0 = src0.as_tensor_ref(); + let src1 = src1.as_tensor_ref(); + let mut wdata = wdata.as_tensor_mut(); + + assert_eq!(wdata.len(), src0.len() + src1.len()); + + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src0.layout())?; + shapes.insert(backend, src1.layout())?; + shapes.insert(backend, wdata.layout())?; + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src0 = shapes.get(src0.layout()).unwrap(); + let shape_src1 = shapes.get(src1.layout()).unwrap(); + let shape_wdata = shapes.get(wdata.layout()).unwrap(); + + // init_dest: shape_dest, dest + { + let dest_len = dest.len() as u32; + let mut buf_dest = dest.buffer_mut(); + self.init_dest.call( + pass, + [dest_len, 1, 1], + &shape_dest.as_slice(), + &mut buf_dest, + )?; + } + + // init_wdata: shape_wdata, wdata + { + let wdata_len = wdata.len() as u32; + let mut buf_wdata = wdata.buffer_mut(); + self.init_wdata.call( + pass, + [wdata_len, 1, 1], + &shape_wdata.as_slice(), + &mut buf_wdata, + )?; + } + + // init_src_a: shape_src0, src0, wdata + { + let mut buf_wdata = wdata.buffer_mut(); + self.init_src_a.call( + pass, + [src0.len() as u32, 1, 1], + &shape_src0.as_slice(), + &src0.buffer(), + &mut buf_wdata, + )?; + } + + // init_src_b: shape_src0, shape_src1, src1, wdata + { + let mut buf_wdata = wdata.buffer_mut(); + self.init_src_b.call( + pass, + [src1.len() as u32, 1, 1], + &shape_src0.as_slice(), + &shape_src1.as_slice(), + &src1.buffer(), + &mut buf_wdata, + )?; + } + + // conv_transpose_2d_ref: shape_src0, shape_src1, shape_dest, stride, wdata, dest + { + let dest_size2 = dest.size(2); + let mut buf_dest = dest.buffer_mut(); + self.conv_transpose_2d_ref.call( + pass, + [dest_size2, 1, 1], + &shape_src0.as_slice(), + &shape_src1.as_slice(), + &shape_dest.as_slice(), + &stride.buffer().as_slice(), + &wdata.buffer(), + &mut buf_dest, + )?; + } + + Ok(()) + } + + pub fn launch( + &self, + backend: &GpuBackend, + pass: &mut GpuPass, + shapes: &mut TensorLayoutBuffers, + stride: &mut Tensor, + mut dest: impl AsTensorMut, + src0: impl AsTensorRef, + src1: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src0 = src0.as_tensor_ref(); + let src1 = src1.as_tensor_ref(); + + let src0 = src0.permute([3, 0, 1, 2]); + let src1 = src1.permute([2, 0, 1, 3]); + + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src0.layout())?; + shapes.insert(backend, src1.layout())?; + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src0 = shapes.get(src0.layout()).unwrap(); + let shape_src1 = shapes.get(src1.layout()).unwrap(); + + // conv_transpose_2d: shape_src1, shape_src0, shape_dest, stride, src1, src0, dest + let dest_size2 = dest.size(2); + let mut buf_dest = dest.buffer_mut(); + + self.conv_transpose_2d.call( + pass, + [dest_size2, 1, 1], + &shape_src1.as_slice(), + &shape_src0.as_slice(), + &shape_dest.as_slice(), + &stride.buffer().as_slice(), + &src1.buffer(), + &src0.buffer(), + &mut buf_dest, + )?; + + Ok(()) + } +} diff --git a/src/ml/gather.rs b/src/ml/gather.rs new file mode 100644 index 0000000..eb04562 --- /dev/null +++ b/src/ml/gather.rs @@ -0,0 +1,79 @@ +//! Gather operation: gathers elements from a tensor based on indices along an axis. + +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; + +/// Shader for the Gather operation. +#[derive(Shader)] +pub struct Gather { + pub gather: vortx_shaders::ml::gather::Gather, +} + +impl Gather { + /// Launch the gather operation. + /// + /// Gathers elements from `src` along `axis` based on `indices`, writing to `dest`. + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut dest: impl AsTensorMut, + src: impl AsTensorRef, + indices: impl AsTensorRef, + axis: u32, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src = src.as_tensor_ref(); + let indices = indices.as_tensor_ref(); + let len = dest.len() as u32; + let max_threads = 65535u32; + + // Upload axis parameter + let axis_buf = TensorBuilder::scalar(BufferUsages::STORAGE | BufferUsages::COPY_DST) + .build_init(backend, &[axis])?; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src.layout())?; + + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src = shapes.get(src.layout()).unwrap(); + let mut buf_dest = dest.buffer_mut(); + + self.gather.call( + pass, + [len.min(max_threads), 1, 1], + &shape_dest.as_slice(), + &shape_src.as_slice(), + &mut buf_dest, + &src.buffer(), + &indices.buffer(), + &axis_buf.buffer().as_slice(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = crate::shaders::linalg::Shapes2 { + shape_a: dest.layout().into(), + shape_b: src.layout().into(), + }; + let mut buf_dest = dest.buffer_mut(); + + self.gather.call( + pass, + [len.min(max_threads), 1, 1], + &mut buf_dest, + &src.buffer(), + &indices.buffer(), + &axis_buf.buffer().as_slice(), + shapes_val, + ) + } + } +} diff --git a/src/ml/gemv_quant.rs b/src/ml/gemv_quant.rs new file mode 100644 index 0000000..401ad08 --- /dev/null +++ b/src/ml/gemv_quant.rs @@ -0,0 +1,1161 @@ +use crate::ml::quantization::{BlockQ4K, BlockQ5K, BlockQ8K}; +use crate::ml::quantized_matrix::GpuQuantTensor; +use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; +use crate::Gemm; + +#[cfg(feature = "rand")] +use rand::distr::{Distribution, StandardUniform}; +#[cfg(feature = "rand")] +use rand::{Rng, RngExt}; + +pub trait QuantizedValue { + /// Number of dequantized elements the quantized value represents. + const DEQUANTIZED_LEN: usize; +} + +impl QuantizedValue for f32 { + const DEQUANTIZED_LEN: usize = 1; +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ8_0x2([u32; 17]); + +impl QuantizedValue for GpuBlockQ8_0x2 { + const DEQUANTIZED_LEN: usize = 64; +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ4_0x2([u32; 9]); + +impl QuantizedValue for GpuBlockQ4_0x2 { + const DEQUANTIZED_LEN: usize = 64; +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ4_1x2([u32; 10]); + +impl QuantizedValue for GpuBlockQ4_1x2 { + const DEQUANTIZED_LEN: usize = 64; +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ5_0x2([u32; 11]); + +impl QuantizedValue for GpuBlockQ5_0x2 { + const DEQUANTIZED_LEN: usize = 64; +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ5_1x2([u32; 12]); + +impl QuantizedValue for GpuBlockQ5_1x2 { + const DEQUANTIZED_LEN: usize = 64; +} + +pub type GpuBlockQ8K = BlockQ8K; +pub type GpuBlockQ5K = BlockQ5K; +pub type GpuBlockQ4K = BlockQ4K; + +#[derive(Copy, Clone, Debug, PartialEq)] +#[repr(C)] +pub struct GpuBlockQ6Kx2([u32; 105]); + +impl QuantizedValue for GpuBlockQ6Kx2 { + const DEQUANTIZED_LEN: usize = 512; +} + +// SAFETY: These impls are safe, they don't exist in bytemuck because they don't +// provide impls for non-power-of-two largeish arrays. +unsafe impl bytemuck::Zeroable for GpuBlockQ6Kx2 {} +unsafe impl bytemuck::Pod for GpuBlockQ6Kx2 {} + +#[cfg(feature = "rand")] +macro_rules! impl_rand { + ($($t: ident, $len: literal);*) => {$( + impl Distribution<$t> for StandardUniform { + fn sample(&self, rng: &mut R) -> $t { + // TODO: are all bit representations valid? + $t([0; $len].map(|_| rng.random())) + } + } + )*}; +} + +#[cfg(feature = "rand")] +impl_rand!( + GpuBlockQ8_0x2, 17; + GpuBlockQ4_0x2, 9; + GpuBlockQ4_1x2, 10; + GpuBlockQ5_0x2, 11; + GpuBlockQ5_1x2, 12; + GpuBlockQ6Kx2, 105 +); + +pub struct GemvQuant { + pub gemm_f32: Gemm, + pub gemv_q8: GemvQ8_0x2, + pub gemv_q5: GemvQ5_0x2, + pub gemv_q4: GemvQ4_0x2, + pub gemv_q5_1: GemvQ5_1x2, + pub gemv_q4_1: GemvQ4_1x2, + pub gemv_q8_k: GemvQ8K, + pub gemv_q6_k: GemvQ6Kx2, + pub gemv_q5_k: GemvQ5K, + pub gemv_q4_k: GemvQ4K, +} + +impl GemvQuant { + pub fn from_backend(backend: &GpuBackend) -> Result { + Ok(Self { + gemm_f32: Gemm::from_backend(backend).unwrap(), // ?, + gemv_q5: GemvQ5_0x2::from_backend(backend).unwrap(), // ?, + gemv_q5_1: GemvQ5_1x2::from_backend(backend).unwrap(), // ?, + gemv_q4: GemvQ4_0x2::from_backend(backend).unwrap(), // ?, + gemv_q4_1: GemvQ4_1x2::from_backend(backend).unwrap(), // ?, + gemv_q8_k: GemvQ8K::from_backend(backend).unwrap(), // ?, + gemv_q6_k: GemvQ6Kx2::from_backend(backend).unwrap(), // ?, + gemv_q5_k: GemvQ5K::from_backend(backend).unwrap(), // ?, + gemv_q4_k: GemvQ4K::from_backend(backend).unwrap(), // ?, + gemv_q8: GemvQ8_0x2::from_backend(backend).unwrap(), // ?, + }) + } +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ8_0x2 { + pub gemv: vortx_shaders::ml::gemv_quant_q8_0x2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ5_0x2 { + pub gemv: vortx_shaders::ml::gemv_quant_q5_0x2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ5_1x2 { + pub gemv: vortx_shaders::ml::gemv_quant_q5_1x2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ4_0x2 { + pub gemv: vortx_shaders::ml::gemv_quant_q4_0x2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ4_1x2 { + pub gemv: vortx_shaders::ml::gemv_quant_q4_1x2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ8K { + pub gemv: vortx_shaders::ml::gemv_quant_q8_k::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ6Kx2 { + pub gemv: vortx_shaders::ml::gemv_quant_q6_kx2::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ5K { + pub gemv: vortx_shaders::ml::gemv_quant_q5_k::Gemv, +} + +#[derive(Shader)] +/// Shader for computing the product of a matrix and a vector. +pub struct GemvQ4K { + pub gemv: vortx_shaders::ml::gemv_quant_q4_k::Gemv, +} + +impl GemvQuant { + /// Queues this shader to compute `out = m * v`. + pub fn launch( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut out: impl AsTensorMut, + m: &GpuQuantTensor, + v: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut out = out.as_tensor_mut(); + + // TODO: add a function to convert a View to a View> + // then remove `TensorLayout::f32_to_vec4`. + let mut v = v.as_tensor_ref(); + assert_eq!(m.rank(), 2); + assert_eq!(v.rank(), 1); + assert_eq!(out.rank(), 1); + assert!(v.is_contiguous()); + assert!(out.is_contiguous()); + + // Special case: if using f32, use the f32 kernel launcher instead. + if let GpuQuantTensor::F32(m_f32) = m { + v = v.unsqueeze(1); + out = out.unsqueeze(1); + return self.gemm_f32.dispatch(backend, shapes, pass, out, m_f32, v); + } + + // assert_eq!( + // m.layout().size[1], + // v.layout().size[0], + // "Gemv: dimension mismatch." + // ); + // assert_eq!( + // out.layout().size[0], + // m.layout().size[0], + // "Gemv: dimension mismatch." + // ); + + let v_shape = match m { + GpuQuantTensor::F32(_) => unreachable!(), + GpuQuantTensor::Q8_0(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q5_0(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q5_1(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q4_0(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q4_1(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q8K(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q6K(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q5K(_) => v.layout().f32_to_vec4(), + GpuQuantTensor::Q4K(_) => v.layout().f32_to_vec4(), + }; + + let out_shape = match m { + GpuQuantTensor::F32(_) => unreachable!(), + // Optimized shaders use workgroup reduction → one Vec4 per workgroup. + GpuQuantTensor::Q8_0(_) + | GpuQuantTensor::Q4_0(_) + | GpuQuantTensor::Q6K(_) + | GpuQuantTensor::Q5K(_) + | GpuQuantTensor::Q4K(_) => out.layout().f32_to_vec4(), + // Non-optimized shaders (Q4_1, Q5_0, Q5_1, Q8K) index output by + // global_invocation_id and write Vec4::splat per row. On WebGPU the + // OOB writes are clamped; on CUDA they corrupt memory. These shader + // paths are currently broken for CUDA and only work by accident on + // WebGPU. They are rarely used in practice (most models use Q4K/Q5K). + _ => out.layout(), + }; + + // Canonicalize the shape for coherence with other matmul shaders. + let shape_m_canon = m.layout().canonicalize(); + #[cfg(not(feature = "push_constants"))] + shapes.insert(backend, out_shape)?; // TODO: propagate error + #[cfg(not(feature = "push_constants"))] + shapes.insert(backend, v_shape)?; // TODO: propagate error + #[cfg(not(feature = "push_constants"))] + shapes.insert(backend, shape_m_canon)?; // TODO: propagate error + #[cfg(not(feature = "push_constants"))] + let _shape_out = shapes.get(out_shape).unwrap(); + #[cfg(not(feature = "push_constants"))] + let _shape_v = shapes.get(v_shape).unwrap(); + #[cfg(not(feature = "push_constants"))] + let shape_m = shapes.get(shape_m_canon).unwrap(); + + let launch = match m { + GpuQuantTensor::F32(_) => unreachable!(), + GpuQuantTensor::Q8_0(_) + | GpuQuantTensor::Q4_0(_) + | GpuQuantTensor::Q6K(_) + | GpuQuantTensor::Q5K(_) + | GpuQuantTensor::Q4K(_) => out.layout().size[0] / 4, + _ => m.layout().size[0].div_ceil(64), + }; + + let grid = DispatchGrid::Grid([launch, 1, 1]); + + // Dispatch to the appropriate kernel based on quantization type. + // Each variant has its own generated args type, but all share the same field names + // (shape_m, out, m, v) since the shader functions have the same signature. + macro_rules! dispatch_gemv { + ($kernel:expr, $tensor:expr) => {{ + #[cfg(not(feature = "push_constants"))] + { + let mut buf_out = out.buffer_mut().reinterpret(); + $kernel.gemv.call( + pass, + grid, + &shape_m.as_slice(), + &mut buf_out, + &$tensor.buffer().as_slice().reinterpret(), + &v.buffer().reinterpret(), + ) + } + #[cfg(feature = "push_constants")] + { + let shapes_val: vortx_shaders::linalg::Shapes1 = shape_m_canon.into(); + let mut buf_out = out.buffer_mut().reinterpret(); + $kernel.gemv.call( + pass, + grid, + &mut buf_out, + &$tensor.buffer().as_slice().reinterpret(), + &v.buffer().reinterpret(), + shapes_val, + ) + } + }}; + } + + match m { + GpuQuantTensor::F32(_) => unreachable!(), + GpuQuantTensor::Q8_0(tensor) => { + dispatch_gemv!(&self.gemv_q8, tensor) + } + GpuQuantTensor::Q5_0(tensor) => { + dispatch_gemv!(&self.gemv_q5, tensor) + } + GpuQuantTensor::Q5_1(tensor) => { + dispatch_gemv!(&self.gemv_q5_1, tensor) + } + GpuQuantTensor::Q4_0(tensor) => { + dispatch_gemv!(&self.gemv_q4, tensor) + } + GpuQuantTensor::Q4_1(tensor) => { + dispatch_gemv!(&self.gemv_q4_1, tensor) + } + GpuQuantTensor::Q8K(tensor) => { + dispatch_gemv!(&self.gemv_q8_k, tensor) + } + GpuQuantTensor::Q6K(tensor) => { + dispatch_gemv!(&self.gemv_q6_k, tensor) + } + GpuQuantTensor::Q5K(tensor) => { + dispatch_gemv!(&self.gemv_q5_k, tensor) + } + GpuQuantTensor::Q4K(tensor) => { + dispatch_gemv!(&self.gemv_q4_k, tensor) + } + } + } +} + +#[cfg(test)] +#[cfg(feature = "rand")] +mod test { + use super::*; + use crate::ml::quantization::*; + use crate::ml::quantized_matrix::GpuQuantTensor; + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; + use khal::BufferUsages; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + /// Generate a random valid f16 scale (avoids NaN/Inf). + /// Includes subnormal values to cover the full f16 range. + fn rand_f16_scale() -> u16 { + let val: f32 = rand::random::() * 2.0 - 1.0; // [-1, 1] + // Scale down ~25% of values into the subnormal f16 range (< 2^-14). + let val = if rand::random::() < 64 { + val * 1e-5 + } else { + val + }; + half::f16::from_f32(val).to_bits() + } + + fn rand_block_q8_0() -> BlockQ8_0 { + BlockQ8_0 { + scale: rand_f16_scale(), + data: rand::random(), + } + } + + fn rand_block_q4_0() -> BlockQ4_0 { + BlockQ4_0 { + d: rand_f16_scale(), + qs: rand::random(), + } + } + + fn rand_block_q8k() -> BlockQ8K { + BlockQ8K { + d: rand::random::() * 2.0 - 1.0, + qs: [0i8; 256].map(|_| rand::random()), + bsums: rand::random(), + } + } + + fn rand_block_q4k() -> BlockQ4K { + BlockQ4K { + d: rand_f16_scale(), + dmin: rand_f16_scale(), + scales: rand::random(), + qs: [0u8; 128].map(|_| rand::random()), + } + } + + fn rand_block_q4_1() -> BlockQ4_1 { + BlockQ4_1 { + d: rand_f16_scale(), + m: rand_f16_scale(), + qs: rand::random(), + } + } + + fn rand_block_q5_0() -> BlockQ5_0 { + BlockQ5_0 { + d: rand_f16_scale(), + qh: rand::random(), + qs: rand::random(), + } + } + + fn rand_block_q5_1() -> BlockQ5_1 { + BlockQ5_1 { + d: rand_f16_scale(), + m: rand_f16_scale(), + qh: rand::random(), + qs: rand::random(), + } + } + + fn rand_block_q6k() -> BlockQ6K { + BlockQ6K { + ql: [0u8; 128].map(|_| rand::random()), + qh: [0u8; 64].map(|_| rand::random()), + scales: [0i8; 16].map(|_| rand::random()), + d: rand_f16_scale(), + } + } + + fn rand_block_q5k() -> BlockQ5K { + BlockQ5K { + d: rand_f16_scale(), + dmin: rand_f16_scale(), + scales: rand::random(), + qh: rand::random(), + qs: [0u8; 128].map(|_| rand::random()), + } + } + + /// CPU reference: dequantize quantized blocks into a full f32 matrix, then + /// compute `out = matrix * vec`. + fn cpu_gemv(dequantized_matrix: &[f32], rows: usize, cols: usize, vec: &[f32]) -> Vec { + assert_eq!(dequantized_matrix.len(), rows * cols); + assert_eq!(vec.len(), cols); + let mut out = vec![0.0f32; rows]; + for r in 0..rows { + let mut sum = 0.0f64; // accumulate in f64 for reference accuracy + for c in 0..cols { + sum += dequantized_matrix[r * cols + c] as f64 * vec[c] as f64; + } + out[r] = sum as f32; + } + out + } + + /// Helper: run a gemv_quant launch and read back the result. + async fn run_gemv( + backend: &GpuBackend, + quant_tensor: &GpuQuantTensor, + v_data: &[f32], + rows: u32, + ) -> Vec { + let gemv = GemvQuant::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + let v = Tensor::vector(backend, v_data, BufferUsages::STORAGE).unwrap(); + let mut out = Tensor::::vector_uninit( + backend, + rows, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test_gemv", None); + gemv.launch(backend, &mut shapes, &mut pass, &mut out, quant_tensor, &v) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + let mut result = vec![0.0f32; rows as usize]; + backend + .slow_read_buffer(out.buffer(), &mut result) + .await + .unwrap(); + result + } + + // ========================================================================= + // Q8_0 + // ========================================================================= + + /// Dequantize a GpuBlockQ8_0x2 slice into flat f32s. Each GPU block = 2 CPU blocks = 64 f32s. + fn dequantize_q8_0x2(blocks: &[GpuBlockQ8_0x2]) -> Vec { + let cpu_blocks: &[BlockQ8_0] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q8_0_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; // must be multiple of 64 + const BLOCKS_PER_ROW: usize = COLS / 64; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 32); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q8_0()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q8_0x2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q8_0(m); + let actual = run_gemv(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q8_0 row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q8_0_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q8_0_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q8_0_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q8_0_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q8_0_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q8_0_generic(&backend).await; + } + + // ========================================================================= + // Q4_0 + // ========================================================================= + + fn dequantize_q4_0x2(blocks: &[GpuBlockQ4_0x2]) -> Vec { + let cpu_blocks: &[BlockQ4_0] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q4_0_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 64; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 32); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q4_0()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q4_0x2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q4_0(m); + let actual = run_gemv(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q4_0 row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4_0_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q4_0_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q4_0_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q4_0_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4_0_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q4_0_generic(&backend).await; + } + + // ========================================================================= + // Q4K + // ========================================================================= + + fn dequantize_q4k(blocks: &[GpuBlockQ4K]) -> Vec { + blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q4k_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 256; + const TOTAL_BLOCKS: usize = ROWS * BLOCKS_PER_ROW; + + let blocks: Vec = (0..TOTAL_BLOCKS).map(|_| rand_block_q4k()).collect(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q4k(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q4K(m); + let actual = run_gemv(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q4K row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4k_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q4k_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q4k_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q4k_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4k_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q4k_generic(&backend).await; + } + + // ========================================================================= + // Q5K + // ========================================================================= + + fn dequantize_q5k(blocks: &[GpuBlockQ5K]) -> Vec { + blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q5k_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 256; + const TOTAL_BLOCKS: usize = ROWS * BLOCKS_PER_ROW; + + let blocks: Vec = (0..TOTAL_BLOCKS).map(|_| rand_block_q5k()).collect(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q5k(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q5K(m); + let actual = run_gemv(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q5K row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5k_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q5k_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q5k_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q5k_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5k_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q5k_generic(&backend).await; + } + + // ========================================================================= + // Q6K (optimized path, uses shared memory + workgroup reduction) + // ========================================================================= + + fn dequantize_q6kx2(blocks: &[GpuBlockQ6Kx2]) -> Vec { + let cpu_blocks: &[BlockQ6K] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q6k_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 512; // must be multiple of 512 (Q6Kx2 = 2 blocks of 256) + const BLOCKS_PER_ROW: usize = COLS / 512; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 256); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q6k()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q6kx2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q6K(m); + let actual = run_gemv(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q6K row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q6k_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q6k_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q6k_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q6k_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q6k_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q6k_generic(&backend).await; + } + + /// Helper for non-optimized shaders (Q4_1, Q5_0, Q5_1, Q8K) which write + /// Vec4::splat(sum) per row. Allocates 4x output to prevent OOB, then + /// extracts the x-component of each Vec4. + async fn run_gemv_vec4_per_row( + backend: &GpuBackend, + quant_tensor: &GpuQuantTensor, + v_data: &[f32], + rows: u32, + ) -> Vec { + let gemv = GemvQuant::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + let v = Tensor::vector(backend, v_data, BufferUsages::STORAGE).unwrap(); + let mut out = Tensor::::vector_uninit( + backend, + rows * 4, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test_gemv", None); + gemv.launch(backend, &mut shapes, &mut pass, &mut out, quant_tensor, &v) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + let mut raw = vec![0.0f32; (rows * 4) as usize]; + backend + .slow_read_buffer(out.buffer(), &mut raw) + .await + .unwrap(); + (0..rows as usize).map(|i| raw[i * 4]).collect() + } + + // --- Q4_1 --- + + fn dequantize_q4_1x2(blocks: &[GpuBlockQ4_1x2]) -> Vec { + let cpu_blocks: &[BlockQ4_1] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q4_1_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 64; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 32); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q4_1()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q4_1x2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q4_1(m); + let actual = run_gemv_vec4_per_row(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q4_1 row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4_1_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q4_1_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q4_1_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q4_1_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q4_1_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q4_1_generic(&backend).await; + } + + // --- Q5_0 --- + + fn dequantize_q5_0x2(blocks: &[GpuBlockQ5_0x2]) -> Vec { + let cpu_blocks: &[BlockQ5_0] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q5_0_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 64; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 32); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q5_0()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q5_0x2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q5_0(m); + let actual = run_gemv_vec4_per_row(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q5_0 row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5_0_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q5_0_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q5_0_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q5_0_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5_0_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q5_0_generic(&backend).await; + } + + // --- Q5_1 --- + + fn dequantize_q5_1x2(blocks: &[GpuBlockQ5_1x2]) -> Vec { + let cpu_blocks: &[BlockQ5_1] = bytemuck::cast_slice(blocks); + cpu_blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q5_1_generic(backend: &GpuBackend) { + const ROWS: usize = 64; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 64; + const TOTAL_CPU_BLOCKS: usize = ROWS * (COLS / 32); + + let cpu_blocks: Vec = (0..TOTAL_CPU_BLOCKS).map(|_| rand_block_q5_1()).collect(); + let blocks: Vec = bytemuck::cast_slice(&cpu_blocks).to_vec(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q5_1x2(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q5_1(m); + let actual = run_gemv_vec4_per_row(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q5_1 row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5_1_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q5_1_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q5_1_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q5_1_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q5_1_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q5_1_generic(&backend).await; + } + + // --- Q8K --- + + fn dequantize_q8k(blocks: &[GpuBlockQ8K]) -> Vec { + blocks.iter().flat_map(|b| b.dequantize()).collect() + } + + async fn test_gemv_q8k_generic(backend: &GpuBackend) { + // Q8K has workgroup size 32 but dispatch uses div_ceil(64), so only 32 + // rows are computed per workgroup. Use 32 rows to stay within bounds. + const ROWS: usize = 32; + const COLS: usize = 256; + const BLOCKS_PER_ROW: usize = COLS / 256; + const TOTAL_BLOCKS: usize = ROWS * BLOCKS_PER_ROW; + + let blocks: Vec = (0..TOTAL_BLOCKS).map(|_| rand_block_q8k()).collect(); + let v: Vec = (0..COLS) + .map(|_| rand::random::() * 2.0 - 1.0) + .collect(); + + let deq = dequantize_q8k(&blocks); + let expected = cpu_gemv(&deq, ROWS, COLS, &v); + + let m = Tensor::matrix( + backend, + ROWS as u32, + BLOCKS_PER_ROW as u32, + &blocks, + BufferUsages::STORAGE, + ) + .unwrap(); + let qt = GpuQuantTensor::Q8K(m); + let actual = run_gemv_vec4_per_row(backend, &qt, &v, ROWS as u32).await; + + for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() { + let diff = (a - e).abs(); + let denom = e.abs().max(1.0); + assert!( + diff / denom < 0.01, + "Q8K row {i}: gpu={a} cpu={e} diff={diff}" + ); + } + } + + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q8k_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_gemv_q8k_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gemv_q8k_cpu() { + let backend = GpuBackend::Cpu; + test_gemv_q8k_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gemv_q8k_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_gemv_q8k_generic(&backend).await; + } +} diff --git a/src/ml/get_rel_pos.rs b/src/ml/get_rel_pos.rs new file mode 100644 index 0000000..a264d8c --- /dev/null +++ b/src/ml/get_rel_pos.rs @@ -0,0 +1,87 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +#[derive(Shader)] +pub struct GetRelPos { + pub get_rel_pos: vortx_shaders::ml::get_rel_pos::GetRelPos, + pub add_rel_pos_phase_a: vortx_shaders::ml::get_rel_pos::AddRelPosPhaseA, + pub add_rel_pos_phase_b: vortx_shaders::ml::get_rel_pos::AddRelPosPhaseB, +} + +impl GetRelPos { + pub fn launch( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut result: impl AsTensorMut, + source: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut result = result.as_tensor_mut(); + let source = source.as_tensor_ref(); + shapes.insert(backend, result.layout())?; + shapes.insert(backend, source.layout())?; + let shape_result = shapes.get(result.layout()).unwrap(); + let shape_source = shapes.get(source.layout()).unwrap(); + + let result_len = result.len() as u32; + let mut buf_result = result.buffer_mut(); + + self.get_rel_pos.call( + pass, + [result_len, 1, 1], + &shape_result.as_slice(), + &shape_source.as_slice(), + &mut buf_result, + &source.buffer(), + ) + } + + pub fn launch_add_rel_pos( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut dst: impl AsTensorMut, + src1: impl AsTensorRef, + src2: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut dst = dst.as_tensor_mut(); + let src1 = src1.as_tensor_ref(); + let src2 = src2.as_tensor_ref(); + + assert_eq!(src1.layout().size, src2.layout().size); + assert_eq!(src1.size(3), dst.size(2)); + assert_eq!(src1.size(1) * src1.size(1), dst.size(1)); + assert_eq!(src1.size(0) * src1.size(1), dst.size(0)); + + shapes.insert(backend, src1.layout())?; + let shape_src1 = shapes.get(src1.layout()).unwrap(); + + // Phase A: add_rel_pos_phase_a(shape_src1, dst, src1) + { + let mut buf_dst = dst.buffer_mut(); + self.add_rel_pos_phase_a.call( + pass, + [src1.len() as u32, 1, 1], + &shape_src1.as_slice(), + &mut buf_dst, + &src1.buffer(), + )?; + } + + // Phase B: add_rel_pos_phase_b(shape_src1, dst, src2) + { + let mut buf_dst = dst.buffer_mut(); + self.add_rel_pos_phase_b.call( + pass, + [src1.len() as u32, 1, 1], + &shape_src1.as_slice(), + &mut buf_dst, + &src2.buffer(), + ) + } + } +} diff --git a/src/ml/im2col.rs b/src/ml/im2col.rs new file mode 100644 index 0000000..613ffcd --- /dev/null +++ b/src/ml/im2col.rs @@ -0,0 +1,94 @@ +use khal::backend::{Backend, DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; + + +pub type Im2ColConfig = vortx_shaders::ml::im2col::Im2ColParams; + +#[derive(Shader)] +pub struct Im2Col { + pub im2col: vortx_shaders::ml::im2col::Im2col, +} + +impl Im2Col { + // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] + // kernel: [OC, IC, KH, KW] + // input: [N, IC, IH, IW] + // result: [N, OH, OW, IC*KH*KW] + pub fn launch( + &self, + backend: &GpuBackend, + pass: &mut GpuPass, + params: &mut Tensor, + mut result: impl AsTensorMut, + kernel: impl AsTensorRef, + input: impl AsTensorRef, + s0: u32, + s1: u32, + p0: u32, + p1: u32, + d0: u32, + d1: u32, + is_2d: bool, + ) -> Result<(), GpuBackendError> { + let mut result = result.as_tensor_mut(); + let kernel = kernel.as_tensor_ref(); + let input = input.as_tensor_ref(); + + let ilayout = input.layout(); + let klayout = kernel.layout(); + let rlayout = result.layout(); + + let ic = ilayout.size[if is_2d { 2 } else { 0 }]; + let ih = if is_2d { ilayout.size[0] } else { 1 }; + let iw = ilayout.size[1]; + + let kh = if is_2d { klayout.size[0] } else { 1 }; + let kw = klayout.size[1]; + + let oh = if is_2d { rlayout.size[2] } else { 1 }; + let ow = rlayout.size[0]; + + let offset_delta = ilayout.stride[if is_2d { 2 } else { 0 }]; + let batch_offset = ilayout.stride[if is_2d { 3 } else { 2 }]; + + let pelements = ow * kw * kh; + let chw = ic * kh * kw; + + let config = Im2ColConfig { + batch_offset, + offset_delta, + IC: ic, + IW: iw, + IH: ih, + OW: ow, + OH: oh, + KW: kw, + KH: kh, + pelements, + CHW: chw, + s0: s0 as i32, + s1: s1 as i32, + p0: p0 as i32, + p1: p1 as i32, + d0: d0 as i32, + d1: d1 as i32, + }; + + backend.write_buffer(params.buffer_mut(), 0, &[config])?; + + let batch = ilayout.size[if is_2d { 3 } else { 2 }]; + let grid = [(ow * kw * kh).div_ceil(32), oh, batch * ic]; + let mut buf_result = result.buffer_mut(); + + self.im2col.call( + pass, + DispatchGrid::Grid(grid), + ¶ms.buffer().as_slice(), + &input.buffer(), + &mut buf_result, + )?; + + Ok(()) + } +} diff --git a/src/ml/layernorm.rs b/src/ml/layernorm.rs new file mode 100644 index 0000000..ff8d666 --- /dev/null +++ b/src/ml/layernorm.rs @@ -0,0 +1,221 @@ +use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::DVector; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +#[derive(Shader)] +/// Shader implementing the layer normalization kernel. +pub struct LayerNorm { + pub layernorm_cols: vortx_shaders::ml::layernorm::LayernormCols, + pub layernorm_rows: vortx_shaders::ml::layernorm::LayernormRows, +} + +impl LayerNorm { + pub fn launch_cols( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut output: impl AsTensorMut, + input: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let input = input.as_tensor_ref().canonicalize(); + let mut output = output.as_tensor_mut().canonicalize(); + let shape_input = input.layout(); + let shape_output = output.layout(); + assert_eq!( + shape_input.size, shape_output.size, + "LayerNorm: dimension mismatch." + ); + + let grid = [ + shape_input.size[3], + shape_input.size[1], + shape_input.size[0], + ]; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, shape_input)?; + shapes.insert(backend, shape_output)?; + let in_shape = shapes.get(shape_input).unwrap(); + let out_shape = shapes.get(shape_output).unwrap(); + let mut out_buf = output.buffer_mut(); + + self.layernorm_cols.call( + pass, + DispatchGrid::Grid(grid), + &in_shape.as_slice(), + &out_shape.as_slice(), + &input.buffer(), + &mut out_buf, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = shape_input.into(); + let mut out_buf = output.buffer_mut(); + + self.layernorm_cols.call( + pass, + DispatchGrid::Grid(grid), + &input.buffer(), + &mut out_buf, + shapes_val, + ) + } + } + + pub fn launch_rows( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut output: impl AsTensorMut, + input: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let input = input.as_tensor_ref().canonicalize(); + let mut output = output.as_tensor_mut().canonicalize(); + let shape_input = input.layout(); + let shape_output = output.layout(); + + assert_eq!( + shape_input.size, shape_output.size, + "LayerNorm: dimension mismatch." + ); + + let grid = [ + shape_input.size[2], + shape_input.size[1], + shape_input.size[0], + ]; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, shape_input)?; + shapes.insert(backend, shape_output)?; + let in_shape = shapes.get(shape_input).unwrap(); + let out_shape = shapes.get(shape_output).unwrap(); + let mut out_buf = output.buffer_mut(); + + self.layernorm_rows.call( + pass, + DispatchGrid::Grid(grid), + &in_shape.as_slice(), + &out_shape.as_slice(), + &input.buffer(), + &mut out_buf, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = shape_input.into(); + let mut out_buf = output.buffer_mut(); + + self.layernorm_rows.call( + pass, + DispatchGrid::Grid(grid), + &input.buffer(), + &mut out_buf, + shapes_val, + ) + } + } + + /// The layernorm function. + /// + /// See for details on the + /// math. + pub fn run_cpu(res: &mut DVector, v: &DVector) { + const NUDGE_FACTOR: f32 = 1.0e-5; + let mean = v.mean(); + res.zip_apply(v, |y, v| *y = v - mean); + let variance = res.norm_squared() / (res.len() as f32); + let scale = 1.0 / (variance + NUDGE_FACTOR).sqrt(); + *res *= scale; + } +} + +#[cfg(test)] +mod test { + use crate::ml::LayerNorm; + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_layernorm_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_layernorm_generic(&backend).await; + } + + async fn gpu_layernorm_generic(backend: &GpuBackend) { + let layernorm = super::LayerNorm::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + const LEN: u32 = 1757; + + let v0 = DVector::new_random(LEN as usize); + let out = DVector::new_random(LEN as usize); + let mut out_read = DVector::zeros(LEN as usize); + let gpu_v0 = + Tensor::vector(backend, &v0, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + let mut gpu_out = + Tensor::vector(backend, &v0, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + layernorm + .launch_rows( + backend, + &mut shapes, + &mut pass, + &mut gpu_out, + gpu_v0.as_view(), + ) + .unwrap(); + drop(pass); + + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_out.buffer(), out_read.as_mut_slice()) + .await + .unwrap(); + + let mut cpu_result = out; + LayerNorm::run_cpu(&mut cpu_result, &v0); + + approx::assert_relative_eq!(out_read, cpu_result, epsilon = 1.0e-3); + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_layernorm_cpu() { + let backend = GpuBackend::Cpu; + gpu_layernorm_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_layernorm_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_layernorm_generic(&backend).await; + } +} diff --git a/src/ml/mod.rs b/src/ml/mod.rs new file mode 100644 index 0000000..63bf926 --- /dev/null +++ b/src/ml/mod.rs @@ -0,0 +1,47 @@ +//! Primitives for building LLM inferences. + +mod batched_multiquery_attention; +mod concat; +mod conv2d_nchw; +mod conv_transpose_2d; +mod gather; +mod gemv_quant; +mod get_rel_pos; +mod im2col; +mod layernorm; +mod pool2d; +mod reduce_axis; +mod rms_norm; +mod rope; +mod select; +mod silu; +mod softmax; +mod unary; +mod win_part; +mod quantized_matrix; +pub mod quantization; + +pub use quantized_matrix::*; +pub use batched_multiquery_attention::{ + FusedAttention, +}; +pub use concat::Concat; +pub use conv2d_nchw::{conv_output_size, Conv2dNchw}; +pub use conv_transpose_2d::ConvTranspose2d; +pub use gather::Gather; +pub use gemv_quant::{ + GemvQuant, GpuBlockQ4K, GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ5K, GpuBlockQ5_0x2, + GpuBlockQ5_1x2, GpuBlockQ6Kx2, GpuBlockQ8K, GpuBlockQ8_0x2, QuantizedValue, +}; +pub use get_rel_pos::GetRelPos; +pub use im2col::{Im2Col, Im2ColConfig}; +pub use layernorm::LayerNorm; +pub use pool2d::{pool_output_size, GlobalPool2dConfig, Pool2d, Pool2dConfig}; +pub use reduce_axis::{ReduceAxis, ReduceOp}; +pub use rms_norm::{RmsNorm, RmsNormConfig}; +pub use rope::{RoPE, RoPEConfig, RoPEVariant}; +pub use select::Select; +pub use silu::Silu; +pub use softmax::SoftMax; +pub use unary::{Unary, UnaryOp}; +pub use win_part::WinPart; diff --git a/src/ml/pool2d.rs b/src/ml/pool2d.rs new file mode 100644 index 0000000..65ee3c8 --- /dev/null +++ b/src/ml/pool2d.rs @@ -0,0 +1,172 @@ +//! 2D Pooling operations (MaxPool2d, AvgPool2d, GlobalAvgPool2d, GlobalMaxPool2d). + +use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; +use khal::Shader; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +/// Pool2d configuration parameters. +#[derive(Copy, Clone, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable, Debug, Default)] +#[repr(C)] +pub struct Pool2dConfig { + pub input_h: u32, + pub input_w: u32, + pub output_h: u32, + pub output_w: u32, + pub kernel_h: u32, + pub kernel_w: u32, + pub stride_h: u32, + pub stride_w: u32, + pub pad_h: u32, + pub pad_w: u32, + pub channels: u32, + pub batch_size: u32, + pub count_include_pad: u32, // For avg pool only + pub _padding: [u32; 3], +} + +/// Global pooling configuration. +#[derive(Copy, Clone, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable, Debug, Default)] +#[repr(C)] +pub struct GlobalPool2dConfig { + pub input_h: u32, + pub input_w: u32, + pub channels: u32, + pub batch_size: u32, +} + +#[derive(Shader)] +pub struct Pool2d { + pub max_pool_2d: vortx_shaders::ml::pool2d::MaxPool2d, + pub avg_pool_2d: vortx_shaders::ml::pool2d::AvgPool2d, + pub global_avg_pool_2d: vortx_shaders::ml::pool2d::GlobalAvgPool2d, + pub global_max_pool_2d: vortx_shaders::ml::pool2d::GlobalMaxPool2d, +} + +impl Pool2d { + /// Launch MaxPool2d operation. + /// + /// Input: [N, C, H, W], Output: [N, C, H_out, W_out] + pub fn launch_max_pool( + &self, + pass: &mut GpuPass, + params: &GpuBuffer, + input: impl AsTensorRef, + mut output: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut output = output.as_tensor_mut(); + let input = input.as_tensor_ref(); + + let output_len = output.len() as u32; + let mut buf_output = output.buffer_mut(); + + self.max_pool_2d.call( + pass, + [output_len, 1, 1], + &mut buf_output, + &input.buffer(), + ¶ms.as_slice(), + )?; + + Ok(()) + } + + /// Launch AvgPool2d operation. + /// + /// Input: [N, C, H, W], Output: [N, C, H_out, W_out] + pub fn launch_avg_pool( + &self, + pass: &mut GpuPass, + params: &GpuBuffer, + input: impl AsTensorRef, + mut output: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut output = output.as_tensor_mut(); + let input = input.as_tensor_ref(); + + let output_len = output.len() as u32; + let mut buf_output = output.buffer_mut(); + + self.avg_pool_2d.call( + pass, + [output_len, 1, 1], + &mut buf_output, + &input.buffer(), + ¶ms.as_slice(), + )?; + + Ok(()) + } + + /// Launch GlobalAvgPool2d operation. + /// + /// Input: [N, C, H, W], Output: [N, C, 1, 1] (stored as [N, C]) + pub fn launch_global_avg_pool( + &self, + pass: &mut GpuPass, + params: &GpuBuffer, + input: impl AsTensorRef, + mut output: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut output = output.as_tensor_mut(); + let input = input.as_tensor_ref(); + + let output_len = output.len() as u32; + let mut buf_output = output.buffer_mut(); + + self.global_avg_pool_2d.call( + pass, + [output_len, 1, 1], + &mut buf_output, + &input.buffer(), + ¶ms.as_slice(), + )?; + + Ok(()) + } + + /// Launch GlobalMaxPool2d operation. + /// + /// Input: [N, C, H, W], Output: [N, C, 1, 1] (stored as [N, C]) + pub fn launch_global_max_pool( + &self, + pass: &mut GpuPass, + params: &GpuBuffer, + input: impl AsTensorRef, + mut output: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut output = output.as_tensor_mut(); + let input = input.as_tensor_ref(); + + let output_len = output.len() as u32; + let mut buf_output = output.buffer_mut(); + + self.global_max_pool_2d.call( + pass, + [output_len, 1, 1], + &mut buf_output, + &input.buffer(), + ¶ms.as_slice(), + )?; + + Ok(()) + } +} + +/// Compute output dimensions for pooling. +pub fn pool_output_size( + input_size: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + ceil_mode: bool, +) -> u32 { + let effective_kernel = dilation * (kernel_size - 1) + 1; + let numerator = input_size + 2 * padding - effective_kernel; + + if ceil_mode { + numerator.div_ceil(stride) + 1 + } else { + numerator / stride + 1 + } +} diff --git a/src/ml/quantization.rs b/src/ml/quantization.rs new file mode 100644 index 0000000..80d4ca9 --- /dev/null +++ b/src/ml/quantization.rs @@ -0,0 +1,433 @@ +//! Quantization and unquantization structures. +//! +//! This is inspired heavily from [ggml-common.h](https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-common.h#L144). + +#![allow(clippy::identity_op)] +#![allow(clippy::explicit_counter_loop)] + +use crate::ml::QuantizedValue; + +#[cfg(feature = "rand")] +use rand::distr::StandardUniform; +#[cfg(feature = "rand")] +use rand::prelude::Distribution; +#[cfg(feature = "rand")] +use rand::{Rng, RngExt}; + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug)] +#[repr(C)] +/// A single `f16` value. +pub struct BlockF16 { + pub data: u16, +} + +impl BlockF16 { + pub fn dequantize(self) -> f32 { + decode_f16(self.data) + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug)] +#[repr(C)] +/// A single `f16` value. +pub struct BlockBF16 { + pub data: u16, +} + +impl BlockBF16 { + pub fn dequantize(self) -> f32 { + decode_bf16(self.data) + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L43-L46 +pub struct BlockQ8_0 { + pub scale: u16, // f16 + pub data: [i8; 32], +} + +impl BlockQ8_0 { + pub const ELEMENTS_PER_BLOCK: usize = 32; + + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L1609 + pub fn dequantize(self) -> [f32; 32] { + let scale = decode_f16(self.scale); + self.data.map(|v| v as f32 * scale) + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L11-L14 +pub struct BlockQ4_0 { + pub d: u16, // f16 + pub qs: [u8; 32 / 2], +} + +impl BlockQ4_0 { + pub const ELEMENTS_PER_BLOCK: usize = 32; + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L1515 + pub fn dequantize(self) -> [f32; 32] { + let mut result = [0.0; 32]; + let d = decode_f16(self.d); + + for j in 0..Self::ELEMENTS_PER_BLOCK / 2 { + let x0 = (self.qs[j] & 0x0F) as i32 - 8; + let x1 = (self.qs[j] >> 4) as i32 - 8; + + result[j] = x0 as f32 * d; + result[j + Self::ELEMENTS_PER_BLOCK / 2] = x1 as f32 * d; + } + + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L18-L22 +pub struct BlockQ4_1 { + pub d: u16, // f16 + pub m: u16, + pub qs: [u8; 32 / 2], +} + +impl BlockQ4_1 { + pub const ELEMENTS_PER_BLOCK: usize = 32; + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L1535 + pub fn dequantize(self) -> [f32; 32] { + let mut result = [0.0; 32]; + let d = decode_f16(self.d); + let m = decode_f16(self.m); + + for j in 0..Self::ELEMENTS_PER_BLOCK / 2 { + let x0 = self.qs[j] & 0x0F; + let x1 = self.qs[j] >> 4; + + result[j] = x0 as f32 * d + m; + result[j + Self::ELEMENTS_PER_BLOCK / 2] = x1 as f32 * d + m; + } + + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L26-L30 +pub struct BlockQ5_0 { + pub d: u16, // f16 + pub qh: [u8; 4], + pub qs: [u8; 32 / 2], +} + +impl BlockQ5_0 { + pub const ELEMENTS_PER_BLOCK: usize = 32; + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L1556 + pub fn dequantize(self) -> [f32; 32] { + let mut result = [0.0; 32]; + let d = decode_f16(self.d); + let qh: u32 = bytemuck::cast(self.qh); + + for j in 0..Self::ELEMENTS_PER_BLOCK / 2 { + let xh_0 = ((qh >> j) << 4) & 0x10; + let xh_1 = (qh >> (j + 12)) & 0x10; + let x0 = ((self.qs[j] as u32 & 0x0F) | xh_0) as i32 - 16; + let x1 = ((self.qs[j] as u32 >> 4) | xh_1) as i32 - 16; + + result[j] = x0 as f32 * d; + result[j + Self::ELEMENTS_PER_BLOCK / 2] = x1 as f32 * d; + } + + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L34-L39 +pub struct BlockQ5_1 { + pub d: u16, // delta + pub m: u16, // min + pub qh: [u8; 4], // 5-th bit of quants + pub qs: [u8; 32 / 2], // nibbles / quants +} + +impl BlockQ5_1 { + pub const ELEMENTS_PER_BLOCK: usize = 32; + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L1582 + pub fn dequantize(self) -> [f32; 32] { + let mut result = [0.0; 32]; + let d = decode_f16(self.d); + let m = decode_f16(self.m); + let qh: u32 = bytemuck::cast(self.qh); + + for j in 0..Self::ELEMENTS_PER_BLOCK / 2 { + let xh_0 = ((qh >> j) << 4) & 0x10; + let xh_1 = (qh >> (j + 12)) & 0x10; + let x0 = (self.qs[j] as u32 & 0x0F) | xh_0; + let x1 = (self.qs[j] as u32 >> 4) | xh_1; + + result[j] = x0 as f32 * d + m; + result[j + Self::ELEMENTS_PER_BLOCK / 2] = x1 as f32 * d + m; + } + + result + } +} + +const QK_K: usize = 256; +const K_SCALE_SIZE: usize = 12; + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L161-L165 +pub struct BlockQ8K { + pub d: f32, // delta + pub qs: [i8; QK_K], // quants + pub bsums: [i16; QK_K / 16], // sum of quants in groups of 16 +} + +impl QuantizedValue for BlockQ8K { + const DEQUANTIZED_LEN: usize = QK_K; +} + +#[cfg(feature = "rand")] +impl Distribution for StandardUniform { + fn sample(&self, rng: &mut R) -> BlockQ8K { + // TODO: are all bit representations valid? + BlockQ8K { + d: rng.random(), + qs: [0; QK_K].map(|_| rng.random()), + bsums: rng.random(), + } + } +} + +impl BlockQ8K { + pub fn dequantize(self) -> [f32; QK_K] { + let mut result = [0.0; QK_K]; + for j in 0..QK_K { + result[j] = self.d * self.qs[j] as f32; + } + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L152-L157 +pub struct BlockQ6K { + pub ql: [u8; QK_K / 2], // quants, lower 4 bits + pub qh: [u8; QK_K / 4], // quants, upper 2 bits + pub scales: [i8; QK_K / 16], // scales, quantized with 8 bits + pub d: u16, // super-block scale +} + +impl BlockQ6K { + pub const ELEMENTS_PER_BLOCK: usize = QK_K; + // https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L2970 + pub fn dequantize(self) -> [f32; QK_K] { + let mut result = [0.0; QK_K]; + + let d = decode_f16(self.d); + let mut i = 0; + + for _ in (0..QK_K).step_by(128) { + for l in 0..32 { + let is = l / 16; + + let ql0 = self.ql[i * 64 + l + 0]; + let ql32 = self.ql[i * 64 + l + 32]; + let qh = self.qh[i * 32 + l]; + + let q1 = ((ql0 & 0xF) | (((qh >> 0) & 3) << 4)) as i8 - 32; + let q2 = ((ql32 & 0xF) | (((qh >> 2) & 3) << 4)) as i8 - 32; + let q3 = ((ql0 >> 4) | (((qh >> 4) & 3) << 4)) as i8 - 32; + let q4 = ((ql32 >> 4) | (((qh >> 6) & 3) << 4)) as i8 - 32; + + result[i * 128 + l + 0] = d * self.scales[i * 8 + is + 0] as f32 * q1 as f32; + result[i * 128 + l + 32] = d * self.scales[i * 8 + is + 2] as f32 * q2 as f32; + result[i * 128 + l + 64] = d * self.scales[i * 8 + is + 4] as f32 * q3 as f32; + result[i * 128 + l + 96] = d * self.scales[i * 8 + is + 6] as f32 * q4 as f32; + } + + i += 1; + } + + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L130-L135 +pub struct BlockQ5K { + pub d: u16, // super-block scale + pub dmin: u16, // super-block scale for quantized mins + pub scales: [u8; K_SCALE_SIZE], // scales and mins, quantized with 6 bits + pub qh: [u8; QK_K / 8], // quants, high bit + pub qs: [u8; QK_K / 2], // quants, low 4 bits +} + +impl QuantizedValue for BlockQ5K { + const DEQUANTIZED_LEN: usize = QK_K; +} + +#[cfg(feature = "rand")] +impl Distribution for StandardUniform { + fn sample(&self, rng: &mut R) -> BlockQ5K { + // TODO: are all bit representations valid? + BlockQ5K { + d: rng.random(), + dmin: rng.random(), + scales: rng.random(), + qh: rng.random(), + qs: [0; QK_K / 2].map(|_| rng.random()), + } + } +} + +impl BlockQ5K { + pub fn dequantize(self) -> [f32; QK_K] { + let mut result = [0.0; QK_K]; + let mut iq = 0; + let mut is = 0; + + let d = decode_f16(self.d); + let min = decode_f16(self.dmin); + + let mut sc = 0; + let mut m = 0; + let mut u1 = 1; + let mut u2 = 2; + + for j in (0..QK_K).step_by(64) { + get_scale_min_k4(is, &self.scales, &mut sc, &mut m); + let d1 = d * sc as f32; + let m1 = min * m as f32; + get_scale_min_k4(is + 1, &self.scales, &mut sc, &mut m); + let d2 = d * sc as f32; + let m2 = min * m as f32; + + for l in 0..32 { + result[j + l] = d1 + * ((self.qs[iq + l] & 0xF) + (if self.qh[l] & u1 != 0 { 16 } else { 0 })) + as f32 + - m1; + } + + for l in 0..32 { + result[j + l + 32] = d2 + * ((self.qs[iq + l] >> 4) + (if self.qh[l] & u2 != 0 { 16 } else { 0 })) as f32 + - m2; + } + + iq += 32; + is += 2; + u1 <<= 2; + u2 <<= 2; + } + + result + } +} + +#[derive(bytemuck::Pod, bytemuck::Zeroable, Copy, Clone, Debug, PartialEq)] +#[repr(C)] +// See https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.h#L109-L113 +pub struct BlockQ4K { + pub d: u16, // super-block scales for quantized scales + pub dmin: u16, // super-block scale for quantized mins + pub scales: [u8; K_SCALE_SIZE], // scales and mins, quantized with 6 bits + pub qs: [u8; QK_K / 2], // 4-bit quants +} + +impl QuantizedValue for BlockQ4K { + const DEQUANTIZED_LEN: usize = QK_K; +} + +#[cfg(feature = "rand")] +impl Distribution for StandardUniform { + fn sample(&self, rng: &mut R) -> BlockQ4K { + // TODO: are all bit representations valid? + BlockQ4K { + d: rng.random(), + dmin: rng.random(), + scales: rng.random(), + qs: [0; QK_K / 2].map(|_| rng.random()), + } + } +} + +impl BlockQ4K { + // See https://github.com/ggerganov/ggml/blob/a3c0188a4b5d3dec052ff87c9f773baa53631d70/src/ggml-quants.c#L2548 + pub fn dequantize(self) -> [f32; QK_K] { + let mut result = [0.0; QK_K]; + let d = decode_f16(self.d); + let min = decode_f16(self.dmin); + + let mut is = 0; + let mut sc = 0u8; + let mut m = 0u8; + let mut iq = 0; + + for j in (0..QK_K).step_by(64) { + get_scale_min_k4(is, &self.scales, &mut sc, &mut m); + let d1 = d * sc as f32; + let m1 = min * m as f32; + get_scale_min_k4(is + 1, &self.scales, &mut sc, &mut m); + let d2 = d * sc as f32; + let m2 = min * m as f32; + + for l in 0..32 { + result[j + l] = d1 * (self.qs[iq + l] & 0xF) as f32 - m1; + } + + for l in 0..32 { + result[j + l + 32] = d2 * (self.qs[iq + l] >> 4) as f32 - m2; + } + + iq += 32; + is += 2; + } + + result + } +} + +fn get_scale_min_k4(j: usize, q: &[u8], d: &mut u8, m: &mut u8) { + if j < 4 { + *d = q[j] & 63; + *m = q[j + 4] & 63; + } else { + *d = (q[j + 4] & 0xf) | ((q[j - 4] >> 6) << 4); + *m = (q[j + 4] >> 4) | ((q[j] >> 6) << 4); + } +} + +// From https://stackoverflow.com/questions/36008434/how-can-i-decode-f16-to-f32-using-only-the-stable-standard-library +pub fn decode_f16(half: u16) -> f32 { + let exp: u16 = (half >> 10) & 0x1f; + let mant: u16 = half & 0x3ff; + let val: f32 = if exp == 0 { + (mant as f32) * (2.0f32).powi(-24) + } else if exp != 31 { + (mant as f32 + 1024f32) * (2.0f32).powi(exp as i32 - 25) + } else if mant == 0 { + f32::INFINITY + } else { + f32::NAN + }; + if half & 0x8000 != 0 { + -val + } else { + val + } +} + +pub fn decode_bf16(half: u16) -> f32 { + f32::from_bits((half as u32) << 16) +} diff --git a/src/ml/quantized_matrix.rs b/src/ml/quantized_matrix.rs new file mode 100644 index 0000000..aa845cd --- /dev/null +++ b/src/ml/quantized_matrix.rs @@ -0,0 +1,118 @@ +use crate::ml::{ + GpuBlockQ4K, GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ5K, GpuBlockQ5_0x2, GpuBlockQ5_1x2, + GpuBlockQ6Kx2, GpuBlockQ8K, GpuBlockQ8_0x2, +}; +use khal::backend::{GpuDispatch, ShaderBinding}; +use khal::shader::ShaderArgsError; +use khal::ShaderArgs; +use crate::shapes::TensorLayout; +use crate::tensor::Tensor; + +pub enum GpuQuantTensor { + F32(Tensor), + Q8_0(Tensor), + Q5_0(Tensor), + Q5_1(Tensor), + Q4_0(Tensor), + Q4_1(Tensor), + Q8K(Tensor), + Q6K(Tensor), + Q5K(Tensor), + Q4K(Tensor), +} + +impl GpuQuantTensor { + pub fn rank(&self) -> u32 { + match self { + GpuQuantTensor::F32(x) => x.rank(), + GpuQuantTensor::Q8_0(x) => x.rank(), + GpuQuantTensor::Q5_0(x) => x.rank(), + GpuQuantTensor::Q5_1(x) => x.rank(), + GpuQuantTensor::Q4_0(x) => x.rank(), + GpuQuantTensor::Q4_1(x) => x.rank(), + GpuQuantTensor::Q8K(x) => x.rank(), + GpuQuantTensor::Q6K(x) => x.rank(), + GpuQuantTensor::Q5K(x) => x.rank(), + GpuQuantTensor::Q4K(x) => x.rank(), + } + } + + pub fn layout(&self) -> TensorLayout { + match self { + GpuQuantTensor::F32(x) => x.layout(), + GpuQuantTensor::Q8_0(x) => x.layout(), + GpuQuantTensor::Q5_0(x) => x.layout(), + GpuQuantTensor::Q5_1(x) => x.layout(), + GpuQuantTensor::Q4_0(x) => x.layout(), + GpuQuantTensor::Q4_1(x) => x.layout(), + GpuQuantTensor::Q8K(x) => x.layout(), + GpuQuantTensor::Q6K(x) => x.layout(), + GpuQuantTensor::Q5K(x) => x.layout(), + GpuQuantTensor::Q4K(x) => x.layout(), + } + } +} + +impl<'b> ShaderArgs<'b> for GpuQuantTensor { + fn write_arg<'a>( + &'b self, + binding: ShaderBinding, + dispatch: &mut GpuDispatch<'a>, + ) -> Result<(), ShaderArgsError> + where + 'b: 'a, + { + match self { + GpuQuantTensor::F32(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q8_0(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q5_0(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q5_1(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q4_0(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q4_1(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q8K(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q6K(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q5K(matrix) => matrix.buffer().write_arg(binding, dispatch), + GpuQuantTensor::Q4K(matrix) => matrix.buffer().write_arg(binding, dispatch), + } + } +} + +macro_rules! impl_from( + ($($variant: ident, $scalar: ident);*) => {$( + impl From> for GpuQuantTensor { + fn from(value: Tensor<$scalar>) -> Self { + Self::$variant(value) + } + } + )*} +); + +impl_from!( + F32, f32; + Q8_0, GpuBlockQ8_0x2; + Q5_0, GpuBlockQ5_0x2; + Q5_1, GpuBlockQ5_1x2; + Q4_0, GpuBlockQ4_0x2; + Q4_1, GpuBlockQ4_1x2; + Q8K, GpuBlockQ8K; + Q6K, GpuBlockQ6Kx2; + Q5K, GpuBlockQ5K; + Q4K, GpuBlockQ4K +); + +impl GpuQuantTensor { + pub fn shape(&self) -> TensorLayout { + match self { + Self::F32(m) => m.as_view().layout(), + Self::Q8_0(m) => m.as_view().layout(), + Self::Q5_0(m) => m.as_view().layout(), + Self::Q5_1(m) => m.as_view().layout(), + Self::Q4_0(m) => m.as_view().layout(), + Self::Q4_1(m) => m.as_view().layout(), + Self::Q8K(m) => m.as_view().layout(), + Self::Q6K(m) => m.as_view().layout(), + Self::Q5K(m) => m.as_view().layout(), + Self::Q4K(m) => m.as_view().layout(), + } + } +} diff --git a/src/ml/reduce_axis.rs b/src/ml/reduce_axis.rs new file mode 100644 index 0000000..069ef36 --- /dev/null +++ b/src/ml/reduce_axis.rs @@ -0,0 +1,269 @@ +//! Axis-based reduction operations (ReduceSum, ReduceMean, etc.) + +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; + +/// Type of reduction operation. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum ReduceOp { + Sum, + Mean, + Max, + Min, +} + +/// Shader for axis-based reduction operations. +#[derive(Shader)] +pub struct ReduceAxis { + pub reduce_sum_axis: vortx_shaders::ml::reduce_axis::ReduceSumAxis, + pub reduce_mean_axis: vortx_shaders::ml::reduce_axis::ReduceMeanAxis, + pub reduce_max_axis: vortx_shaders::ml::reduce_axis::ReduceMaxAxis, + pub reduce_min_axis: vortx_shaders::ml::reduce_axis::ReduceMinAxis, +} + +impl ReduceAxis { + /// Launch a reduction operation along a single axis. + /// + /// `axis` is the axis to reduce along (0-3). + /// `reduce_size` is the size of the dimension being reduced. + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + op: ReduceOp, + mut dest: impl AsTensorMut, + src: impl AsTensorRef, + axis: u32, + reduce_size: u32, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src = src.as_tensor_ref(); + let len = dest.len() as u32; + let max_threads = 65535u32; + + // Upload params [axis, reduce_size] + let params_buf = TensorBuilder::scalar(BufferUsages::STORAGE | BufferUsages::COPY_DST) + .build_init(backend, &[axis, reduce_size])?; + + macro_rules! dispatch_reduce { + ($wrapper:expr) => {{ + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src.layout())?; + + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src = shapes.get(src.layout()).unwrap(); + let mut buf_dest = dest.buffer_mut(); + + $wrapper.call( + pass, + [len.min(max_threads), 1, 1], + &shape_dest.as_slice(), + &shape_src.as_slice(), + &mut buf_dest, + &src.buffer(), + ¶ms_buf.buffer().as_slice(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = crate::shaders::linalg::Shapes2 { + shape_a: dest.layout().into(), + shape_b: src.layout().into(), + }; + let mut buf_dest = dest.buffer_mut(); + + $wrapper.call( + pass, + [len.min(max_threads), 1, 1], + &mut buf_dest, + &src.buffer(), + ¶ms_buf.buffer().as_slice(), + shapes_val, + ) + } + }}; + } + + match op { + ReduceOp::Sum => dispatch_reduce!(&self.reduce_sum_axis), + ReduceOp::Mean => dispatch_reduce!(&self.reduce_mean_axis), + ReduceOp::Max => dispatch_reduce!(&self.reduce_max_axis), + ReduceOp::Min => dispatch_reduce!(&self.reduce_min_axis), + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; + use khal::{BufferUsages, Shader}; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + async fn test_reduce_sum_axis_generic(backend: &GpuBackend) { + let reduce = ReduceAxis::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + // Input: 3x4 matrix + let src_data: Vec = vec![ + 1.0, 2.0, 3.0, 4.0, // row 0: sum=10 + 5.0, 6.0, 7.0, 8.0, // row 1: sum=26 + 9.0, 10.0, 11.0, 12.0, // row 2: sum=42 + ]; + let src = Tensor::matrix(backend, 3, 4, &src_data, BufferUsages::STORAGE).unwrap(); + + // Reduce along axis 1 (columns) → output shape [3, 1] + let mut dest = Tensor::::vector_uninit( + backend, + 3, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + reduce + .launch( + backend, + &mut shapes, + &mut pass, + ReduceOp::Sum, + &mut dest, + &src, + 1, + 4, + ) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + let mut result = vec![0.0f32; 3]; + backend + .slow_read_buffer(dest.buffer(), &mut result) + .await + .unwrap(); + + let expected = [10.0, 26.0, 42.0]; + for (i, (a, e)) in result.iter().zip(expected.iter()).enumerate() { + assert!( + (a - e).abs() < 1e-4, + "ReduceSum row {i}: gpu={a} expected={e}" + ); + } + } + + async fn test_reduce_max_axis_generic(backend: &GpuBackend) { + let reduce = ReduceAxis::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + let src_data: Vec = vec![ + 3.0, 1.0, 4.0, 1.0, // row 0: max=4 + 5.0, 9.0, 2.0, 6.0, // row 1: max=9 + ]; + let src = Tensor::matrix(backend, 2, 4, &src_data, BufferUsages::STORAGE).unwrap(); + + let mut dest = Tensor::::vector_uninit( + backend, + 2, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + reduce + .launch( + backend, + &mut shapes, + &mut pass, + ReduceOp::Max, + &mut dest, + &src, + 1, + 4, + ) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + let mut result = vec![0.0f32; 2]; + backend + .slow_read_buffer(dest.buffer(), &mut result) + .await + .unwrap(); + + assert!( + (result[0] - 4.0).abs() < 1e-4, + "ReduceMax row 0: {}", + result[0] + ); + assert!( + (result[1] - 9.0).abs() < 1e-4, + "ReduceMax row 1: {}", + result[1] + ); + } + + #[futures_test::test] + #[serial_test::serial] + async fn reduce_sum_axis_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_reduce_sum_axis_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn reduce_sum_axis_cpu() { + let backend = GpuBackend::Cpu; + test_reduce_sum_axis_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn reduce_sum_axis_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_reduce_sum_axis_generic(&backend).await; + } + + #[futures_test::test] + #[serial_test::serial] + async fn reduce_max_axis_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_reduce_max_axis_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn reduce_max_axis_cpu() { + let backend = GpuBackend::Cpu; + test_reduce_max_axis_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn reduce_max_axis_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_reduce_max_axis_generic(&backend).await; + } +} diff --git a/src/ml/rms_norm.rs b/src/ml/rms_norm.rs new file mode 100644 index 0000000..77a9044 --- /dev/null +++ b/src/ml/rms_norm.rs @@ -0,0 +1,174 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::{DVector, Dyn, Storage, Vector}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; + +#[derive(Shader)] +/// Shader implementing the RMS norm kernel. +pub struct RmsNorm { + pub rms_norm: vortx_shaders::ml::rms_norm::RmsNorm, +} + +pub use vortx_shaders::ml::rms_norm::RmsNormConfig; + +impl RmsNorm { + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + config: &Tensor, + mut result: impl AsTensorMut, + value: impl AsTensorRef, + weight: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let value = value.as_tensor_ref().canonicalize(); + let weight = weight.as_tensor_ref().canonicalize(); + let mut result = result.as_tensor_mut().canonicalize(); + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, value.layout())?; + shapes.insert(backend, weight.layout())?; + shapes.insert(backend, result.layout())?; + let shape_v = shapes.get(value.layout()).unwrap(); + let shape_w = shapes.get(weight.layout()).unwrap(); + let shape_out = shapes.get(result.layout()).unwrap(); + let mut out_buf = result.buffer_mut(); + + self.rms_norm.call( + pass, + [1; 3], + &shape_v.as_slice(), + &shape_w.as_slice(), + &shape_out.as_slice(), + &value.buffer(), + &weight.buffer(), + &mut out_buf, + &config.buffer().as_slice(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = value.layout().into(); + let mut out_buf = result.buffer_mut(); + + self.rms_norm.call( + pass, + [1; 3], + &value.buffer(), + &weight.buffer(), + &mut out_buf, + &config.buffer().as_slice(), + shapes_val, + ) + } + } + + pub fn run_cpu>( + out: &mut DVector, + a: &DVector, + w: &Vector, + ) { + const NUDGE_FACTOR: f32 = 1.0e-5; + let rms = 1.0 / (a.norm_squared() / (a.nrows() as f32) + NUDGE_FACTOR).sqrt(); + out.zip_zip_apply(a, w, |o, a, w| *o = (a * rms) * w); + } +} + +#[cfg(test)] +mod test { + use crate::ml::{RmsNorm, RmsNormConfig}; + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_rms_norm_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_rms_norm_generic(&backend).await; + } + + async fn gpu_rms_norm_generic(backend: &GpuBackend) { + let rmsnorm = super::RmsNorm::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + const LEN: u32 = 1757; + + let result = DVector::new_random(LEN as usize); + let value = DVector::new_random(LEN as usize); + let weight = DVector::new_random(LEN as usize); + let mut gpu_result_read = DVector::zeros(LEN as usize); + + let mut gpu_result = Tensor::vector( + backend, + &result, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + let gpu_value = Tensor::vector(backend, &value, BufferUsages::STORAGE).unwrap(); + let gpu_weight = Tensor::vector(backend, &weight, BufferUsages::STORAGE).unwrap(); + let config = Tensor::scalar( + backend, + RmsNormConfig { + nudge_factor: 1.0e-6, + }, + BufferUsages::UNIFORM | BufferUsages::STORAGE, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + rmsnorm + .launch( + backend, + &mut shapes, + &mut pass, + &config, + &mut gpu_result, + gpu_value.as_view(), + gpu_weight.as_view(), + ) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_result.buffer(), gpu_result_read.as_mut_slice()) + .await + .unwrap(); + + let mut cpu_result = result; + RmsNorm::run_cpu(&mut cpu_result, &value, &weight); + + approx::assert_relative_eq!(gpu_result_read, cpu_result, epsilon = 1.0e-4); + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_rms_norm_cpu() { + let backend = GpuBackend::Cpu; + gpu_rms_norm_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_rms_norm_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_rms_norm_generic(&backend).await; + } +} diff --git a/src/ml/rope.rs b/src/ml/rope.rs new file mode 100644 index 0000000..a4d5573 --- /dev/null +++ b/src/ml/rope.rs @@ -0,0 +1,246 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::{vector, DVector, DVectorViewMut, Rotation2}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, Tensor}; + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum RoPEVariant { + // The original version of RoPE, where the rotated entries are adjacent. + Original, + // A variant of RoPE where the rotated entries are separated by `head_size / 2` elements. + Neox, +} + +#[derive(Shader)] +/// Shader implementing the Rotary Positional Encoding kernel. +pub struct RoPE { + pub rope_neox: vortx_shaders::ml::rope::RopeNeox, + pub rope: vortx_shaders::ml::rope::Rope, +} + +pub use vortx_shaders::ml::rope::RoPEConfig; + +impl RoPE { + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + variant: RoPEVariant, + config: &Tensor, + mut in_out_q: impl AsTensorMut, + mut in_out_k: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut in_out_q = in_out_q.as_tensor_mut(); + let mut in_out_k = in_out_k.as_tensor_mut(); + + assert_eq!(in_out_q.len() % 2, 0); + assert_eq!(in_out_k.len() % 2, 0); + assert!( + in_out_q.len() >= in_out_k.len(), + "The Query vector must be larger than, or as large as, the Key vector." + ); + + let num_threads = [in_out_q.len() as u32 / 2, 1, 1]; + + macro_rules! dispatch_rope { + ($wrapper:expr) => {{ + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, in_out_q.layout()).unwrap(); + shapes.insert(backend, in_out_k.layout()).unwrap(); + let shape_q = shapes.get(in_out_q.layout()).unwrap(); + let shape_k = shapes.get(in_out_k.layout()).unwrap(); + let mut buf_q = in_out_q.buffer_mut(); + let mut buf_k = in_out_k.buffer_mut(); + + $wrapper.call( + pass, + num_threads, + &shape_q.as_slice(), + &shape_k.as_slice(), + &config.buffer().as_slice(), + &mut buf_q, + &mut buf_k, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = crate::shaders::linalg::Shapes2 { + shape_a: in_out_q.layout().into(), + shape_b: in_out_k.layout().into(), + }; + let mut buf_q = in_out_q.buffer_mut(); + let mut buf_k = in_out_k.buffer_mut(); + + $wrapper.call( + pass, + num_threads, + &config.buffer().as_slice(), + &mut buf_q, + &mut buf_k, + shapes_val, + ) + } + }}; + } + + match variant { + RoPEVariant::Original => dispatch_rope!(&self.rope), + RoPEVariant::Neox => dispatch_rope!(&self.rope_neox), + } + } + + // Rotary Positional Encoding (RoPE): complex-valued rotate q and k in each head. + pub fn run_cpu( + q: &mut DVector, + k: &mut DVectorViewMut, + head_size: usize, + dim: usize, + kv_dim: usize, + pos: usize, + ) { + for i in (0..dim).step_by(2) { + // For RoPE, we have one rotation matrix like https://youtu.be/Mn_9W1nCFLo?si=GLIXuFLGVG8q6v2u&t=1963 + // for each head. So we need to transform `i` into the corresponding index within + // the head. + let head_dim = (i % head_size) as f32; + // Not that the formulae from the video linked above would be: + // 10000.0.powf(-2.0 * ((i / 2) as f32 - 1.0) / dim as f32) + // Although in the paper shown in the video, their index is 1-based which his why thy + // have to subtract 1.0 whereas we don't need to.The `i / 2` and multiplication by 2.0 + // are both accounted for by stepping only on even values for `i`. + // Therefore, the formulae below is equivalent to the RoPE paper's formulae. + let theta = 10000.0_f32.powf(-head_dim / head_size as f32); + let m_theta = pos as f32 * theta; + let rot = Rotation2::new(m_theta); + + let qi = vector![q[i], q[i + 1]]; + let mut out_q = q.fixed_rows_mut::<2>(i); + out_q.copy_from(&(rot * qi)); + + // When i >= kv_dim, we are done rotating all the elements from the keys. That's + // because there are less key heads than query heads, but each key head sub-vector has + // the same dimension as the query head (they loose dimension when multiplied with the + // key weight matrices). + if i < kv_dim { + let ki = vector![k[i], k[i + 1]]; + let mut out_k = k.fixed_rows_mut::<2>(i); + out_k.copy_from(&(rot * ki)); + } + } + } +} + +#[cfg(test)] +mod test { + use super::RoPEConfig; + use crate::ml::{RoPE, RoPEVariant}; + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_rope_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_rope_generic(&backend).await; + } + + async fn gpu_rope_generic(backend: &GpuBackend) { + let rope = super::RoPE::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + const HEAD_SIZE: u32 = 128; + const LEN_Q: u32 = 13 * HEAD_SIZE; + const LEN_K: u32 = 9 * HEAD_SIZE; + + let rope_indices = RoPEConfig { + head_size: HEAD_SIZE, + kv_dim: LEN_K, + pos: 10, + base_freq: 1.0e4, + }; + + let mut q = DVector::new_random(LEN_Q as usize); + let mut k = DVector::new_random(LEN_K as usize); + let mut result_q = DVector::zeros(LEN_Q as usize); + let mut result_k = DVector::zeros(LEN_K as usize); + + let gpu_indices = Tensor::scalar( + backend, + rope_indices, + BufferUsages::UNIFORM | BufferUsages::STORAGE, + ) + .unwrap(); + let mut gpu_q = + Tensor::vector(backend, &q, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + let mut gpu_k = + Tensor::vector(backend, &k, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("rope_test", None); + rope.launch( + backend, + &mut shapes, + &mut pass, + RoPEVariant::Original, + &gpu_indices, + &mut gpu_q, + &mut gpu_k, + ) + .unwrap(); + drop(pass); + + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_q.buffer(), result_q.as_mut_slice()) + .await + .unwrap(); + backend + .slow_read_buffer(gpu_k.buffer(), result_k.as_mut_slice()) + .await + .unwrap(); + + RoPE::run_cpu( + &mut q, + &mut k.rows_mut(0, LEN_K as usize), + rope_indices.head_size as usize, + LEN_Q as usize, + rope_indices.kv_dim as usize, + rope_indices.pos as usize, + ); + + // TODO: why is the epsilon so high? Is it a difference in sin/cos implementations? + approx::assert_relative_eq!(result_q, q, epsilon = 1.0e-5); + approx::assert_relative_eq!(result_k, k, epsilon = 1.0e-5); + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_rope_cpu() { + let backend = GpuBackend::Cpu; + gpu_rope_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_rope_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_rope_generic(&backend).await; + } +} diff --git a/src/ml/select.rs b/src/ml/select.rs new file mode 100644 index 0000000..a687f50 --- /dev/null +++ b/src/ml/select.rs @@ -0,0 +1,147 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +#[derive(Shader)] +pub struct Select { + pub select: vortx_shaders::ml::select::Select, +} + +impl Select { + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut dest: impl AsTensorMut, + src: impl AsTensorRef, + idx: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src = src.as_tensor_ref(); + let idx = idx.as_tensor_ref(); + let len = dest.len() as u32; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, dest.layout())?; + shapes.insert(backend, src.layout())?; + + let shape_dest = shapes.get(dest.layout()).unwrap(); + let shape_src = shapes.get(src.layout()).unwrap(); + let mut buf_dest = dest.buffer_mut(); + + self.select.call( + pass, + [len, 1, 1], + &shape_dest.as_slice(), + &shape_src.as_slice(), + &mut buf_dest, + &src.buffer(), + &idx.buffer(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = crate::shaders::linalg::Shapes2 { + shape_a: dest.layout().into(), + shape_b: src.layout().into(), + }; + let mut buf_dest = dest.buffer_mut(); + + self.select.call( + pass, + [len, 1, 1], + &mut buf_dest, + &src.buffer(), + &idx.buffer(), + shapes_val, + ) + } + } +} + +#[cfg(test)] +mod test { + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; + use khal::{BufferUsages, Shader}; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + /// Select rows from a matrix by index: dest[i] = src[idx[i], :]. + async fn test_select_generic(backend: &GpuBackend) { + let select = super::Select::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + // Source matrix 8x4 + let src_data: Vec = (0..32).map(|i| i as f32).collect(); + let src = Tensor::matrix(backend, 8, 4, &src_data, BufferUsages::STORAGE).unwrap(); + + // Index vector: pick rows [2, 0, 5, 7] + let idx_data: Vec = vec![2, 0, 5, 7]; + let idx = Tensor::vector(backend, &idx_data, BufferUsages::STORAGE).unwrap(); + + // Destination matrix 4x4 + let mut dest = Tensor::::matrix_uninit( + backend, + 4, + 4, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + select + .launch(backend, &mut shapes, &mut pass, &mut dest, &src, &idx) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + let mut result = vec![0.0f32; 16]; + backend + .slow_read_buffer(dest.buffer(), &mut result) + .await + .unwrap(); + + // Expected: rows 2, 0, 5, 7 of the source + let expected: Vec = vec![ + 8.0, 9.0, 10.0, 11.0, // row 2 + 0.0, 1.0, 2.0, 3.0, // row 0 + 20.0, 21.0, 22.0, 23.0, // row 5 + 28.0, 29.0, 30.0, 31.0, // row 7 + ]; + assert_eq!(result, expected); + } + + #[futures_test::test] + #[serial_test::serial] + async fn select_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + test_select_generic(&backend).await; + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn select_cpu() { + let backend = GpuBackend::Cpu; + test_select_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn select_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + test_select_generic(&backend).await; + } +} diff --git a/src/ml/silu.rs b/src/ml/silu.rs new file mode 100644 index 0000000..34bcfbf --- /dev/null +++ b/src/ml/silu.rs @@ -0,0 +1,141 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::DVector; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef}; + +#[derive(Shader)] +/// Shader implementing the Silu activation function. +pub struct Silu { + pub silu: vortx_shaders::ml::silu::Silu, +} + +impl Silu { + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut in_out_h1: impl AsTensorMut, + in_h2: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut h1 = in_out_h1.as_tensor_mut().canonicalize(); + let h2 = in_h2.as_tensor_ref().canonicalize(); + let len = h1.len() as u32; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, h1.layout())?; + shapes.insert(backend, h2.layout())?; + let shape_a = shapes.get(h1.layout()).unwrap(); + let shape_b = shapes.get(h2.layout()).unwrap(); + let mut in_out_a = h1.buffer_mut(); + + self.silu.call( + pass, + [len, 1, 1], + &shape_a.as_slice(), + &shape_b.as_slice(), + &mut in_out_a, + &h2.buffer(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = h1.layout().into(); + let mut in_out_a = h1.buffer_mut(); + + self.silu + .call(pass, [len, 1, 1], &mut in_out_a, &h2.buffer(), shapes_val) + } + } + + pub fn run_cpu(h1: &mut DVector, h2: &DVector) { + // SwiGLU non-linearity. + fn swish(x: f32, beta: f32) -> f32 { + // This is the swish function from https://youtu.be/Mn_9W1nCFLo?si=LT6puSAfzgpP6ydz&t=3973 + x / (1.0 + (-beta * x).exp()) + } + + h1.zip_apply(h2, |h, h2| *h = h2 * swish(*h, 1.0)); + } +} + +#[cfg(test)] +mod test { + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_silu_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_silu_generic(&backend).await; + } + + async fn gpu_silu_generic(backend: &GpuBackend) { + let silu = super::Silu::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + const LEN: u32 = 1757; + + let h1 = DVector::new_random(LEN as usize); + let h2 = DVector::new_random(LEN as usize); + let mut h1_read = DVector::zeros(LEN as usize); + + let mut gpu_h1 = + Tensor::vector(backend, &h1, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + let gpu_h2 = Tensor::vector(backend, &h2, BufferUsages::STORAGE).unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("silu_test", None); + silu.launch( + backend, + &mut shapes, + &mut pass, + &mut gpu_h1, + gpu_h2.as_view(), + ) + .unwrap(); + drop(pass); + + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_h1.buffer(), h1_read.as_mut_slice()) + .await + .unwrap(); + + let mut cpu_result = h1; + super::Silu::run_cpu(&mut cpu_result, &h2); + + approx::assert_relative_eq!(h1_read, cpu_result, epsilon = 1.0e-5); + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_silu_cpu() { + let backend = GpuBackend::Cpu; + gpu_silu_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_silu_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_silu_generic(&backend).await; + } +} diff --git a/src/ml/softmax.rs b/src/ml/softmax.rs new file mode 100644 index 0000000..b3e7ea2 --- /dev/null +++ b/src/ml/softmax.rs @@ -0,0 +1,214 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::{Dyn, StorageMut, Vector}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::AsTensorMut; + +/* +layout (push_constant) uniform parameter +{ + uint KX; + uint KY; + uint ne00; + uint ne01; + uint ne02; + uint ne12; + uint ne13; + uint nb11; + uint nb12; + uint nb13; + float scale; + float max_bias; + float m0; + float m1; + uint n_head_log2; + uint nrows_x; + uint has_sinks; +} p; + +#include "types.glsl" + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer Y {B_TYPE data_b[];}; +layout (binding = 2) readonly buffer Z {float data_c[];}; +layout (binding = 3) buffer D {D_TYPE data_d[];}; + +struct SoftMaxGgmlUniforms { + +} + +struct SoftMaxGgmlArgs<'a> { + +} + */ + +#[derive(Shader)] +/// Shader implementing the softmax kernel. +pub struct SoftMax { + pub softmax: vortx_shaders::ml::softmax::Softmax, + pub log_softmax: vortx_shaders::ml::softmax::LogSoftmax, +} + +impl SoftMax { + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut in_out_mat: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut in_out_mat = in_out_mat.as_tensor_mut().canonicalize(); + let size = in_out_mat.layout().size; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, in_out_mat.layout())?; + let shape_buf = shapes.get(in_out_mat.layout()).unwrap(); + let mut mat_buf = in_out_mat.buffer_mut(); + + self.softmax.call( + pass, + [size[2], size[1], size[0]], + &shape_buf.as_slice(), + &mut mat_buf, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = in_out_mat.layout().into(); + let mut mat_buf = in_out_mat.buffer_mut(); + + self.softmax + .call(pass, [size[2], size[1], size[0]], &mut mat_buf, shapes_val) + } + } + + pub fn launch_log( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut in_out_mat: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let mut in_out_mat = in_out_mat.as_tensor_mut().canonicalize(); + let size = in_out_mat.layout().size; + + #[cfg(not(feature = "push_constants"))] + { + shapes.insert(backend, in_out_mat.layout())?; + let shape_buf = shapes.get(in_out_mat.layout()).unwrap(); + let mut mat_buf = in_out_mat.buffer_mut(); + + self.log_softmax.call( + pass, + [size[2], size[1], size[0]], + &shape_buf.as_slice(), + &mut mat_buf, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = in_out_mat.layout().into(); + let mut mat_buf = in_out_mat.buffer_mut(); + + self.log_softmax + .call(pass, [size[2], size[1], size[0]], &mut mat_buf, shapes_val) + } + } + + /// The softmax function. + /// + /// Converts a set of real number into a probability distribution. + /// See + pub fn run_cpu>(vals: &mut Vector) { + // Note that llama2.c also introduces a bias based on the max value + // to improve numerical stability. So it is effectively computing: + // softmax(z) = (e^z - max) / (e^z - max).sum() + let max_val = vals.max(); + let mut sum = 0.0; + + vals.apply(|x| { + *x = (*x - max_val).exp(); + sum += *x; + }); + + *vals /= sum; + } +} + +#[cfg(test)] +mod test { + use crate::ml::SoftMax; + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_softmax_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_softmax_generic(&backend).await; + } + + async fn gpu_softmax_generic(backend: &GpuBackend) { + let softmax = super::SoftMax::from_backend(backend).unwrap(); + let mut shapes = TensorLayoutBuffers::new(backend); + + const LEN: u32 = 1757; + + let v0 = DVector::from_fn(LEN as usize, |i, _| i as f32); + let mut gpu_v0_read = DVector::zeros(LEN as usize); + let mut gpu_v0 = + Tensor::vector(backend, &v0, BufferUsages::STORAGE | BufferUsages::COPY_SRC).unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + softmax + .launch(backend, &mut shapes, &mut pass, &mut gpu_v0) + .unwrap(); + drop(pass); + + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_v0.buffer(), gpu_v0_read.as_mut_slice()) + .await + .unwrap(); + + let mut cpu_result = v0; + SoftMax::run_cpu(&mut cpu_result); + + approx::assert_relative_eq!(gpu_v0_read, cpu_result, epsilon = 1.0e-7); + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_softmax_cpu() { + let backend = GpuBackend::Cpu; + gpu_softmax_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_softmax_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_softmax_generic(&backend).await; + } +} diff --git a/src/ml/unary.rs b/src/ml/unary.rs new file mode 100644 index 0000000..1e3ac77 --- /dev/null +++ b/src/ml/unary.rs @@ -0,0 +1,1060 @@ +use khal_std::glamx::Vec4; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use nalgebra::{Dyn, StorageMut, Vector}; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +#[non_exhaustive] +/// Listing of all unary operations that can be applied by the [`Unary`] kernel. +pub enum UnaryOp { + Abs, + Sgn, + Neg, + Step, + Elu, + Gelu, + GeluQuick, + Silu, + Tanh, + Sin, + Cos, + Relu, + Sigmoid, + HardSigmoid, + // HardSwish, + Sqr, + Sqrt, + Log, + Exp, + Reciprocal, + Erf, + // Unary ops with extra args. + LeakyRelu, + Clamp, + Scale, + AddScalar, // Named GGML_OP_ADD1 in ggml. + Pow, +} + +impl UnaryOp { + const fn has_args(self) -> bool { + match self { + Self::Abs + | Self::Sgn + | Self::Neg + | Self::Step + | Self::Elu + | Self::Gelu + | Self::GeluQuick + | Self::Silu + | Self::Tanh + | Self::Relu + | Self::Sigmoid + | Self::HardSigmoid + // | Self::HardSwish + | Self::Sqr + | Self::Sqrt + | Self::Log + | Self::Exp + | Self::Reciprocal + | Self::Erf + | Self::Sin + | Self::Cos => false, + Self::LeakyRelu | Self::Clamp | Self::Scale | Self::AddScalar | Self::Pow => true, + } + } + + pub fn eval(self, x: f32, args: Vec4) -> f32 { + match self { + Self::Abs => x.abs(), + Self::Sgn => x.signum(), + Self::Neg => -x, + Self::Step => { + if x > 0.0 { + 1.0 + } else { + 0.0 + } + } + Self::Elu => { + if x > 0.0 { + x + } else { + x.exp() - 1.0 + } + } + Self::Gelu => { + const GELU_COEF_A: f32 = 0.044715; + const SQRT_2_OVER_PI: f32 = 0.7978846; + 0.5 * x * (1.0 + (SQRT_2_OVER_PI * x * (1.0 + GELU_COEF_A * x * x)).tanh()) + } + Self::GeluQuick => { + const GELU_QUICK_COEF: f32 = -1.702; + x * (1.0 / (1.0 + (GELU_QUICK_COEF * x).exp())) + } + Self::Silu => x / (1.0 + (-x).exp()), + Self::Tanh => x.tanh(), + Self::Relu => x.max(0.0), + Self::Sigmoid => 1.0 / (1.0 + (-x).exp()), + Self::HardSigmoid => 1.0f32.min(0.0f32.max((x + 3.0) / 6.0)), + // Self::HardSwish => x * 1.0f32.min(0.0f32.max((x + 3.0) / 6.0)), + Self::Sqr => x * x, + Self::Sqrt => x.sqrt(), + Self::Sin => x.sin(), + Self::Cos => x.cos(), + Self::Log => x.ln(), + Self::Exp => x.exp(), + Self::Reciprocal => 1.0 / x, + Self::Erf => { + // Abramowitz and Stegun approximation + #[allow(clippy::excessive_precision)] + let a1: f32 = 0.254829592; + #[allow(clippy::excessive_precision)] + let a2: f32 = -0.284496736; + #[allow(clippy::excessive_precision)] + let a3: f32 = 1.421413741; + #[allow(clippy::excessive_precision)] + let a4: f32 = -1.453152027; + #[allow(clippy::excessive_precision)] + let a5: f32 = 1.061405429; + #[allow(clippy::excessive_precision)] + let p: f32 = 0.3275911; + let sign = if x >= 0.0 { 1.0 } else { -1.0 }; + let x = x.abs(); + let t = 1.0 / (1.0 + p * x); + let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp(); + sign * y + } + Self::LeakyRelu => x.max(0.0) + x.min(0.0) * args.x, + Self::Clamp => x.clamp(args.x, args.y), + Self::Scale => x * args.x, + Self::AddScalar => x + args.x, + Self::Pow => x.powf(args.x), + } + } +} + +/// Shader implementing various unary operations selected with [`UnaryOp`]. +#[derive(Shader)] +pub struct Unary { + pub abs_op: vortx_shaders::ml::unary::AbsOp, + pub abs_inplace: vortx_shaders::ml::unary::AbsInplace, + pub sgn_op: vortx_shaders::ml::unary::SgnOp, + pub sgn_inplace: vortx_shaders::ml::unary::SgnInplace, + pub neg_op: vortx_shaders::ml::unary::NegOp, + pub neg_inplace: vortx_shaders::ml::unary::NegInplace, + pub step_op: vortx_shaders::ml::unary::StepOp, + pub step_inplace: vortx_shaders::ml::unary::StepInplace, + pub elu_op: vortx_shaders::ml::unary::EluOp, + pub elu_inplace: vortx_shaders::ml::unary::EluInplace, + pub gelu_op: vortx_shaders::ml::unary::GeluOp, + pub gelu_inplace: vortx_shaders::ml::unary::GeluInplace, + pub gelu_quick_op: vortx_shaders::ml::unary::GeluQuickOp, + pub gelu_quick_inplace: vortx_shaders::ml::unary::GeluQuickInplace, + pub silu_op: vortx_shaders::ml::unary::SiluOp, + pub silu_inplace: vortx_shaders::ml::unary::SiluInplace, + pub tanh_op: vortx_shaders::ml::unary::TanhOp, + pub tanh_inplace: vortx_shaders::ml::unary::TanhInplace, + pub relu_op: vortx_shaders::ml::unary::ReluOp, + pub relu_inplace: vortx_shaders::ml::unary::ReluInplace, + pub sigmoid_op: vortx_shaders::ml::unary::SigmoidOp, + pub sigmoid_inplace: vortx_shaders::ml::unary::SigmoidInplace, + pub hard_sigmoid_op: vortx_shaders::ml::unary::HardSigmoidOp, + pub hard_sigmoid_inplace: vortx_shaders::ml::unary::HardSigmoidInplace, + // pub hard_swish_op: vortx_shaders::ml::unary::HardSwishOp, + // pub hard_swish_inplace: vortx_shaders::ml::unary::HardSwishInplace, + pub sqr_op: vortx_shaders::ml::unary::SqrOp, + pub sqr_inplace: vortx_shaders::ml::unary::SqrInplace, + pub sqrt_op: vortx_shaders::ml::unary::SqrtOp, + pub sqrt_inplace: vortx_shaders::ml::unary::SqrtInplace, + pub log_op: vortx_shaders::ml::unary::LogOp, + pub log_inplace: vortx_shaders::ml::unary::LogInplace, + pub leaky_relu_op: vortx_shaders::ml::unary::LeakyReluOp, + pub leaky_relu_inplace: vortx_shaders::ml::unary::LeakyReluInplace, + pub clamp_op: vortx_shaders::ml::unary::ClampOp, + pub clamp_inplace: vortx_shaders::ml::unary::ClampInplace, + pub scale_op: vortx_shaders::ml::unary::ScaleOp, + pub scale_inplace: vortx_shaders::ml::unary::ScaleInplace, + pub add_scalar_op: vortx_shaders::ml::unary::AddScalarOp, + pub add_scalar_inplace: vortx_shaders::ml::unary::AddScalarInplace, + pub sin_op: vortx_shaders::ml::unary::SinOp, + pub sin_inplace: vortx_shaders::ml::unary::SinInplace, + pub cos_op: vortx_shaders::ml::unary::CosOp, + pub cos_inplace: vortx_shaders::ml::unary::CosInplace, + pub exp_op: vortx_shaders::ml::unary::ExpOp, + pub exp_inplace: vortx_shaders::ml::unary::ExpInplace, + pub reciprocal_op: vortx_shaders::ml::unary::ReciprocalOp, + pub reciprocal_inplace: vortx_shaders::ml::unary::ReciprocalInplace, + pub erf_op: vortx_shaders::ml::unary::ErfOp, + pub erf_inplace: vortx_shaders::ml::unary::ErfInplace, + pub pow_op: vortx_shaders::ml::unary::PowOp, + pub pow_inplace: vortx_shaders::ml::unary::PowInplace, +} + +/// Helper macro for dispatching inplace unary ops without args. +/// Each kernel has its own generated args type with the same fields. +macro_rules! dispatch_inplace_no_args { + ($self:expr, $shapes:expr, $backend:expr, $pass:expr, $src:expr, $len:expr, $kernel:ident, $($args_ty:ident)::+) => {{ + #[cfg(not(feature = "push_constants"))] + { + $shapes.insert($backend, $src.layout())?; + let shape_src = $shapes.get($src.layout()).unwrap(); + let mut src_buf = $src.buffer_mut(); + + $self + .$kernel + .call($pass, [$len, 1, 1], &shape_src.as_slice(), &mut src_buf) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = $src.layout().into(); + let mut src_buf = $src.buffer_mut(); + + $self + .$kernel + .call($pass, [$len, 1, 1], &mut src_buf, shapes_val) + } + }}; +} + +/// Helper macro for dispatching inplace unary ops with args. +macro_rules! dispatch_inplace_with_args { + ($self:expr, $shapes:expr, $backend:expr, $pass:expr, $src:expr, $args:expr, $len:expr, $kernel:ident, $($args_ty:ident)::+) => {{ + #[cfg(not(feature = "push_constants"))] + { + $shapes.insert($backend, $src.layout())?; + let shape_src = $shapes.get($src.layout()).unwrap(); + let mut src_buf = $src.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &shape_src.as_slice(), + &mut src_buf, + $args.unwrap().buffer(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = $src.layout().into(); + let mut src_buf = $src.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &mut src_buf, + $args.unwrap().buffer(), + shapes_val, + ) + } + }}; +} + +/// Helper macro for dispatching non-inplace unary ops without args. +macro_rules! dispatch_op_no_args { + ($self:expr, $shapes:expr, $backend:expr, $pass:expr, $src:expr, $dest:expr, $len:expr, $kernel:ident, $($args_ty:ident)::+) => {{ + #[cfg(not(feature = "push_constants"))] + { + $shapes.insert($backend, $dest.layout())?; + $shapes.insert($backend, $src.layout())?; + let shape_dst = $shapes.get($dest.layout()).unwrap(); + let shape_src = $shapes.get($src.layout()).unwrap(); + let mut dst_buf = $dest.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &shape_src.as_slice(), + &$src.buffer(), + &shape_dst.as_slice(), + &mut dst_buf, + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = $src.layout().into(); + let mut dst_buf = $dest.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &$src.buffer(), + &mut dst_buf, + shapes_val, + ) + } + }}; +} + +/// Helper macro for dispatching non-inplace unary ops with args. +macro_rules! dispatch_op_with_args { + ($self:expr, $shapes:expr, $backend:expr, $pass:expr, $src:expr, $dest:expr, $args:expr, $len:expr, $kernel:ident, $($args_ty:ident)::+) => {{ + #[cfg(not(feature = "push_constants"))] + { + $shapes.insert($backend, $dest.layout())?; + $shapes.insert($backend, $src.layout())?; + let shape_dst = $shapes.get($dest.layout()).unwrap(); + let shape_src = $shapes.get($src.layout()).unwrap(); + let mut dst_buf = $dest.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &shape_src.as_slice(), + &$src.buffer(), + &shape_dst.as_slice(), + &mut dst_buf, + &$args.unwrap().buffer().as_slice(), + ) + } + + #[cfg(feature = "push_constants")] + { + let shapes_val = $src.layout().into(); + let mut dst_buf = $dest.buffer_mut(); + + $self.$kernel.call( + $pass, + [$len, 1, 1], + &$src.buffer(), + &mut dst_buf, + &$args.unwrap().buffer().as_slice(), + shapes_val, + ) + } + }}; +} + +impl Unary { + pub fn launch_inplace( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + op: UnaryOp, + mut src: impl AsTensorMut, + args: Option<&Tensor>, + ) -> Result<(), GpuBackendError> { + let mut src = src.as_tensor_mut(); + let len = src.len() as u32; + + assert_eq!( + op.has_args(), + args.is_some(), + "Unary ops argument mismatch." + ); + + match op { + UnaryOp::Abs => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + abs_inplace, + vortx_shaders::ml::unary::AbsInplaceArgs + ), + UnaryOp::Sgn => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + sgn_inplace, + vortx_shaders::ml::unary::SgnInplaceArgs + ), + UnaryOp::Neg => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + neg_inplace, + vortx_shaders::ml::unary::NegInplaceArgs + ), + UnaryOp::Step => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + step_inplace, + vortx_shaders::ml::unary::StepInplaceArgs + ), + UnaryOp::Elu => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + elu_inplace, + vortx_shaders::ml::unary::EluInplaceArgs + ), + UnaryOp::Gelu => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + gelu_inplace, + vortx_shaders::ml::unary::GeluInplaceArgs + ), + UnaryOp::GeluQuick => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + gelu_quick_inplace, + vortx_shaders::ml::unary::GeluQuickInplaceArgs + ), + UnaryOp::Silu => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + silu_inplace, + vortx_shaders::ml::unary::SiluInplaceArgs + ), + UnaryOp::Tanh => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + tanh_inplace, + vortx_shaders::ml::unary::TanhInplaceArgs + ), + UnaryOp::Relu => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + relu_inplace, + vortx_shaders::ml::unary::ReluInplaceArgs + ), + UnaryOp::Sigmoid => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + sigmoid_inplace, + vortx_shaders::ml::unary::SigmoidInplaceArgs + ), + UnaryOp::HardSigmoid => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + hard_sigmoid_inplace, + vortx_shaders::ml::unary::HardSigmoidInplaceArgs + ), + // UnaryOp::HardSwish => dispatch_inplace_no_args!(self, shapes, backend, pass, src, len, hard_swish_inplace, vortx_shaders::ml::unary::HardSwishInplaceArgs), + UnaryOp::Sqr => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + sqr_inplace, + vortx_shaders::ml::unary::SqrInplaceArgs + ), + UnaryOp::Sqrt => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + sqrt_inplace, + vortx_shaders::ml::unary::SqrtInplaceArgs + ), + UnaryOp::Sin => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + sin_inplace, + vortx_shaders::ml::unary::SinInplaceArgs + ), + UnaryOp::Cos => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + cos_inplace, + vortx_shaders::ml::unary::CosInplaceArgs + ), + UnaryOp::Log => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + log_inplace, + vortx_shaders::ml::unary::LogInplaceArgs + ), + UnaryOp::Exp => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + exp_inplace, + vortx_shaders::ml::unary::ExpInplaceArgs + ), + UnaryOp::Reciprocal => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + reciprocal_inplace, + vortx_shaders::ml::unary::ReciprocalInplaceArgs + ), + UnaryOp::Erf => dispatch_inplace_no_args!( + self, + shapes, + backend, + pass, + src, + len, + erf_inplace, + vortx_shaders::ml::unary::ErfInplaceArgs + ), + UnaryOp::LeakyRelu => dispatch_inplace_with_args!( + self, + shapes, + backend, + pass, + src, + args, + len, + leaky_relu_inplace, + vortx_shaders::ml::unary::LeakyReluInplaceArgs + ), + UnaryOp::Clamp => dispatch_inplace_with_args!( + self, + shapes, + backend, + pass, + src, + args, + len, + clamp_inplace, + vortx_shaders::ml::unary::ClampInplaceArgs + ), + UnaryOp::Scale => dispatch_inplace_with_args!( + self, + shapes, + backend, + pass, + src, + args, + len, + scale_inplace, + vortx_shaders::ml::unary::ScaleInplaceArgs + ), + UnaryOp::AddScalar => dispatch_inplace_with_args!( + self, + shapes, + backend, + pass, + src, + args, + len, + add_scalar_inplace, + vortx_shaders::ml::unary::AddScalarInplaceArgs + ), + UnaryOp::Pow => dispatch_inplace_with_args!( + self, + shapes, + backend, + pass, + src, + args, + len, + pow_inplace, + vortx_shaders::ml::unary::PowInplaceArgs + ), + } + } + + pub fn launch( + &self, + backend: &GpuBackend, + #[cfg_attr(feature = "push_constants", allow(unused_variables))] + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + op: UnaryOp, + mut dest: impl AsTensorMut, + src: impl AsTensorRef, + args: Option<&Tensor>, + ) -> Result<(), GpuBackendError> { + let mut dest = dest.as_tensor_mut(); + let src = src.as_tensor_ref(); + let len = dest.len() as u32; + + assert_eq!( + op.has_args(), + args.is_some(), + "Unary ops argument mismatch." + ); + + match op { + UnaryOp::Abs => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + abs_op, + vortx_shaders::ml::unary::AbsOpArgs + ), + UnaryOp::Sgn => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + sgn_op, + vortx_shaders::ml::unary::SgnOpArgs + ), + UnaryOp::Neg => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + neg_op, + vortx_shaders::ml::unary::NegOpArgs + ), + UnaryOp::Step => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + step_op, + vortx_shaders::ml::unary::StepOpArgs + ), + UnaryOp::Elu => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + elu_op, + vortx_shaders::ml::unary::EluOpArgs + ), + UnaryOp::Gelu => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + gelu_op, + vortx_shaders::ml::unary::GeluOpArgs + ), + UnaryOp::GeluQuick => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + gelu_quick_op, + vortx_shaders::ml::unary::GeluQuickOpArgs + ), + UnaryOp::Silu => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + silu_op, + vortx_shaders::ml::unary::SiluOpArgs + ), + UnaryOp::Tanh => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + tanh_op, + vortx_shaders::ml::unary::TanhOpArgs + ), + UnaryOp::Relu => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + relu_op, + vortx_shaders::ml::unary::ReluOpArgs + ), + UnaryOp::Sigmoid => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + sigmoid_op, + vortx_shaders::ml::unary::SigmoidOpArgs + ), + UnaryOp::HardSigmoid => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + hard_sigmoid_op, + vortx_shaders::ml::unary::HardSigmoidOpArgs + ), + // UnaryOp::HardSwish => dispatch_op_no_args!(self, shapes, backend, pass, src, dest, len, hard_swish_op, vortx_shaders::ml::unary::HardSwishOpArgs), + UnaryOp::Sqr => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + sqr_op, + vortx_shaders::ml::unary::SqrOpArgs + ), + UnaryOp::Sqrt => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + sqrt_op, + vortx_shaders::ml::unary::SqrtOpArgs + ), + UnaryOp::Sin => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + sin_op, + vortx_shaders::ml::unary::SinOpArgs + ), + UnaryOp::Cos => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + cos_op, + vortx_shaders::ml::unary::CosOpArgs + ), + UnaryOp::Log => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + log_op, + vortx_shaders::ml::unary::LogOpArgs + ), + UnaryOp::Exp => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + exp_op, + vortx_shaders::ml::unary::ExpOpArgs + ), + UnaryOp::Reciprocal => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + reciprocal_op, + vortx_shaders::ml::unary::ReciprocalOpArgs + ), + UnaryOp::Erf => dispatch_op_no_args!( + self, + shapes, + backend, + pass, + src, + dest, + len, + erf_op, + vortx_shaders::ml::unary::ErfOpArgs + ), + UnaryOp::LeakyRelu => dispatch_op_with_args!( + self, + shapes, + backend, + pass, + src, + dest, + args, + len, + leaky_relu_op, + vortx_shaders::ml::unary::LeakyReluOpArgs + ), + UnaryOp::Clamp => dispatch_op_with_args!( + self, + shapes, + backend, + pass, + src, + dest, + args, + len, + clamp_op, + vortx_shaders::ml::unary::ClampOpArgs + ), + UnaryOp::Scale => dispatch_op_with_args!( + self, + shapes, + backend, + pass, + src, + dest, + args, + len, + scale_op, + vortx_shaders::ml::unary::ScaleOpArgs + ), + UnaryOp::AddScalar => dispatch_op_with_args!( + self, + shapes, + backend, + pass, + src, + dest, + args, + len, + add_scalar_op, + vortx_shaders::ml::unary::AddScalarOpArgs + ), + UnaryOp::Pow => dispatch_op_with_args!( + self, + shapes, + backend, + pass, + src, + dest, + args, + len, + pow_op, + vortx_shaders::ml::unary::PowOpArgs + ), + } + } + + pub fn run_cpu>( + &self, + op: UnaryOp, + vals: &mut Vector, + args: Vec4, + ) { + vals.apply(|x| *x = op.eval(*x, args)); + } +} + +#[cfg(test)] +#[cfg(feature = "rand")] +mod test { + use crate::ml::UnaryOp; + use khal_std::glamx::Vec4; + use khal::backend::WebGpu; + use khal::backend::{Backend, Encoder, GpuBackend}; + use khal::{BufferUsages, Shader}; + use nalgebra::DVector; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; + use wgpu::{Features, Limits}; + + #[futures_test::test] + #[serial_test::serial] + async fn gpu_unary_ops_webgpu() { + let webgpu = WebGpu::new(Features::default(), Limits::default()) + .await + .unwrap(); + let backend = GpuBackend::WebGpu(webgpu); + gpu_unary_ops_generic(&backend).await; + } + + async fn gpu_unary_ops_generic(backend: &GpuBackend) { + let unop = super::Unary::from_backend(backend).unwrap(); + + let ops = [ + UnaryOp::Abs, + UnaryOp::Sgn, + UnaryOp::Neg, + UnaryOp::Step, + UnaryOp::Elu, + UnaryOp::Gelu, + UnaryOp::GeluQuick, + UnaryOp::Silu, + UnaryOp::Tanh, + UnaryOp::Relu, + UnaryOp::Sigmoid, + UnaryOp::HardSigmoid, + // UnaryOp::HardSwish, + UnaryOp::Sqr, + UnaryOp::Sqrt, + UnaryOp::Sin, + UnaryOp::Cos, + UnaryOp::Log, + UnaryOp::Exp, + UnaryOp::Reciprocal, + UnaryOp::Erf, + UnaryOp::LeakyRelu, + UnaryOp::Clamp, + UnaryOp::Scale, + UnaryOp::AddScalar, + UnaryOp::Pow, + ]; + + for op in ops { + let mut shapes = TensorLayoutBuffers::new(backend); + + println!("Checking {:?}", op); + + const LEN: u32 = 1757; + + let src = DVector::new_random(LEN as usize); + let dst = DVector::zeros(LEN as usize); + let mut dst_read = DVector::zeros(LEN as usize); + let mut args = Vec4::new( + rand::random(), + rand::random(), + rand::random(), + rand::random(), + ); + if args.y < args.x { + let (x, y) = (args.x, args.y); + args.x = y; + args.y = x; // Ensure min <= max for clamp. + } + let gpu_args = op + .has_args() + .then(|| Tensor::scalar(backend, args, BufferUsages::STORAGE).unwrap()); + let gpu_src = Tensor::vector(backend, &src, BufferUsages::STORAGE).unwrap(); + let mut gpu_dst = Tensor::vector( + backend, + &dst, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("test", None); + unop.launch( + backend, + &mut shapes, + &mut pass, + op, + &mut gpu_dst, + gpu_src.as_view(), + gpu_args.as_ref(), + ) + .unwrap(); + drop(pass); + + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + + backend + .slow_read_buffer(gpu_dst.buffer(), dst_read.as_mut_slice()) + .await + .unwrap(); + + let mut cpu_result = src; + unop.run_cpu(op, &mut cpu_result, args); + + approx::assert_relative_eq!(dst_read, cpu_result, epsilon = 1.0e-5); + } + } + + #[cfg(feature = "cpu")] + #[futures_test::test] + async fn gpu_unary_ops_cpu() { + let backend = GpuBackend::Cpu; + gpu_unary_ops_generic(&backend).await; + } + + #[cfg(feature = "cuda")] + #[futures_test::test] + #[serial_test::serial] + async fn gpu_unary_ops_cuda() { + let cuda = khal::backend::Cuda::new(0).unwrap(); + let backend = GpuBackend::Cuda(cuda); + gpu_unary_ops_generic(&backend).await; + } +} diff --git a/src/ml/win_part.rs b/src/ml/win_part.rs new file mode 100644 index 0000000..7503308 --- /dev/null +++ b/src/ml/win_part.rs @@ -0,0 +1,70 @@ +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::Shader; +use crate::shapes::TensorLayoutBuffers; +use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; + +#[derive(Shader)] +pub struct WinPart { + pub win_part: vortx_shaders::ml::win_part::WinPart, + pub win_unpart: vortx_shaders::ml::win_part::WinUnpart, +} + +impl WinPart { + pub fn launch( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut result: impl AsTensorMut, + source: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut result = result.as_tensor_mut(); + let source = source.as_tensor_ref(); + shapes.insert(backend, result.layout())?; + shapes.insert(backend, source.layout())?; + let shape_result = shapes.get(result.layout()).unwrap(); + let shape_source = shapes.get(source.layout()).unwrap(); + + let result_len = result.len() as u32; + let mut buf_result = result.buffer_mut(); + + self.win_part.call( + pass, + [result_len, 1, 1], + &shape_result.as_slice(), + &shape_source.as_slice(), + &mut buf_result, + &source.buffer(), + ) + } + + pub fn launch_unpart( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + window_size: &Tensor, + mut result: impl AsTensorMut, + source: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut result = result.as_tensor_mut(); + let source = source.as_tensor_ref(); + shapes.insert(backend, result.layout())?; + shapes.insert(backend, source.layout())?; + let shape_result = shapes.get(result.layout()).unwrap(); + let shape_source = shapes.get(source.layout()).unwrap(); + + let result_len = result.len() as u32; + let mut buf_result = result.buffer_mut(); + + self.win_unpart.call( + pass, + [result_len, 1, 1], + &shape_result.as_slice(), + &shape_source.as_slice(), + &window_size.buffer().as_slice(), + &mut buf_result, + &source.buffer(), + ) + } +} diff --git a/vortx-shaders/Cargo.toml b/vortx-shaders/Cargo.toml index 3ababf5..26a55ba 100644 --- a/vortx-shaders/Cargo.toml +++ b/vortx-shaders/Cargo.toml @@ -22,6 +22,8 @@ subgroup_ops = [] cpu = ["khal-std/cpu", "khal/cpu"] cpu-parallel = ["cpu", "khal-std/cpu-parallel"] cuda = ["khal-std/cuda", "khal/cuda"] +# Enables machine-learning (llm inference, reinforcement learning, etc.) operators. +ml = [] [dependencies] khal-std = { workspace = true } diff --git a/vortx-shaders/src/lib.rs b/vortx-shaders/src/lib.rs index 6ab2b13..f7128b1 100644 --- a/vortx-shaders/src/lib.rs +++ b/vortx-shaders/src/lib.rs @@ -13,3 +13,5 @@ extern crate std; pub mod linalg; pub mod utils; +#[cfg(feature = "ml")] +pub mod ml; \ No newline at end of file diff --git a/vortx-shaders/src/ml/batched_multiquery_attention.rs b/vortx-shaders/src/ml/batched_multiquery_attention.rs new file mode 100644 index 0000000..abdef91 --- /dev/null +++ b/vortx-shaders/src/ml/batched_multiquery_attention.rs @@ -0,0 +1,57 @@ +//! Batched multi-query attention. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; + +const WORKGROUP_SIZE: u32 = 64; + +/// Attention parameters. +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct AttentionParams { + /// Maximum sequence length (for KV cache sizing). + pub seq_len: u32, + /// KV dimension (n_kv_heads * head_size). + pub kv_dim: u32, + /// Number of query heads per KV head (for grouped-query attention). + pub kv_mul: u32, + /// Total number of query heads. + pub n_heads: u32, + /// Size of each attention head. + pub head_size: u32, + /// Current position in sequence (0-indexed). + pub pos: u32, +} + +#[inline] +fn div_ceil4(a: u32) -> u32 { + a.div_ceil(4) +} + +/// Multiply and mask attention scores. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn mult_mask_attn( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] params: &[AttentionParams], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] attn: &mut [f32], +) { + // Load params from storage buffer to local variable (enables LICM) + let params = *params.at(0); + + let nonzero_len = params.pos + 1; + let aligned_len = div_ceil4(params.pos + 1) * 4; + if invocation_id.x % aligned_len < nonzero_len { + *attn.at_mut(invocation_id.x as usize) = + *attn.at(invocation_id.x as usize) / (params.head_size as f32).sqrt(); + } else { + *attn.at_mut(invocation_id.x as usize) = 0.0; + } +} diff --git a/vortx-shaders/src/ml/concat.rs b/vortx-shaders/src/ml/concat.rs new file mode 100644 index 0000000..4ab9c58 --- /dev/null +++ b/vortx-shaders/src/ml/concat.rs @@ -0,0 +1,63 @@ +//! Concat operation: concatenates tensors along a given axis. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; +use crate::utils::limits::MAX_NUM_WORKGROUPS; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Copy a source tensor into a slice of the destination tensor along a given axis. +/// +/// This is used to implement concat by calling it once per input tensor, +/// with offset indicating where this tensor's data should be placed in the output. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn concat_copy( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] params: &[u32], // [axis, offset_along_axis] +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *params.at(0); + let offset = *params.at(1); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index in source + let id_src = shape_src.decompose(thread_id); + + // Build destination coordinates by adding offset along the axis + let mut id_dest = id_src; + match axis { + 0 => id_dest.x += offset, + 1 => id_dest.y += offset, + 2 => id_dest.z += offset, + _ => id_dest.w += offset, + } + + let i_src = shape_src.it_vec(id_src) as usize; + let i_dest = shape_dest.it_vec(id_dest) as usize; + + *dest.at_mut(i_dest) = *src.at(i_src); + } +} diff --git a/vortx-shaders/src/ml/conv2d.rs b/vortx-shaders/src/ml/conv2d.rs new file mode 100644 index 0000000..6d5081e --- /dev/null +++ b/vortx-shaders/src/ml/conv2d.rs @@ -0,0 +1,127 @@ +//! 2D Convolution operation (NCHW format). +//! +//! Input X shape: [N, C_in, H, W] +//! Weight W shape: [C_out, C_in/groups, K_H, K_W] +//! Output Y shape: [N, C_out, H_out, W_out] +//! +//! Parameters in params buffer: +//! \[0\] input_height +//! \[1\] input_width +//! \[2\] output_height +//! \[3\] output_width +//! \[4\] kernel_h +//! \[5\] kernel_w +//! \[6\] stride_h +//! \[7\] stride_w +//! \[8\] pad_h +//! \[9\] pad_w +//! \[10\] dilation_h +//! \[11\] dilation_w +//! \[12\] in_channels +//! \[13\] out_channels +//! \[14\] batch_size +//! \[15\] groups (must be 1 for now) + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::utils::limits::MAX_NUM_WORKGROUPS; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Conv2d - compute 2D convolution. +/// +/// This is a straightforward implementation, not optimized for performance. +/// For each output element, iterate over the kernel and compute the convolution. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn conv_2d_nchw( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] output: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] input: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] weight: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] params: &[u32], +) { + let input_h = *params.at(0); + let input_w = *params.at(1); + let output_h = *params.at(2); + let output_w = *params.at(3); + let kernel_h = *params.at(4); + let kernel_w = *params.at(5); + let stride_h = *params.at(6); + let stride_w = *params.at(7); + let pad_h = *params.at(8); + let pad_w = *params.at(9); + let dilation_h = *params.at(10); + let dilation_w = *params.at(11); + let in_channels = *params.at(12); + let out_channels = *params.at(13); + let batch_size = *params.at(14); + let groups = *params.at(15); + + // For grouped convolution: + // - Input channels per group: in_channels / groups + // - Output channels per group: out_channels / groups + // - Weight shape: [out_channels, in_channels/groups, kernel_h, kernel_w] + let in_channels_per_group = in_channels / groups; + let out_channels_per_group = out_channels / groups; + + let output_len = batch_size * out_channels * output_h * output_w; + + for thread_id in (invocation_id.x..output_len).step_by(MAX_NUM_THREADS as usize) { + // Decompose output index into [n, oc, oh, ow] + let ow = thread_id % output_w; + let oh = (thread_id / output_w) % output_h; + let oc = (thread_id / (output_w * output_h)) % out_channels; + let n = thread_id / (output_w * output_h * out_channels); + + // Determine which group this output channel belongs to + let group = oc / out_channels_per_group; + + // Input channels for this group: [group * in_channels_per_group, (group + 1) * in_channels_per_group) + let ic_start = group * in_channels_per_group; + + let mut sum: f32 = 0.0; + + // Iterate over input channels in this group and kernel + for ic_local in 0..in_channels_per_group { + let ic = ic_start + ic_local; + + for kh in 0..kernel_h { + for kw in 0..kernel_w { + // Compute input position with dilation + let ih_signed = (oh * stride_h + kh * dilation_h) as i32 - pad_h as i32; + let iw_signed = (ow * stride_w + kw * dilation_w) as i32 - pad_w as i32; + + // Check if within bounds + if ih_signed >= 0 + && ih_signed < input_h as i32 + && iw_signed >= 0 + && iw_signed < input_w as i32 + { + let ih = ih_signed as u32; + let iw = iw_signed as u32; + + // Input index: [n, ic, ih, iw] in NCHW + let i_input = (n * in_channels * input_h * input_w + + ic * input_h * input_w + + ih * input_w + + iw) as usize; + + // Weight index: [oc, ic_local, kh, kw] + // Note: ic_local is used because weight has shape [out_channels, in_channels/groups, kh, kw] + let i_weight = (oc * in_channels_per_group * kernel_h * kernel_w + + ic_local * kernel_h * kernel_w + + kh * kernel_w + + kw) as usize; + + sum += *input.at(i_input) * *weight.at(i_weight); + } + } + } + } + + *output.at_mut(thread_id as usize) = sum; + } +} diff --git a/vortx-shaders/src/ml/conv_transpose_2d.rs b/vortx-shaders/src/ml/conv_transpose_2d.rs new file mode 100644 index 0000000..81d0e16 --- /dev/null +++ b/vortx-shaders/src/ml/conv_transpose_2d.rs @@ -0,0 +1,275 @@ +//! Transposed 2D convolution. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::{Shapes1, Shapes2, Shapes3}; + +const WORKGROUP_SIZE: u32 = 64; + +/// Initialize destination buffer to zero. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn init_dest( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] dest: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_dest = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + + if invocation_id.x >= shape_dest.len() { + return; + } + + *dest.at_mut(invocation_id.x as usize) = 0.0; +} + +/// Initialize working data buffer to zero. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn init_wdata( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_wdata: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] wdata: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_wdata = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_wdata = *shape_wdata.at(0); + + if invocation_id.x >= shape_wdata.len() { + return; + } + + *wdata.at_mut(invocation_id.x as usize) = 0.0; +} + +/// Initialize src0 (permute kernel data). +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn init_src_a( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src0: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src0: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] wdata: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src0 = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src0 = *shape_src0.at(0); + + if invocation_id.x >= shape_src0.len() { + return; + } + + // permute kernel data (src0) from (Kw x Kh x Cout x Cin) to (Cin x Kw x Kh x Cout) + let id = shape_src0.decompose(invocation_id.x); + let id_wdata = id.z * shape_src0.h * shape_src0.w * shape_src0.n + + id.x * shape_src0.w * shape_src0.n + + id.y * shape_src0.n + + id.w; + *wdata.at_mut(id_wdata as usize) = *src0.at(invocation_id.x as usize); +} + +/// Initialize src1 (permute source data). +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn init_src_b( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src0: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src1: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] src1: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] wdata: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src0, shape_src1) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src0 = *shape_src0.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src1 = *shape_src1.at(0); + + if invocation_id.x >= shape_src1.len() { + return; + } + + // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh) + let nk = shape_src0.len(); + let id = shape_src1.decompose(invocation_id.x); + let id_wdata = nk + id.x * shape_src1.w * shape_src1.c + id.y * shape_src1.c + id.z; + *wdata.at_mut(id_wdata as usize) = *src1.at(invocation_id.x as usize); +} + +/// Reference implementation of transposed 2D convolution. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn conv_transpose_2d_ref( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes3, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src0: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src1: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dest: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] stride: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] wdata: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] dest: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src0, shape_src1, shape_dest) = + (shapes.shape_out, shapes.shape_lhs, shapes.shape_rhs); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src0 = *shape_src0.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src1 = *shape_src1.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + let stride = *stride.at(0); + + let i2 = invocation_id.x; + + if i2 >= shape_dest.c { + return; + } + + let nk = shape_src0.len(); + let ne0 = shape_dest.w; + let nb2 = shape_dest.c_stride; + + let ne00 = shape_src0.w; + let ne01 = shape_src0.h; + let ne03 = shape_src0.n; + let ne10 = shape_src1.w; + let ne11 = shape_src1.h; + let ne12 = shape_src1.c; + + for i11 in 0..ne11 as i32 { + for i10 in 0..ne10 as i32 { + let i1n = i11 * ne10 as i32 * ne12 as i32 + i10 * ne12 as i32; + for i01 in 0..ne01 as i32 { + for i00 in 0..ne00 as i32 { + let mut v = 0.0f32; + + for k in 0..ne03 as i32 { + v += *wdata.at((nk as i32 + i1n + k) as usize) + * *wdata.at((i2 as i32 * ne01 as i32 * ne00 as i32 * ne03 as i32 + + i01 * ne00 as i32 * ne03 as i32 + + i00 * ne03 as i32 + + k) as usize); + } + let dest_idx = (i2 * nb2 + + (i11 * stride as i32 + i01) as u32 * ne0 + + (i10 * stride as i32 + i00) as u32) + as usize; + *dest.at_mut(dest_idx) = *dest.at(dest_idx) + v; + } + } + } + } +} + +/// Transposed 2D convolution. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn conv_transpose_2d( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes3, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src1: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src0: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dest: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] stride: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] src1: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] src0: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dest: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src1, shape_src0, shape_dest) = + (shapes.shape_out, shapes.shape_lhs, shapes.shape_rhs); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src1 = *shape_src1.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src0 = *shape_src0.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + let stride = *stride.at(0); + + let i2 = invocation_id.x; + + if i2 >= shape_dest.c { + return; + } + + for k in 0..(shape_dest.h * shape_dest.w) as i32 { + *dest.at_mut((i2 * shape_dest.c + k as u32) as usize) = 0.0; + } + + for i11 in 0..shape_src1.h as i32 { + for i10 in 0..shape_src1.w as i32 { + for i01 in 0..shape_src0.h as i32 { + for i00 in 0..shape_src0.w as i32 { + let mut v = 0.0f32; + + for k in 0..shape_src0.c as i32 { + v += *src1.at(shape_src1.it(0, i11 as u32, i10 as u32, k as u32) as usize) + * *src0 + .at(shape_src0.it(i2, i01 as u32, i00 as u32, k as u32) as usize); + } + + let dest_id = shape_dest.it( + 0, + i2, + (i11 * stride as i32 + i01) as u32, + (i10 * stride as i32 + i00) as u32, + ) as usize; + *dest.at_mut(dest_id) = *dest.at(dest_id) + v; + } + } + } + } +} diff --git a/vortx-shaders/src/ml/fused_attention.rs b/vortx-shaders/src/ml/fused_attention.rs new file mode 100644 index 0000000..300fd7e --- /dev/null +++ b/vortx-shaders/src/ml/fused_attention.rs @@ -0,0 +1,617 @@ +//! Fused multi-query attention kernel. +//! +//! This kernel fuses the following operations into a single dispatch: +//! 1. Q × K^T (dot products) +//! 2. Scale by 1/sqrt(head_size) +//! 3. Causal masking +//! 4. Softmax +//! 5. Attention × V (weighted sum) +//! +//! Each workgroup processes one query head. + +use crate::ml::batched_multiquery_attention::AttentionParams; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; + +/// Workgroup size - should be >= head_size for efficient V accumulation. +#[cfg(feature = "subgroup_ops")] +const WORKGROUP_SIZE: usize = 32; +#[cfg(not(feature = "subgroup_ops"))] +const WORKGROUP_SIZE: usize = 128; + +/// Maximum sequence length we can handle in shared memory for attention scores. +/// For longer sequences, we use online softmax to avoid storing all scores. +const MAX_SEQ_LEN: usize = 2048; + +/// Block size for Flash Attention - number of KV tokens processed per iteration. +/// Must be tuned to fit shared memory: kv_tile uses BLOCK_KV * WORKGROUP_SIZE * 4 bytes. +/// With BLOCK_KV=32 and WORKGROUP_SIZE=128: 32 * 128 * 4 = 16KB for kv_tile alone. +const BLOCK_KV: usize = 32; + +#[inline] +fn reduce_max(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + if index < stride { + workspace.write( + index, + workspace.read(index).max(workspace.read(index + stride)), + ); + } +} + +#[inline] +fn reduce_sum(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + if index < stride { + workspace.write( + index, + workspace.read(index) + workspace.read(index + stride), + ); + } +} + +/// Fused attention kernel for single-token inference. +/// +/// Workgroup layout: [WORKGROUP_SIZE, 1, 1] +/// Dispatch: [n_kv_heads * kv_mul, 1, 1] workgroups +/// +/// Each workgroup computes attention for one query head: +/// out\[head\] = softmax(Q\[head\] · K^T / sqrt(d)) · V +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn fused_attention( + #[spirv(workgroup_id)] wg_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] attn_scores: &mut [f32; MAX_SEQ_LEN], + #[spirv(workgroup)] max_score: &mut f32, + #[spirv(workgroup)] sum_exp: &mut f32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] params: &[AttentionParams], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] q: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] key_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] value_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] out: &mut [f32], +) { + let params = params.read(0); + let tid = local_id.x as usize; + let head_idx = wg_id.x; // Which query head we're processing + + let head_size = params.head_size as usize; + let kv_mul = params.kv_mul; + let n_kv_heads = params.n_heads / kv_mul; + let seq_len = (params.pos + 1) as usize; // Number of tokens to attend to + + // Which KV head this query head uses + let kv_head = head_idx / kv_mul; + + // Base offsets for this head + let q_base = (head_idx * params.head_size) as usize; + let kv_base = (kv_head * params.head_size) as usize; + + // ========================================================================== + // Phase 1: Compute Q · K^T for all positions, scale, and find max + // ========================================================================== + + // Each thread computes dot products for a subset of positions + // Max iterations: ceil(MAX_SEQ_LEN / WORKGROUP_SIZE) = ceil(2048/128) = 16 + let mut my_max = -1.0e38f32; + + for iter in 0..16 { + let t = tid + iter * WORKGROUP_SIZE; + if t < seq_len { + // Compute Q · K[t] for position t + // Key cache layout: [seq_len, kv_dim] where kv_dim = n_kv_heads * head_size + let k_base = t * (n_kv_heads * params.head_size) as usize + kv_base; + + let mut dot = 0.0f32; + for d in 0..head_size { + let q_val = q.read(q_base + d); + let k_val = key_cache.read(k_base + d); + dot += q_val * k_val; + } + + // Scale by 1/sqrt(head_size) + let score = dot / (head_size as f32).sqrt(); + + // Store score in shared memory + attn_scores.write(t, score); + my_max = my_max.max(score); + } + } + + // Reduce to find global max + workspace.write(tid, my_max); + + #[cfg(feature = "subgroup_ops")] + let max_val = khal_std::sync::subgroup_f_max(my_max); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_max(tid, 64, workspace); + reduce_max(tid, 32, workspace); + reduce_max(tid, 16, workspace); + reduce_max(tid, 8, workspace); + reduce_max(tid, 4, workspace); + reduce_max(tid, 2, workspace); + reduce_max(tid, 1, workspace); + } + + if tid == 0 { + #[cfg(feature = "subgroup_ops")] + { + *max_score = max_val; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *max_score = workspace.read(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ========================================================================== + // Phase 2: Compute exp(score - max) and sum + // ========================================================================== + + let the_max = *max_score; + let mut my_sum = 0.0f32; + + for iter in 0..16 { + let t = tid + iter * WORKGROUP_SIZE; + if t < seq_len { + let score = attn_scores.read(t); + let exp_score = (score - the_max).exp(); + attn_scores.write(t, exp_score); + my_sum += exp_score; + } + } + + // Reduce to find sum + workspace.write(tid, my_sum); + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(my_sum); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(tid, 64, workspace); + reduce_sum(tid, 32, workspace); + reduce_sum(tid, 16, workspace); + reduce_sum(tid, 8, workspace); + reduce_sum(tid, 4, workspace); + reduce_sum(tid, 2, workspace); + reduce_sum(tid, 1, workspace); + } + + if tid == 0 { + #[cfg(feature = "subgroup_ops")] + { + *sum_exp = sum; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *sum_exp = workspace.read(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ========================================================================== + // Phase 3: Normalize attention weights (divide by sum) + // ========================================================================== + + let the_sum = *sum_exp; + let inv_sum = 1.0 / the_sum; + + for iter in 0..16 { + let t = tid + iter * WORKGROUP_SIZE; + if t < seq_len { + let exp_score = attn_scores.read(t); + attn_scores.write(t, exp_score * inv_sum); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ========================================================================== + // Phase 4: Compute weighted sum of values + // ========================================================================== + + // Each thread computes output for a subset of head dimensions + let out_base = (head_idx * params.head_size) as usize; + + // Since head_size <= WORKGROUP_SIZE typically, each thread handles at most one dimension + if tid < head_size { + let mut weighted_sum = 0.0f32; + + for t in 0..seq_len { + // Value cache layout: [seq_len, kv_dim] + let v_base = t * (n_kv_heads * params.head_size) as usize + kv_base; + let v_val = value_cache.read(v_base + tid); + let attn_weight = attn_scores.read(t); + weighted_sum += attn_weight * v_val; + } + + out.write(out_base + tid, weighted_sum); + } +} + +/// Fused attention with online softmax for long sequences. +/// +/// This variant uses online softmax to avoid storing all attention scores, +/// making it memory-efficient for arbitrarily long sequences. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn fused_attention_online( + #[spirv(workgroup_id)] wg_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] out_accum: &mut [f32; WORKGROUP_SIZE], // Accumulator for output + #[spirv(workgroup)] max_score: &mut f32, + #[spirv(workgroup)] sum_exp: &mut f32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] params: &[AttentionParams], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] q: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] key_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] value_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] out: &mut [f32], +) { + let params = params.read(0); + let tid = local_id.x as usize; + let head_idx = wg_id.x; + + let head_size = params.head_size as usize; + let kv_mul = params.kv_mul; + let n_kv_heads = params.n_heads / kv_mul; + let seq_len = (params.pos + 1) as usize; + + let kv_head = head_idx / kv_mul; + let q_base = (head_idx * params.head_size) as usize; + let kv_base = (kv_head * params.head_size) as usize; + let out_base = (head_idx * params.head_size) as usize; + + // Initialize accumulators + if tid < head_size { + out_accum.write(tid, 0.0); + } + if tid == 0 { + *max_score = -1.0e38f32; + *sum_exp = 0.0; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Process sequence one token at a time, using online softmax + for t in 0..seq_len { + // Compute Q · K[t] + let k_base = t * (n_kv_heads * params.head_size) as usize + kv_base; + + // Collaborative dot product - each thread handles part of the dimensions + let mut partial_dot = 0.0f32; + if tid < head_size { + let q_val = q.read(q_base + tid); + let k_val = key_cache.read(k_base + tid); + partial_dot = q_val * k_val; + } + workspace.write(tid, partial_dot); + + // Reduce dot product + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(partial_dot); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(tid, 64, workspace); + reduce_sum(tid, 32, workspace); + reduce_sum(tid, 16, workspace); + reduce_sum(tid, 8, workspace); + reduce_sum(tid, 4, workspace); + reduce_sum(tid, 2, workspace); + reduce_sum(tid, 1, workspace); + } + + if tid == 0 { + #[cfg(feature = "subgroup_ops")] + { + workspace.write(0, sum); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Scale score + let score = workspace.read(0) / (head_size as f32).sqrt(); + + // Online softmax update + let old_max = *max_score; + let new_max = old_max.max(score); + let old_sum = *sum_exp; + + // Rescale old accumulator and sum + let rescale = (old_max - new_max).exp(); + let new_weight = (score - new_max).exp(); + + if tid == 0 { + *max_score = new_max; + *sum_exp = old_sum * rescale + new_weight; + } + + // Update output accumulator with rescaling + let v_base = t * (n_kv_heads * params.head_size) as usize + kv_base; + if tid < head_size { + let old_val = out_accum.read(tid); + let v_val = value_cache.read(v_base + tid); + out_accum.write(tid, old_val * rescale + new_weight * v_val); + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + } + + // Final normalization and write output + let final_sum = *sum_exp; + if tid < head_size { + let val = out_accum.read(tid) / final_sum; + out.write(out_base + tid, val); + } +} + +/// Flash Attention kernel with tiled/block-wise processing. +/// +/// This kernel processes KV in blocks of BLOCK_KV tokens, using online softmax +/// to maintain O(1) memory per softmax row. This is ~100x more efficient than +/// the fused_attention_online kernel which processes one token at a time. +/// +/// Workgroup layout: [WORKGROUP_SIZE, 1, 1] +/// Dispatch: [n_heads, 1, 1] workgroups +/// +/// Each workgroup computes attention for one query head using Flash Attention: +/// - Processes KV cache in blocks of BLOCK_KV tokens +/// - Uses online softmax with rescaling between blocks +/// - Accumulates weighted V values incrementally +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn flash_attention( + #[spirv(workgroup_id)] wg_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] q_shared: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] kv_tile: &mut [f32; BLOCK_KV * WORKGROUP_SIZE], + #[spirv(workgroup)] scores: &mut [f32; BLOCK_KV], + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] out_accum: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] running_max: &mut f32, + #[spirv(workgroup)] running_sum: &mut f32, + #[spirv(workgroup)] block_max_shared: &mut f32, + #[spirv(workgroup)] block_sum_shared: &mut f32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] params: &[AttentionParams], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] q: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] key_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] value_cache: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] out: &mut [f32], +) { + let params = params.read(0); + let tid = local_id.x as usize; + let head_idx = wg_id.x; + + let head_size = params.head_size as usize; + let kv_mul = params.kv_mul; + let n_kv_heads = params.n_heads / kv_mul; + let seq_len = (params.pos + 1) as usize; + + let kv_head = head_idx / kv_mul; + let q_base = (head_idx * params.head_size) as usize; + let kv_base = (kv_head * params.head_size) as usize; + let kv_stride = (n_kv_heads * params.head_size) as usize; + let out_base = (head_idx * params.head_size) as usize; + + // ========================================================================== + // Phase 0: Load Q into shared memory and initialize accumulators + // ========================================================================== + if tid < head_size { + q_shared.write(tid, q.read(q_base + tid)); + out_accum.write(tid, 0.0); + } + if tid == 0 { + *running_max = -1.0e38f32; + *running_sum = 0.0; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ========================================================================== + // Main loop: Process KV in blocks of BLOCK_KV tokens + // ========================================================================== + // Maximum number of blocks we might process + let num_blocks = seq_len.div_ceil(BLOCK_KV); + + for block_idx in 0..num_blocks { + let block_start = block_idx * BLOCK_KV; + let block_end = if block_start + BLOCK_KV < seq_len { + block_start + BLOCK_KV + } else { + seq_len + }; + let block_len = block_end - block_start; + + // ---------------------------------------------------------------------- + // Step 1: Load K block into shared memory + // ---------------------------------------------------------------------- + // Layout: kv_tile[pos_in_block * head_size + dim] + // Max elements = BLOCK_KV * head_size, max iterations = ceil(BLOCK_KV * head_size / WORKGROUP_SIZE) + for iter in 0..BLOCK_KV { + let load_idx = tid + iter * WORKGROUP_SIZE; + if load_idx < block_len * head_size { + let pos_in_block = load_idx / head_size; + let dim = load_idx % head_size; + let global_pos = block_start + pos_in_block; + let k_idx = global_pos * kv_stride + kv_base + dim; + kv_tile.write(load_idx, key_cache.read(k_idx)); + } + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ---------------------------------------------------------------------- + // Step 2: Compute Q · K[t] for all positions in block + // ---------------------------------------------------------------------- + // Each thread handles at most one position (since BLOCK_KV <= WORKGROUP_SIZE) + if tid < block_len { + let mut dot = 0.0f32; + for d in 0..head_size { + let q_val = q_shared.read(d); + let k_val = kv_tile.read(tid * head_size + d); + dot += q_val * k_val; + } + + // Scale by 1/sqrt(head_size) and apply causal mask + let global_pos = block_start + tid; + let score = if global_pos <= params.pos as usize { + dot / (head_size as f32).sqrt() + } else { + -1.0e38f32 // Causal mask: future positions masked out + }; + scores.write(tid, score); + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ---------------------------------------------------------------------- + // Step 3: Find block max via parallel reduction + // ---------------------------------------------------------------------- + let my_max = if tid < block_len { + scores.read(tid) + } else { + -1.0e38f32 + }; + workspace.write(tid, my_max); + + #[cfg(feature = "subgroup_ops")] + let max_val = khal_std::sync::subgroup_f_max(my_max); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_max(tid, 64, workspace); + reduce_max(tid, 32, workspace); + reduce_max(tid, 16, workspace); + reduce_max(tid, 8, workspace); + reduce_max(tid, 4, workspace); + reduce_max(tid, 2, workspace); + reduce_max(tid, 1, workspace); + } + + // Store result in shared variable so all threads can read it after barrier + if tid == 0 { + #[cfg(feature = "subgroup_ops")] + { + *block_max_shared = max_val; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *block_max_shared = workspace.read(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + let block_max = *block_max_shared; + + // ---------------------------------------------------------------------- + // Step 4: Compute exp(score - block_max) and sum + // ---------------------------------------------------------------------- + let my_sum = if tid < block_len { + let exp_score = (scores.read(tid) - block_max).exp(); + scores.write(tid, exp_score); // Overwrite with exp values + exp_score + } else { + 0.0f32 + }; + workspace.write(tid, my_sum); + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(my_sum); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(tid, 64, workspace); + reduce_sum(tid, 32, workspace); + reduce_sum(tid, 16, workspace); + reduce_sum(tid, 8, workspace); + reduce_sum(tid, 4, workspace); + reduce_sum(tid, 2, workspace); + reduce_sum(tid, 1, workspace); + } + + // Store result in shared variable so all threads can read it after barrier + if tid == 0 { + #[cfg(feature = "subgroup_ops")] + { + *block_sum_shared = sum; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *block_sum_shared = workspace.read(0); + } + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + let block_sum = *block_sum_shared; + + // ---------------------------------------------------------------------- + // Step 5: Update running statistics with rescaling + // ---------------------------------------------------------------------- + let old_max = *running_max; + let old_sum = *running_sum; + let new_max = old_max.max(block_max); + let rescale_old = (old_max - new_max).exp(); + let rescale_new = (block_max - new_max).exp(); + + if tid == 0 { + *running_max = new_max; + *running_sum = old_sum * rescale_old + block_sum * rescale_new; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // ---------------------------------------------------------------------- + // Step 6: Load V block and accumulate weighted values + // ---------------------------------------------------------------------- + // Reuse kv_tile for V block + for iter in 0..BLOCK_KV { + let load_idx = tid + iter * WORKGROUP_SIZE; + if load_idx < block_len * head_size { + let pos_in_block = load_idx / head_size; + let dim = load_idx % head_size; + let global_pos = block_start + pos_in_block; + let v_idx = global_pos * kv_stride + kv_base + dim; + kv_tile.write(load_idx, value_cache.read(v_idx)); + } + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Update output accumulator: O = O * rescale_old + (S_exp @ V) * rescale_new + // Each thread handles one output dimension + if tid < head_size { + // Rescale old accumulator + let old_val = out_accum.read(tid) * rescale_old; + + // Accumulate weighted V values for this dimension + let mut new_contrib = 0.0f32; + for pos in 0..block_len { + let weight = scores.read(pos) * rescale_new; + let v_val = kv_tile.read(pos * head_size + tid); + new_contrib += weight * v_val; + } + + out_accum.write(tid, old_val + new_contrib); + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + } + + // ========================================================================== + // Final: Normalize by running sum and write output + // ========================================================================== + let final_sum = *running_sum; + if tid < head_size { + let val = out_accum.read(tid) / final_sum; + out.write(out_base + tid, val); + } +} diff --git a/vortx-shaders/src/ml/gather.rs b/vortx-shaders/src/ml/gather.rs new file mode 100644 index 0000000..4d2fb03 --- /dev/null +++ b/vortx-shaders/src/ml/gather.rs @@ -0,0 +1,77 @@ +//! Gather operation: gathers elements from source tensor based on indices along a given axis. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; +use crate::utils::limits::MAX_NUM_WORKGROUPS; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Gather elements from source based on indices along axis. +/// +/// For axis=0: +/// output\[i, j, k\] = input\[indices\[i\], j, k\] +/// For axis=1: +/// output\[i, j, k\] = input\[i, indices\[j\], k\] +/// etc. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gather( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] indices: &[i32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] axis_buf: &[u32], +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *axis_buf.at(0); + + for thread_id in (invocation_id.x..shape_dest.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index to 4D coordinates in output + let id_dest = shape_dest.decompose(thread_id); + + // Get the coordinate along the gather axis + let idx_coord = match axis { + 0 => id_dest.x, + 1 => id_dest.y, + 2 => id_dest.z, + _ => id_dest.w, + }; + + // Look up the index value + let gathered_idx = *indices.at(idx_coord as usize); + + // Build source coordinates by replacing the axis coordinate with the gathered index + let mut id_src = id_dest; + match axis { + 0 => id_src.x = gathered_idx as u32, + 1 => id_src.y = gathered_idx as u32, + 2 => id_src.z = gathered_idx as u32, + _ => id_src.w = gathered_idx as u32, + } + + let i_dest = shape_dest.it_vec(id_dest) as usize; + let i_src = shape_src.it_vec(id_src) as usize; + + *dest.at_mut(i_dest) = *src.at(i_src); + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs new file mode 100644 index 0000000..06e03ca --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs @@ -0,0 +1,164 @@ +//! Q4_0x2 quantized GEMV shader. +//! +//! BlockQ4_0x2 contains two BlockQ4_0 blocks (f16 scale + 16 x 4-bit quants each). + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{Mat4, UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: usize = 32; +const COLS_STEP: u32 = 4; +const BLOCK_Q4_0X2_SIZE: u32 = 9; + +#[inline] +fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) { + if index < stride { + let val = sketch.read(index + stride); + *sketch.at_mut(index) += val; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +/// Dequantize a part of BlockQ4_0 data. +/// Returns two Vec4s: low nibbles and high nibbles, scaled. +#[inline] +fn dequantize_part(data: u32, scale: f32) -> [Vec4; 2] { + let x0 = (data & 0x0F) as i32 - 8; + let x1 = ((data >> 4) & 0x0F) as i32 - 8; + let x2 = ((data >> 8) & 0x0F) as i32 - 8; + let x3 = ((data >> 12) & 0x0F) as i32 - 8; + let x4 = ((data >> 16) & 0x0F) as i32 - 8; + let x5 = ((data >> 20) & 0x0F) as i32 - 8; + let x6 = ((data >> 24) & 0x0F) as i32 - 8; + let x7 = ((data >> 28) & 0x0F) as i32 - 8; + + [ + Vec4::new(x0 as f32, x2 as f32, x4 as f32, x6 as f32) * scale, + Vec4::new(x1 as f32, x3 as f32, x5 as f32, x7 as f32) * scale, + ] +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + #[cfg(not(feature = "push_constants"))] + let shape_m = shape_m.read(0); + + let mut sum = [Vec4::ZERO]; + let lid = local_id.x; + + for j in 0..(shape_m.w / COLS_STEP) { + // Calculate block indices for 4 matrix rows + let quant0 = (shape_m.it(0, 0, workgroup_id.x * 4, j * COLS_STEP + lid / 8) + * BLOCK_Q4_0X2_SIZE) as usize; + let quant1 = (shape_m.it(0, 0, workgroup_id.x * 4 + 1, j * COLS_STEP + lid / 8) + * BLOCK_Q4_0X2_SIZE) as usize; + let quant2 = (shape_m.it(0, 0, workgroup_id.x * 4 + 2, j * COLS_STEP + lid / 8) + * BLOCK_Q4_0X2_SIZE) as usize; + let quant3 = (shape_m.it(0, 0, workgroup_id.x * 4 + 3, j * COLS_STEP + lid / 8) + * BLOCK_Q4_0X2_SIZE) as usize; + + let j_base = (j * 16 * COLS_STEP) as usize; + let vj_a = v.read(j_base + lid as usize + (lid / 4) as usize * 4); + let vj_b = v.read(j_base + lid as usize + (lid / 4) as usize * 4 + 4); + + if (lid / 4).is_multiple_of(2) { + // Dequantizing block 1 + let llid = (lid % 4) as usize; + + let scale0 = unpack_half2x16(m.read(quant0)).x; + let data0 = m.read(quant0 + llid) >> 16 | m.read(quant0 + llid + 1) << 16; + let parts0 = dequantize_part(data0, scale0); + + let scale1 = unpack_half2x16(m.read(quant1)).x; + let data1 = m.read(quant1 + llid) >> 16 | m.read(quant1 + llid + 1) << 16; + let parts1 = dequantize_part(data1, scale1); + + let scale2 = unpack_half2x16(m.read(quant2)).x; + let data2 = m.read(quant2 + llid) >> 16 | m.read(quant2 + llid + 1) << 16; + let parts2 = dequantize_part(data2, scale2); + + let scale3 = unpack_half2x16(m.read(quant3)).x; + let data3 = m.read(quant3 + llid) >> 16 | m.read(quant3 + llid + 1) << 16; + let parts3 = dequantize_part(data3, scale3); + + // Matrix-vector multiply + let mat0 = Mat4::from_cols(parts0[0], parts1[0], parts2[0], parts3[0]); + let mat1 = Mat4::from_cols(parts0[1], parts1[1], parts2[1], parts3[1]); + sum[0] += mat0.transpose() * vj_a + mat1.transpose() * vj_b; + } else { + // Dequantizing block 2 + let llid = (lid % 4) as usize; + + let scale0 = unpack_half2x16(m.read(quant0 + 4)).y; + let data0 = m.read(quant0 + llid + 5); + let parts0 = dequantize_part(data0, scale0); + + let scale1 = unpack_half2x16(m.read(quant1 + 4)).y; + let data1 = m.read(quant1 + llid + 5); + let parts1 = dequantize_part(data1, scale1); + + let scale2 = unpack_half2x16(m.read(quant2 + 4)).y; + let data2 = m.read(quant2 + llid + 5); + let parts2 = dequantize_part(data2, scale2); + + let scale3 = unpack_half2x16(m.read(quant3 + 4)).y; + let data3 = m.read(quant3 + llid + 5); + let parts3 = dequantize_part(data3, scale3); + + // Matrix-vector multiply + let mat0 = Mat4::from_cols(parts0[0], parts1[0], parts2[0], parts3[0]); + let mat1 = Mat4::from_cols(parts0[1], parts1[1], parts2[1], parts3[1]); + sum[0] += mat0.transpose() * vj_a + mat1.transpose() * vj_b; + } + } + + #[cfg(feature = "subgroup_ops")] + { + let reduced = Vec4::new( + khal_std::sync::subgroup_f_add(sum[0].x), + khal_std::sync::subgroup_f_add(sum[0].y), + khal_std::sync::subgroup_f_add(sum[0].z), + khal_std::sync::subgroup_f_add(sum[0].w), + ); + if lid == 0 { + *out.at_mut(workgroup_id.x as usize) = reduced; + } + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *sketch.at_mut(lid as usize) = sum[0]; + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + reduce_sum(lid as usize, 16, sketch); + reduce_sum(lid as usize, 8, sketch); + reduce_sum(lid as usize, 4, sketch); + reduce_sum(lid as usize, 2, sketch); + reduce_sum(lid as usize, 1, sketch); + + if lid == 0 { + *out.at_mut(workgroup_id.x as usize) = sketch.read(0); + } + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs new file mode 100644 index 0000000..78fb824 --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs @@ -0,0 +1,108 @@ +//! Q4_1x2 quantized GEMV shader. +//! +//! BlockQ4_1x2 contains two BlockQ4_1 blocks (f16 scale + f16 min + 16 x 4-bit quants each). + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{UVec3, Vec2, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: u32 = 64; + +/// Dequantize a part of BlockQ4_1 data. +/// Returns two Vec4s: low nibbles and high nibbles, scaled and shifted. +#[inline] +fn dequantize_part(data: u32, scale_mid: Vec2) -> [Vec4; 2] { + let x0 = data & 0x0F; + let x1 = (data >> 4) & 0x0F; + let x2 = (data >> 8) & 0x0F; + let x3 = (data >> 12) & 0x0F; + let x4 = (data >> 16) & 0x0F; + let x5 = (data >> 20) & 0x0F; + let x6 = (data >> 24) & 0x0F; + let x7 = (data >> 28) & 0x0F; + + [ + Vec4::new(x0 as f32, x2 as f32, x4 as f32, x6 as f32) * scale_mid.x + scale_mid.y, + Vec4::new(x1 as f32, x3 as f32, x5 as f32, x7 as f32) * scale_mid.x + scale_mid.y, + ] +} + +/// Dequantize a full BlockQ4_1x2 block. +#[inline] +fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { + let mut result = [Vec4::ZERO; 16]; + + // First block + let scale_mid_a = unpack_half2x16(*data.at(base)); + for k in 0..4 { + let parts = dequantize_part(*data.at(base + k + 1), scale_mid_a); + result[k] = parts[0]; + result[4 + k] = parts[1]; + } + + // Second block + let scale_mid_b = unpack_half2x16(*data.at(base + 5)); + for k in 0..4 { + let parts = dequantize_part(*data.at(base + k + 6), scale_mid_b); + result[8 + k] = parts[0]; + result[12 + k] = parts[1]; + } + + result +} + +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gemv( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + if invocation_id.x < shape_m.h { + let i_out = invocation_id.x as usize; + let mut sum = 0.0f32; + + for j in 0..shape_m.w { + let block_idx = (shape_m.it(0, 0, invocation_id.x, j) * 10) as usize; // BlockQ4_1x2 is 10 u32s + let dequant = dequantize_block(m, block_idx); + + // Unroll calculation with all block elements + let i_base = (j * 16) as usize; + sum += dequant[0].dot(*v.at(i_base)) + + dequant[1].dot(*v.at(i_base + 1)) + + dequant[2].dot(*v.at(i_base + 2)) + + dequant[3].dot(*v.at(i_base + 3)) + + dequant[4].dot(*v.at(i_base + 4)) + + dequant[5].dot(*v.at(i_base + 5)) + + dequant[6].dot(*v.at(i_base + 6)) + + dequant[7].dot(*v.at(i_base + 7)) + + dequant[8].dot(*v.at(i_base + 8)) + + dequant[9].dot(*v.at(i_base + 9)) + + dequant[10].dot(*v.at(i_base + 10)) + + dequant[11].dot(*v.at(i_base + 11)) + + dequant[12].dot(*v.at(i_base + 12)) + + dequant[13].dot(*v.at(i_base + 13)) + + dequant[14].dot(*v.at(i_base + 14)) + + dequant[15].dot(*v.at(i_base + 15)); + } + + *out.at_mut(i_out) = Vec4::splat(sum); + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q4_k.rs b/vortx-shaders/src/ml/gemv_quant_q4_k.rs new file mode 100644 index 0000000..cc0d5c2 --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q4_k.rs @@ -0,0 +1,182 @@ +//! Q4_K quantized GEMV shader. +//! +//! BlockQ4K: super-block scale, super-block min, 12 bytes scales/mins, 128 bytes quants (256 4-bit values). + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{UVec2, UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: usize = 32; +// BlockQ4K size in u32s: 1 (d_dmin) + 3 (scales) + 32 (qs) = 36 +const BLOCK_Q4K_SIZE: u32 = 36; + +#[inline] +fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) { + if index < stride { + let val = *sketch.at(index + stride); + *sketch.at_mut(index) += val; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +/// Unpack scale and min from the packed scales array. +#[inline] +fn unpack_scale_and_min(j: u32, qj_prev: u32, qj: u32, qj_next: u32) -> UVec2 { + let shift = (j % 4) * 8; + let qj_prev_shifted = (qj_prev >> shift) & 0x00ff; + let qj_shifted = (qj >> shift) & 0x00ff; + let qj_next_shifted = (qj_next >> shift) & 0x00ff; + + if j < 4 { + let d = qj_shifted & 63; + let m = qj_next_shifted & 63; + UVec2::new(d, m) + } else { + let d = (qj_next_shifted & 0xf) | ((qj_prev_shifted >> 6) << 4); + let m = (qj_next_shifted >> 4) | ((qj_shifted >> 6) << 4); + UVec2::new(d, m) + } +} + +/// Dequantize Q4_K block for a workgroup thread. +#[inline] +fn dequantize_q4_k_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 2] { + let d_dmin_id = (block_id * BLOCK_Q4K_SIZE) as usize; + let scales_id = d_dmin_id + 1; + let qs_id = d_dmin_id + 4; + + let d_dmin = unpack_half2x16(*m.at(d_dmin_id)); + let d = d_dmin.x; + let min = d_dmin.y; + + // 32 threads workgroups + let j = k / 8; + let is = j * 2; + + let qj_prev1 = *m.at(scales_id + (is / 4).max(1) as usize - 1); + let qj1 = *m.at(scales_id + (is / 4) as usize); + let qj_next1 = *m.at(scales_id + (is / 4 + 1) as usize); + let sc_m1 = unpack_scale_and_min(is, qj_prev1, qj1, qj_next1); + let d1 = d * sc_m1.x as f32; + let m1 = min * sc_m1.y as f32; + + let qj_prev2 = *m.at(scales_id + ((is + 1) / 4).max(1) as usize - 1); + let qj2 = *m.at(scales_id + ((is + 1) / 4) as usize); + let qj_next2 = *m.at(scales_id + ((is + 1) / 4 + 1) as usize); + let sc_m2 = unpack_scale_and_min(is + 1, qj_prev2, qj2, qj_next2); + let d2 = d * sc_m2.x as f32; + let m2 = min * sc_m2.y as f32; + + let qs = *m.at(qs_id + k as usize); + + let res_a = Vec4::new( + (qs & 0xF) as f32, + ((qs >> 8) & 0xF) as f32, + ((qs >> 16) & 0xF) as f32, + ((qs >> 24) & 0xF) as f32, + ) * d1 + - m1; + + let res_b = Vec4::new( + ((qs >> 4) & 0xF) as f32, + ((qs >> 12) & 0xF) as f32, + ((qs >> 20) & 0xF) as f32, + ((qs >> 28) & 0xF) as f32, + ) * d2 + - m2; + + [res_a, res_b] +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + let mut sum = Vec4::ZERO; + let lid = local_id.x; + + for j in 0..shape_m.w { + let quant0 = shape_m.it(0, 0, workgroup_id.x * 4, j); + let quant1 = shape_m.it(0, 0, workgroup_id.x * 4 + 1, j); + let quant2 = shape_m.it(0, 0, workgroup_id.x * 4 + 2, j); + let quant3 = shape_m.it(0, 0, workgroup_id.x * 4 + 3, j); + + let parts0 = dequantize_q4_k_workgroup(m, quant0, lid); + let parts1 = dequantize_q4_k_workgroup(m, quant1, lid); + let parts2 = dequantize_q4_k_workgroup(m, quant2, lid); + let parts3 = dequantize_q4_k_workgroup(m, quant3, lid); + + let j_base = (j * 64) as usize; + let jj = (lid & 0xfffffff8) as usize; // == (lid / 8) * 8 + let vj_a = *v.at(j_base + lid as usize + jj); + let vj_b = *v.at(j_base + lid as usize + jj + 8); + + sum += Vec4::new( + parts0[0].dot(vj_a), + parts1[0].dot(vj_a), + parts2[0].dot(vj_a), + parts3[0].dot(vj_a), + ); + sum += Vec4::new( + parts0[1].dot(vj_b), + parts1[1].dot(vj_b), + parts2[1].dot(vj_b), + parts3[1].dot(vj_b), + ); + } + + #[cfg(feature = "subgroup_ops")] + { + let reduced = Vec4::new( + khal_std::sync::subgroup_f_add(sum.x), + khal_std::sync::subgroup_f_add(sum.y), + khal_std::sync::subgroup_f_add(sum.z), + khal_std::sync::subgroup_f_add(sum.w), + ); + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = reduced; + } + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *sketch.at_mut(lid as usize) = sum; + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // reduce_sum(lid as usize, 32, sketch); + reduce_sum(lid as usize, 16, sketch); + reduce_sum(lid as usize, 8, sketch); + reduce_sum(lid as usize, 4, sketch); + reduce_sum(lid as usize, 2, sketch); + reduce_sum(lid as usize, 1, sketch); + + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = *sketch.at(0); + } + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs new file mode 100644 index 0000000..5e06dca --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs @@ -0,0 +1,121 @@ +//! Q5_0x2 quantized GEMV shader. +//! +//! BlockQ5_0x2 contains two BlockQ5_0 blocks (f16 scale + u32 high bits + 16 x 4-bit quants each). + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: u32 = 64; + +/// Dequantize a part of BlockQ5_0 data. +#[inline] +fn dequantize_part(j0: u32, qh: u32, data: u32, scale: f32) -> [Vec4; 2] { + let xh0 = ((qh >> j0) << 4) & 0x10; + let x0 = ((data & 0x0F) | xh0) as i32 - 16; + let xh1 = (qh >> (j0 + 12)) & 0x10; + let x1 = (((data >> 4) & 0x0F) | xh1) as i32 - 16; + let xh2 = ((qh >> (j0 + 1)) << 4) & 0x10; + let x2 = (((data >> 8) & 0x0F) | xh2) as i32 - 16; + let xh3 = (qh >> (j0 + 1 + 12)) & 0x10; + let x3 = (((data >> 12) & 0x0F) | xh3) as i32 - 16; + let xh4 = ((qh >> (j0 + 2)) << 4) & 0x10; + let x4 = (((data >> 16) & 0x0F) | xh4) as i32 - 16; + let xh5 = (qh >> (j0 + 2 + 12)) & 0x10; + let x5 = (((data >> 20) & 0x0F) | xh5) as i32 - 16; + let xh6 = ((qh >> (j0 + 3)) << 4) & 0x10; + let x6 = (((data >> 24) & 0x0F) | xh6) as i32 - 16; + let xh7 = (qh >> (j0 + 3 + 12)) & 0x10; + let x7 = (((data >> 28) & 0x0F) | xh7) as i32 - 16; + + [ + Vec4::new(x0 as f32, x2 as f32, x4 as f32, x6 as f32) * scale, + Vec4::new(x1 as f32, x3 as f32, x5 as f32, x7 as f32) * scale, + ] +} + +/// Dequantize a full BlockQ5_0x2 block. +#[inline] +fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { + let mut result = [Vec4::ZERO; 16]; + + // First block + let d1 = unpack_half2x16(*data.at(base)).x; + let qh1 = *data.at(base) >> 16 | *data.at(base + 1) << 16; + + for k in 0u32..4 { + let d = *data.at(base + k as usize + 1) >> 16 | *data.at(base + k as usize + 2) << 16; + let parts = dequantize_part(k * 4, qh1, d, d1); + result[k as usize] = parts[0]; + result[4 + k as usize] = parts[1]; + } + + // Second block + let d2 = unpack_half2x16(*data.at(base + 5)).y; + let qh2 = *data.at(base + 6); + + for k in 0u32..4 { + let d = *data.at(base + k as usize + 7); + let parts = dequantize_part(k * 4, qh2, d, d2); + result[8 + k as usize] = parts[0]; + result[12 + k as usize] = parts[1]; + } + + result +} + +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gemv( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + if invocation_id.x < shape_m.h { + let i_out = invocation_id.x as usize; + let mut sum = 0.0f32; + + for j in 0..shape_m.w { + let block_idx = (shape_m.it(0, 0, invocation_id.x, j) * 11) as usize; // BlockQ5_0x2 is 11 u32s + let dequant = dequantize_block(m, block_idx); + + // Unroll calculation with all block elements + let i_base = (j * 16) as usize; + sum += dequant[0].dot(*v.at(i_base)) + + dequant[1].dot(*v.at(i_base + 1)) + + dequant[2].dot(*v.at(i_base + 2)) + + dequant[3].dot(*v.at(i_base + 3)) + + dequant[4].dot(*v.at(i_base + 4)) + + dequant[5].dot(*v.at(i_base + 5)) + + dequant[6].dot(*v.at(i_base + 6)) + + dequant[7].dot(*v.at(i_base + 7)) + + dequant[8].dot(*v.at(i_base + 8)) + + dequant[9].dot(*v.at(i_base + 9)) + + dequant[10].dot(*v.at(i_base + 10)) + + dequant[11].dot(*v.at(i_base + 11)) + + dequant[12].dot(*v.at(i_base + 12)) + + dequant[13].dot(*v.at(i_base + 13)) + + dequant[14].dot(*v.at(i_base + 14)) + + dequant[15].dot(*v.at(i_base + 15)); + } + + *out.at_mut(i_out) = Vec4::splat(sum); + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs new file mode 100644 index 0000000..bab4ae2 --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs @@ -0,0 +1,121 @@ +//! Q5_1x2 quantized GEMV shader. +//! +//! BlockQ5_1x2 contains two BlockQ5_1 blocks (f16 scale + f16 min + u32 high bits + 16 x 4-bit quants each). + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{UVec3, Vec2, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: u32 = 64; + +/// Dequantize a part of BlockQ5_1 data. +#[inline] +fn dequantize_part(j0: u32, qh: u32, data: u32, d_m: Vec2) -> [Vec4; 2] { + let xh0 = ((qh >> j0) << 4) & 0x10; + let x0 = (data & 0x0F) | xh0; + let xh1 = (qh >> (j0 + 12)) & 0x10; + let x1 = ((data >> 4) & 0x0F) | xh1; + let xh2 = ((qh >> (j0 + 1)) << 4) & 0x10; + let x2 = ((data >> 8) & 0x0F) | xh2; + let xh3 = (qh >> (j0 + 1 + 12)) & 0x10; + let x3 = ((data >> 12) & 0x0F) | xh3; + let xh4 = ((qh >> (j0 + 2)) << 4) & 0x10; + let x4 = ((data >> 16) & 0x0F) | xh4; + let xh5 = (qh >> (j0 + 2 + 12)) & 0x10; + let x5 = ((data >> 20) & 0x0F) | xh5; + let xh6 = ((qh >> (j0 + 3)) << 4) & 0x10; + let x6 = ((data >> 24) & 0x0F) | xh6; + let xh7 = (qh >> (j0 + 3 + 12)) & 0x10; + let x7 = ((data >> 28) & 0x0F) | xh7; + + [ + Vec4::new(x0 as f32, x2 as f32, x4 as f32, x6 as f32) * d_m.x + d_m.y, + Vec4::new(x1 as f32, x3 as f32, x5 as f32, x7 as f32) * d_m.x + d_m.y, + ] +} + +/// Dequantize a full BlockQ5_1x2 block. +#[inline] +fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { + let mut result = [Vec4::ZERO; 16]; + + // First block + let d_m1 = unpack_half2x16(*data.at(base)); + let qh1 = *data.at(base + 1); + + for k in 0u32..4 { + let d = *data.at(base + k as usize + 2); + let parts = dequantize_part(k * 4, qh1, d, d_m1); + result[k as usize] = parts[0]; + result[4 + k as usize] = parts[1]; + } + + // Second block + let d_m2 = unpack_half2x16(*data.at(base + 6)); + let qh2 = *data.at(base + 7); + + for k in 0u32..4 { + let d = *data.at(base + k as usize + 8); + let parts = dequantize_part(k * 4, qh2, d, d_m2); + result[8 + k as usize] = parts[0]; + result[12 + k as usize] = parts[1]; + } + + result +} + +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gemv( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + if invocation_id.x < shape_m.h { + let i_out = invocation_id.x as usize; + let mut sum = 0.0f32; + + for j in 0..shape_m.w { + let block_idx = (shape_m.it(0, 0, invocation_id.x, j) * 12) as usize; // BlockQ5_1x2 is 12 u32s + let dequant = dequantize_block(m, block_idx); + + // Unroll calculation with all block elements + let i_base = (j * 16) as usize; + sum += dequant[0].dot(*v.at(i_base)) + + dequant[1].dot(*v.at(i_base + 1)) + + dequant[2].dot(*v.at(i_base + 2)) + + dequant[3].dot(*v.at(i_base + 3)) + + dequant[4].dot(*v.at(i_base + 4)) + + dequant[5].dot(*v.at(i_base + 5)) + + dequant[6].dot(*v.at(i_base + 6)) + + dequant[7].dot(*v.at(i_base + 7)) + + dequant[8].dot(*v.at(i_base + 8)) + + dequant[9].dot(*v.at(i_base + 9)) + + dequant[10].dot(*v.at(i_base + 10)) + + dequant[11].dot(*v.at(i_base + 11)) + + dequant[12].dot(*v.at(i_base + 12)) + + dequant[13].dot(*v.at(i_base + 13)) + + dequant[14].dot(*v.at(i_base + 14)) + + dequant[15].dot(*v.at(i_base + 15)); + } + + *out.at_mut(i_out) = Vec4::splat(sum); + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q5_k.rs b/vortx-shaders/src/ml/gemv_quant_q5_k.rs new file mode 100644 index 0000000..1de769b --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q5_k.rs @@ -0,0 +1,196 @@ +//! Q5_K quantized GEMV shader. +//! +//! BlockQ5K: super-block scale, super-block min, 12 bytes scales/mins, 32 bytes high bits, 128 bytes quants. + +use crate::utils::half::unpack_half2x16; +use khal_std::glamx::{UVec2, UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: usize = 32; +// BlockQ5K size in u32s: 1 (d_dmin) + 3 (scales) + 8 (qh) + 32 (qs) = 44 +const BLOCK_Q5K_SIZE: u32 = 44; + +#[inline] +fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) { + if index < stride { + let val = *sketch.at(index + stride); + *sketch.at_mut(index) += val; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +/// Unpack scale and min from the packed scales array. +/// Shared with Q4_K. +#[inline] +fn unpack_scale_and_min(j: u32, qj_prev: u32, qj: u32, qj_next: u32) -> UVec2 { + let shift = (j % 4) * 8; + let qj_prev_shifted = (qj_prev >> shift) & 0x00ff; + let qj_shifted = (qj >> shift) & 0x00ff; + let qj_next_shifted = (qj_next >> shift) & 0x00ff; + + if j < 4 { + let d = qj_shifted & 63; + let m = qj_next_shifted & 63; + UVec2::new(d, m) + } else { + let d = (qj_next_shifted & 0xf) | ((qj_prev_shifted >> 6) << 4); + let m = (qj_next_shifted >> 4) | ((qj_shifted >> 6) << 4); + UVec2::new(d, m) + } +} + +/// Dequantize Q5_K block for a workgroup thread. +#[inline] +fn dequantize_q5_k_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 2] { + let d_dmin_id = (block_id * BLOCK_Q5K_SIZE) as usize; + let scales_id = d_dmin_id + 1; + let qh_id = scales_id + 3; + let qs_id = qh_id + 8; + + let d_dmin = unpack_half2x16(*m.at(d_dmin_id)); + let d = d_dmin.x; + let min = d_dmin.y; + + let j = k / 8; + let l = k % 8; + let is = j * 2; + let u1 = 1u32 << (j * 2); + let u2 = 2u32 << (j * 2); + + let qj_prev1 = *m.at(scales_id + (is / 4).max(1) as usize - 1); + let qj1 = *m.at(scales_id + (is / 4) as usize); + let qj_next1 = *m.at(scales_id + (is / 4 + 1) as usize); + let sc_m1 = unpack_scale_and_min(is, qj_prev1, qj1, qj_next1); + let d1 = d * sc_m1.x as f32; + let m1 = min * sc_m1.y as f32; + + let qj_prev2 = *m.at(scales_id + ((is + 1) / 4).max(1) as usize - 1); + let qj2 = *m.at(scales_id + ((is + 1) / 4) as usize); + let qj_next2 = *m.at(scales_id + ((is + 1) / 4 + 1) as usize); + let sc_m2 = unpack_scale_and_min(is + 1, qj_prev2, qj2, qj_next2); + let d2 = d * sc_m2.x as f32; + let m2 = min * sc_m2.y as f32; + + let qs = *m.at(qs_id + k as usize); + let qh = *m.at(qh_id + l as usize); + + #[inline] + fn select_u32(cond: bool, t: u32, f: u32) -> u32 { + if cond { + t + } else { + f + } + } + + let res_a = Vec4::new( + ((qs & 0xF) + select_u32((qh & u1) != 0, 16, 0)) as f32, + (((qs >> 8) & 0xF) + select_u32(((qh >> 8) & u1) != 0, 16, 0)) as f32, + (((qs >> 16) & 0xF) + select_u32(((qh >> 16) & u1) != 0, 16, 0)) as f32, + (((qs >> 24) & 0xF) + select_u32(((qh >> 24) & u1) != 0, 16, 0)) as f32, + ) * d1 + - m1; + + let res_b = Vec4::new( + (((qs >> 4) & 0xF) + select_u32((qh & u2) != 0, 16, 0)) as f32, + (((qs >> 12) & 0xF) + select_u32(((qh >> 8) & u2) != 0, 16, 0)) as f32, + (((qs >> 20) & 0xF) + select_u32(((qh >> 16) & u2) != 0, 16, 0)) as f32, + (((qs >> 28) & 0xF) + select_u32(((qh >> 24) & u2) != 0, 16, 0)) as f32, + ) * d2 + - m2; + + [res_a, res_b] +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + let mut sum = Vec4::ZERO; + let lid = local_id.x; + + for j in 0..shape_m.w { + let quant0 = shape_m.it(0, 0, workgroup_id.x * 4, j); + let quant1 = shape_m.it(0, 0, workgroup_id.x * 4 + 1, j); + let quant2 = shape_m.it(0, 0, workgroup_id.x * 4 + 2, j); + let quant3 = shape_m.it(0, 0, workgroup_id.x * 4 + 3, j); + + let parts0 = dequantize_q5_k_workgroup(m, quant0, lid); + let parts1 = dequantize_q5_k_workgroup(m, quant1, lid); + let parts2 = dequantize_q5_k_workgroup(m, quant2, lid); + let parts3 = dequantize_q5_k_workgroup(m, quant3, lid); + + let j_base = (j * 64) as usize; + let jj = (lid & 0xfffffff8) as usize; // == (lid / 8) * 8 + let vj_a = *v.at(j_base + lid as usize + jj); + let vj_b = *v.at(j_base + lid as usize + jj + 8); + + sum += Vec4::new( + parts0[0].dot(vj_a), + parts1[0].dot(vj_a), + parts2[0].dot(vj_a), + parts3[0].dot(vj_a), + ); + sum += Vec4::new( + parts0[1].dot(vj_b), + parts1[1].dot(vj_b), + parts2[1].dot(vj_b), + parts3[1].dot(vj_b), + ); + } + + #[cfg(feature = "subgroup_ops")] + { + let reduced = Vec4::new( + khal_std::sync::subgroup_f_add(sum.x), + khal_std::sync::subgroup_f_add(sum.y), + khal_std::sync::subgroup_f_add(sum.z), + khal_std::sync::subgroup_f_add(sum.w), + ); + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = reduced; + } + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *sketch.at_mut(lid as usize) = sum; + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // reduce_sum(lid as usize, 32, sketch); + reduce_sum(lid as usize, 16, sketch); + reduce_sum(lid as usize, 8, sketch); + reduce_sum(lid as usize, 4, sketch); + reduce_sum(lid as usize, 2, sketch); + reduce_sum(lid as usize, 1, sketch); + + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = *sketch.at(0); + } + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs new file mode 100644 index 0000000..cec6c19 --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs @@ -0,0 +1,227 @@ +//! Q6_Kx2 quantized GEMV shader. +//! +//! BlockQ6Kx2 contains two BlockQ6K blocks packed together (105 u32s total). +//! Each BlockQ6K: f16 scale, 64 bytes ql (low 4 bits), 32 bytes qh (high 2 bits), 16 bytes scales. + +use crate::utils::half::{unpack_half2x16, unpack_int4x8, unpack_uint4x8}; +use khal_std::glamx::{IVec4, Mat4, UVec3, UVec4, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: usize = 32; +// BlockQ6Kx2 size in u32s: 105 +const BLOCK_Q6KX2_SIZE: u32 = 105; + +#[inline] +fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) { + if index < stride { + let val = *sketch.at(index + stride); + *sketch.at_mut(index) += val; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +/// Select element from Vec4 based on index (0-3). +/// Avoids variable indexing which SPIR-V doesn't support well. +#[inline] +fn vec4_select(v: Vec4, idx: usize) -> f32 { + if idx == 0 { + v.x + } else if idx == 1 { + v.y + } else if idx == 2 { + v.z + } else { + v.w + } +} + +/// Dequantize Q6_Kx2 block for a workgroup thread. +/// Returns 4 Vec4s for each thread. +#[inline] +fn dequantize_q6_kx2_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 4] { + let _0xf = UVec4::splat(0xF); + let splat_6 = UVec4::splat(6); + let splat_4 = UVec4::splat(4); + let splat_3 = UVec4::splat(3); + let splat_2 = UVec4::splat(2); + let splat_32 = IVec4::splat(32); + + let data_id = (block_id * BLOCK_Q6KX2_SIZE) as usize; + + if k / 16 == 0 { + // Block A + // Its data goes from data[0] to half of data[52] + let d_a = unpack_half2x16(*m.at(data_id + 52)).x; + + const QL0: usize = 0; + const QH0: usize = 32; + const SC0: usize = 48; + + let i = ((k / 8) % 2) as usize; + let data0 = Vec4::new(1.0, 1.0, 1.0, 1.0) + * d_a + * unpack_int4x8(*m.at(data_id + SC0 + i * 2)).as_vec4(); + let data1 = Vec4::new(1.0, 1.0, 1.0, 1.0) + * d_a + * unpack_int4x8(*m.at(data_id + SC0 + i * 2 + 1)).as_vec4(); + + let l = (k % 8) as usize; + let is = l / 4; // NOTE: is is either 0 or 1 + + let qh = unpack_uint4x8(*m.at(data_id + l + QH0 + i * 8)); + let ql0 = unpack_uint4x8(*m.at(data_id + l + QL0 + i * 16)); + let ql32 = unpack_uint4x8(*m.at(data_id + l + QL0 + i * 16 + 8)); + + let q1 = ((ql0 & _0xf) | ((qh & splat_3) << splat_4)).as_ivec4() - splat_32; + let q2 = ((ql32 & _0xf) | (((qh >> splat_2) & splat_3) << splat_4)).as_ivec4() - splat_32; + let q3 = + ((ql0 >> splat_4) | (((qh >> splat_4) & splat_3) << splat_4)).as_ivec4() - splat_32; + let q4 = + ((ql32 >> splat_4) | (((qh >> splat_6) & splat_3) << splat_4)).as_ivec4() - splat_32; + + [ + Vec4::splat(vec4_select(data0, is)) * q1.as_vec4(), + Vec4::splat(vec4_select(data0, is + 2)) * q2.as_vec4(), + Vec4::splat(vec4_select(data1, is)) * q3.as_vec4(), + Vec4::splat(vec4_select(data1, is + 2)) * q4.as_vec4(), + ] + } else { + // Block B + // Its data goes from half of data[52] to data[104]. + // All values are starting with the u16 leftmost bits of the previous index. + let d_b = unpack_half2x16(*m.at(data_id + 104)).y; + + const QL0: usize = 53; + const QH0: usize = 53 + 32; + const SC0: usize = 53 + 48; + + let i = ((k / 8) % 2) as usize; + let l = (k % 8) as usize; + let isc0 = SC0 + i * 2; + let isc1 = SC0 + i * 2 + 1; + let data0 = Vec4::new(1.0, 1.0, 1.0, 1.0) + * d_b + * unpack_int4x8((*m.at(data_id + isc0 - 1) >> 16) | (*m.at(data_id + isc0) << 16)) + .as_vec4(); + let data1 = Vec4::new(1.0, 1.0, 1.0, 1.0) + * d_b + * unpack_int4x8((*m.at(data_id + isc1 - 1) >> 16) | (*m.at(data_id + isc1) << 16)) + .as_vec4(); + + let is = l / 4; // NOTE: is either 0 or 1 + + let iqh = l + QH0 + i * 8; + let iql0 = l + QL0 + i * 16; + let iql32 = l + QL0 + i * 16 + 8; + + let qh = unpack_uint4x8((*m.at(data_id + iqh - 1) >> 16) | (*m.at(data_id + iqh) << 16)); + let ql0 = unpack_uint4x8((*m.at(data_id + iql0 - 1) >> 16) | (*m.at(data_id + iql0) << 16)); + let ql32 = + unpack_uint4x8((*m.at(data_id + iql32 - 1) >> 16) | (*m.at(data_id + iql32) << 16)); + + let q1 = ((ql0 & _0xf) | ((qh & splat_3) << splat_4)).as_ivec4() - splat_32; + let q2 = ((ql32 & _0xf) | (((qh >> splat_2) & splat_3) << splat_4)).as_ivec4() - splat_32; + let q3 = + ((ql0 >> splat_4) | (((qh >> splat_4) & splat_3) << splat_4)).as_ivec4() - splat_32; + let q4 = + ((ql32 >> splat_4) | (((qh >> splat_6) & splat_3) << splat_4)).as_ivec4() - splat_32; + + [ + Vec4::splat(vec4_select(data0, is)) * q1.as_vec4(), + Vec4::splat(vec4_select(data0, is + 2)) * q2.as_vec4(), + Vec4::splat(vec4_select(data1, is)) * q3.as_vec4(), + Vec4::splat(vec4_select(data1, is + 2)) * q4.as_vec4(), + ] + } +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + let mut sum = Vec4::ZERO; + let lid = local_id.x; + + for j in 0..shape_m.w { + let quant0 = shape_m.it(0, 0, workgroup_id.x * 4, j); + let quant1 = shape_m.it(0, 0, workgroup_id.x * 4 + 1, j); + let quant2 = shape_m.it(0, 0, workgroup_id.x * 4 + 2, j); + let quant3 = shape_m.it(0, 0, workgroup_id.x * 4 + 3, j); + + let parts0 = dequantize_q6_kx2_workgroup(m, quant0, lid); + let parts1 = dequantize_q6_kx2_workgroup(m, quant1, lid); + let parts2 = dequantize_q6_kx2_workgroup(m, quant2, lid); + let parts3 = dequantize_q6_kx2_workgroup(m, quant3, lid); + + let j_base = (j * 128) as usize; + let jj = ((lid / 16) * 64 + ((lid / 8) % 2) * 32 + (lid % 8)) as usize; + let vj_a = *v.at(j_base + jj); + let vj_b = *v.at(j_base + jj + 8); + let vj_c = *v.at(j_base + jj + 16); + let vj_d = *v.at(j_base + jj + 24); + + let mat0 = Mat4::from_cols(parts0[0], parts1[0], parts2[0], parts3[0]); + let mat1 = Mat4::from_cols(parts0[1], parts1[1], parts2[1], parts3[1]); + let mat2 = Mat4::from_cols(parts0[2], parts1[2], parts2[2], parts3[2]); + let mat3 = Mat4::from_cols(parts0[3], parts1[3], parts2[3], parts3[3]); + sum += mat0.transpose() * vj_a + + mat1.transpose() * vj_b + + mat2.transpose() * vj_c + + mat3.transpose() * vj_d; + } + + #[cfg(feature = "subgroup_ops")] + { + let reduced = Vec4::new( + khal_std::sync::subgroup_f_add(sum.x), + khal_std::sync::subgroup_f_add(sum.y), + khal_std::sync::subgroup_f_add(sum.z), + khal_std::sync::subgroup_f_add(sum.w), + ); + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = reduced; + } + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *sketch.at_mut(lid as usize) = sum; + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // reduce_sum(lid as usize, 32, sketch); + reduce_sum(lid as usize, 16, sketch); + reduce_sum(lid as usize, 8, sketch); + reduce_sum(lid as usize, 4, sketch); + reduce_sum(lid as usize, 2, sketch); + reduce_sum(lid as usize, 1, sketch); + + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = *sketch.at(0); + } + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs new file mode 100644 index 0000000..ddbce38 --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs @@ -0,0 +1,175 @@ +//! Q8_0x2 quantized GEMV shader. +//! +//! BlockQ8_0x2 contains two BlockQ8_0 blocks (f16 scale + 32 x 8-bit signed quants each). + +use crate::utils::half::{unpack_half2x16, unpack_int4x8}; +use khal_std::glamx::{UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: usize = 32; +const BLOCK_Q8_0X2_SIZE: u32 = 17; // 17 u32s + +#[inline] +fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) { + if index < stride { + let val = *sketch.at(index + stride); + *sketch.at_mut(index) += val; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + let mut sum = Vec4::ZERO; + let lid = local_id.x; + + for j in 0..(shape_m.w / 2) { + // Calculate block indices for 4 matrix rows + let quant0 = + (shape_m.it(0, 0, workgroup_id.x * 4, j * 2 + lid / 16) * BLOCK_Q8_0X2_SIZE) as usize; + let quant1 = (shape_m.it(0, 0, workgroup_id.x * 4 + 1, j * 2 + lid / 16) + * BLOCK_Q8_0X2_SIZE) as usize; + let quant2 = (shape_m.it(0, 0, workgroup_id.x * 4 + 2, j * 2 + lid / 16) + * BLOCK_Q8_0X2_SIZE) as usize; + let quant3 = (shape_m.it(0, 0, workgroup_id.x * 4 + 3, j * 2 + lid / 16) + * BLOCK_Q8_0X2_SIZE) as usize; + + let j_base = (j * 32) as usize; + let vj = *v.at(j_base + lid as usize); + + if (lid / 8).is_multiple_of(2) { + // Dequantizing block 1 + let llid = (lid % 8) as usize; + let scale0 = unpack_half2x16(*m.at(quant0)).x; + let data0 = unpack_int4x8(*m.at(quant0 + llid) >> 16 | *m.at(quant0 + llid + 1) << 16); + let scale1 = unpack_half2x16(*m.at(quant1)).x; + let data1 = unpack_int4x8(*m.at(quant1 + llid) >> 16 | *m.at(quant1 + llid + 1) << 16); + let scale2 = unpack_half2x16(*m.at(quant2)).x; + let data2 = unpack_int4x8(*m.at(quant2 + llid) >> 16 | *m.at(quant2 + llid + 1) << 16); + let scale3 = unpack_half2x16(*m.at(quant3)).x; + let data3 = unpack_int4x8(*m.at(quant3 + llid) >> 16 | *m.at(quant3 + llid + 1) << 16); + + let row0 = Vec4::new( + data0.x as f32, + data0.y as f32, + data0.z as f32, + data0.w as f32, + ) * scale0; + let row1 = Vec4::new( + data1.x as f32, + data1.y as f32, + data1.z as f32, + data1.w as f32, + ) * scale1; + let row2 = Vec4::new( + data2.x as f32, + data2.y as f32, + data2.z as f32, + data2.w as f32, + ) * scale2; + let row3 = Vec4::new( + data3.x as f32, + data3.y as f32, + data3.z as f32, + data3.w as f32, + ) * scale3; + + sum += Vec4::new(row0.dot(vj), row1.dot(vj), row2.dot(vj), row3.dot(vj)); + } else { + // Dequantizing block 2 + let llid = (lid % 8) as usize; + let scale0 = unpack_half2x16(*m.at(quant0 + 8)).y; + let data0 = unpack_int4x8(*m.at(quant0 + llid + 9)); + let scale1 = unpack_half2x16(*m.at(quant1 + 8)).y; + let data1 = unpack_int4x8(*m.at(quant1 + llid + 9)); + let scale2 = unpack_half2x16(*m.at(quant2 + 8)).y; + let data2 = unpack_int4x8(*m.at(quant2 + llid + 9)); + let scale3 = unpack_half2x16(*m.at(quant3 + 8)).y; + let data3 = unpack_int4x8(*m.at(quant3 + llid + 9)); + + let row0 = Vec4::new( + data0.x as f32, + data0.y as f32, + data0.z as f32, + data0.w as f32, + ) * scale0; + let row1 = Vec4::new( + data1.x as f32, + data1.y as f32, + data1.z as f32, + data1.w as f32, + ) * scale1; + let row2 = Vec4::new( + data2.x as f32, + data2.y as f32, + data2.z as f32, + data2.w as f32, + ) * scale2; + let row3 = Vec4::new( + data3.x as f32, + data3.y as f32, + data3.z as f32, + data3.w as f32, + ) * scale3; + + sum += Vec4::new(row0.dot(vj), row1.dot(vj), row2.dot(vj), row3.dot(vj)); + } + } + + #[cfg(feature = "subgroup_ops")] + { + let reduced = Vec4::new( + khal_std::sync::subgroup_f_add(sum.x), + khal_std::sync::subgroup_f_add(sum.y), + khal_std::sync::subgroup_f_add(sum.z), + khal_std::sync::subgroup_f_add(sum.w), + ); + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = reduced; + } + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *sketch.at_mut(lid as usize) = sum; + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // reduce_sum(lid as usize, 32, sketch); // Only 32 threads + reduce_sum(lid as usize, 16, sketch); + reduce_sum(lid as usize, 8, sketch); + reduce_sum(lid as usize, 4, sketch); + reduce_sum(lid as usize, 2, sketch); + reduce_sum(lid as usize, 1, sketch); + + if lid == 0 { + let i_out = workgroup_id.x as usize; + *out.at_mut(i_out) = *sketch.at(0); + } + } +} diff --git a/vortx-shaders/src/ml/gemv_quant_q8_k.rs b/vortx-shaders/src/ml/gemv_quant_q8_k.rs new file mode 100644 index 0000000..777babb --- /dev/null +++ b/vortx-shaders/src/ml/gemv_quant_q8_k.rs @@ -0,0 +1,93 @@ +//! Q8_K quantized GEMV shader. +//! +//! BlockQ8K: f32 delta, 256 x 8-bit signed quants, 16 x 16-bit bsums. + +use crate::utils::half::unpack_int4x8; +use khal_std::glamx::{UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +const WORKGROUP_SIZE: u32 = 32; + +// BlockQ8K structure (repr(C), alignment 4): +// - d: f32 (1 u32) +// - qs: [i8; 256] (64 u32s) +// - bsums: [i16; 16] (8 u32s, no padding — 260 is already 2-byte aligned) +// Total: 73 u32s = 292 bytes +const BLOCK_Q8K_SIZE: u32 = 73; + +/// Dequantize a full BlockQ8K block. +#[inline] +fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 64] { + let mut result = [Vec4::ZERO; 64]; + + // d is stored as f32 directly + let d = f32::from_bits(*data.at(base)); + + #[allow(clippy::needless_range_loop)] + for j in 0..64 { + let qs = unpack_int4x8(*data.at(base + 1 + j)); + result[j] = Vec4::new(qs.x as f32, qs.y as f32, qs.z as f32, qs.w as f32) * d; + } + + result +} + +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn gemv( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_m: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] m: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_m = shapes.shape; + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_m = *shape_m.at(0); + + if invocation_id.x < shape_m.h { + let i_out = invocation_id.x as usize; + let mut sum = 0.0f32; + + for j in 0..shape_m.w { + let block_idx = (shape_m.it(0, 0, invocation_id.x, j) * BLOCK_Q8K_SIZE) as usize; + let dequant = dequantize_block(m, block_idx); + + // Unroll calculation with all block elements + let i_base = (j * 64) as usize; + + for k in (0u32..64).step_by(16) { + let k = k as usize; + sum += dequant[k].dot(*v.at(k + i_base)) + + dequant[k + 1].dot(*v.at(k + i_base + 1)) + + dequant[k + 2].dot(*v.at(k + i_base + 2)) + + dequant[k + 3].dot(*v.at(k + i_base + 3)) + + dequant[k + 4].dot(*v.at(k + i_base + 4)) + + dequant[k + 5].dot(*v.at(k + i_base + 5)) + + dequant[k + 6].dot(*v.at(k + i_base + 6)) + + dequant[k + 7].dot(*v.at(k + i_base + 7)) + + dequant[k + 8].dot(*v.at(k + i_base + 8)) + + dequant[k + 9].dot(*v.at(k + i_base + 9)) + + dequant[k + 10].dot(*v.at(k + i_base + 10)) + + dequant[k + 11].dot(*v.at(k + i_base + 11)) + + dequant[k + 12].dot(*v.at(k + i_base + 12)) + + dequant[k + 13].dot(*v.at(k + i_base + 13)) + + dequant[k + 14].dot(*v.at(k + i_base + 14)) + + dequant[k + 15].dot(*v.at(k + i_base + 15)); + } + } + + *out.at_mut(i_out) = Vec4::splat(sum); + } +} diff --git a/vortx-shaders/src/ml/get_rel_pos.rs b/vortx-shaders/src/ml/get_rel_pos.rs new file mode 100644 index 0000000..eab0112 --- /dev/null +++ b/vortx-shaders/src/ml/get_rel_pos.rs @@ -0,0 +1,125 @@ +//! Relative position computation. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::{Shapes1, Shapes2}; + +const WORKGROUP_SIZE: u32 = 128; + +/// Get relative position. +#[spirv_bindgen] +#[spirv(compute(threads(128, 1, 1)))] +pub fn get_rel_pos( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_result: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_source: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] result: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] source: &[f32], +) { + #[cfg(feature = "push_constants")] + let (shape_result, shape_source) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_result = *shape_result.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_source = *shape_source.at(0); + + if invocation_id.x >= shape_result.len() { + return; + } + + let id = shape_result.decompose(invocation_id.x); + let i = shape_result.it_vec(id) as usize; + let w = shape_result.h; + let pos = (w - id.x - 1) + id.z; + let j = shape_source.it(0, 0, pos, id.y) as usize; + + *result.at_mut(i) = *source.at(j); +} + +/// Add relative position phase 2. +#[spirv_bindgen] +#[spirv(compute(threads(128, 1, 1)))] +pub fn add_rel_pos_phase_b( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src1: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] src2: &[f32], +) { + #[cfg(feature = "push_constants")] + let shape_src1 = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src1 = *shape_src1.at(0); + + if invocation_id.x >= shape_src1.len() { + return; + } + + let id = shape_src1.decompose(invocation_id.x); + let jp0 = shape_src1.it_vec(id); + + // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L357-L359 + let src2_e = *src2.at(jp0 as usize); + let ne10 = shape_src1.w; + + let jdh = jp0 * ne10; + + for j in 0..ne10 { + *dst.at_mut((jdh + j) as usize) = *dst.at((jdh + j) as usize) + src2_e; + } +} + +/// Add relative position phase 1. +#[spirv_bindgen] +#[spirv(compute(threads(128, 1, 1)))] +pub fn add_rel_pos_phase_a( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src1: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] src1: &[f32], +) { + #[cfg(feature = "push_constants")] + let shape_src1 = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src1 = *shape_src1.at(0); + + if invocation_id.x >= shape_src1.len() { + return; + } + + let id = shape_src1.decompose(invocation_id.x); + let jp0 = shape_src1.it_vec(id); + + // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L357-L359 + let src1_e = *src1.at(jp0 as usize); + let ne10 = shape_src1.w; + + let jdh = jp0 * ne10; + let jdw = jdh - (ne10 - 1) * id.y; + + for j in 0..ne10 { + *dst.at_mut((jdw + j * ne10) as usize) = *dst.at((jdw + j * ne10) as usize) + src1_e; + } +} diff --git a/vortx-shaders/src/ml/im2col.rs b/vortx-shaders/src/ml/im2col.rs new file mode 100644 index 0000000..962605f --- /dev/null +++ b/vortx-shaders/src/ml/im2col.rs @@ -0,0 +1,118 @@ +//! Image to column transformation. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +const WORKGROUP_SIZE: u32 = 32; +const NUM_ITER: u32 = 512 / WORKGROUP_SIZE; + +/// Im2Col parameters. +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct Im2ColParams { + pub batch_offset: u32, + pub offset_delta: u32, + #[allow(non_snake_case)] + pub IC: u32, + #[allow(non_snake_case)] + pub IW: u32, + #[allow(non_snake_case)] + pub IH: u32, + #[allow(non_snake_case)] + pub OW: u32, + #[allow(non_snake_case)] + pub OH: u32, + #[allow(non_snake_case)] + pub KW: u32, + #[allow(non_snake_case)] + pub KH: u32, + pub pelements: u32, + #[allow(non_snake_case)] + pub CHW: u32, + pub s0: i32, + pub s1: i32, + pub p0: i32, + pub p1: i32, + pub d0: i32, + pub d1: i32, +} + +/// Im2Col transformation. +#[spirv_bindgen] +#[spirv(compute(threads(32, 1, 1)))] +pub fn im2col( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] params: &[Im2ColParams], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] in_tensor: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] out_tensor: &mut [f32], +) { + // Load params from storage buffer to local variable (enables LICM) + let params = *params.at(0); + + let gidx = invocation_id.x; + let oh = invocation_id.y; + let batch = invocation_id.z / params.IC; + let ic = invocation_id.z % params.IC; + + let src_base = ic * params.offset_delta + batch * params.batch_offset; + let dst_base = + ((batch * params.OH + oh) * params.OW) * params.CHW + ic * (params.KW * params.KH); + let oh_s1 = oh as i32 * params.s1; + let ksize = params.OW * if params.KH > 1 { params.KW } else { 1 }; + + let base_linear_idx = gidx * NUM_ITER; + + let max_ky = ksize / params.OW; + + let mut current_kx = base_linear_idx / ksize; + let rem = base_linear_idx - (current_kx * ksize); + let mut current_ky = rem / params.OW; + let mut current_ix = rem % params.OW; + + let mut values = [0.0f32; NUM_ITER as usize]; + let mut offset_dst = [0u32; NUM_ITER as usize]; + + for idx in 0..NUM_ITER { + let linear_idx = base_linear_idx + idx; + + if linear_idx >= params.pelements { + continue; + } + + let iiw = + (current_ix as i32 * params.s0 + current_kx as i32 * params.d0 - params.p0) as u32; + let iih = (oh_s1 + current_ky as i32 * params.d1 - params.p1) as u32; + + offset_dst[idx as usize] = + dst_base + current_ix * params.CHW + current_ky * params.KW + current_kx; + + if iih < params.IH && iiw < params.IW { + values[idx as usize] = *in_tensor.at((src_base + iih * params.IW + iiw) as usize); + } + + current_ix += 1; + if current_ix == params.OW { + current_ix = 0; + current_ky += 1; + if current_ky == max_ky { + current_ky = 0; + current_kx += 1; + } + } + } + + for idx in 0..NUM_ITER { + let linear_idx = base_linear_idx + idx; + + if linear_idx >= params.pelements { + continue; + } + + *out_tensor.at_mut(offset_dst[idx as usize] as usize) = values[idx as usize]; + } +} diff --git a/vortx-shaders/src/ml/layernorm.rs b/vortx-shaders/src/ml/layernorm.rs new file mode 100644 index 0000000..f3c8628 --- /dev/null +++ b/vortx-shaders/src/ml/layernorm.rs @@ -0,0 +1,240 @@ +//! Layer normalization kernels. + +use crate::utils::iterators::StepRng; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; + +#[cfg(feature = "subgroup_ops")] +const WORKGROUP_SIZE: usize = 32; +#[cfg(not(feature = "subgroup_ops"))] +const WORKGROUP_SIZE: usize = 128; +const NUDGE_FACTOR: f32 = 1.0e-6; + +#[inline] +fn reduce_sum(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + if index < stride { + let val = *workspace.at(index + stride); + *workspace.at_mut(index) += val; + } +} + +/// Layer normalization on columns. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn layernorm_cols( + #[spirv(workgroup_id)] wid: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] the_mean: &mut f32, + #[spirv(workgroup)] scale: &mut f32, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + in_shape: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + out_shape: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] input: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] output: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (in_shape, out_shape) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let in_shape = *in_shape.at(0); + #[cfg(not(feature = "push_constants"))] + let out_shape = *out_shape.at(0); + + let thread_id = local_id.x as usize; + + // Compute the MEAN + let data_len = in_shape.h; + *workspace.at_mut(thread_id) = 0.0; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *input.at(in_shape.it(wid.z, wid.y, i, wid.x) as usize); + *workspace.at_mut(thread_id) = *workspace.at(thread_id) + val_i; + } + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(*workspace.at(thread_id)); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 64, workspace); + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *the_mean = sum / data_len as f32; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *the_mean = *workspace.at(0) / data_len as f32; + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Compute the SQUARED NORM + *workspace.at_mut(thread_id) = 0.0; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *input.at(in_shape.it(wid.z, wid.y, i, wid.x) as usize) - *the_mean; + *workspace.at_mut(thread_id) = *workspace.at(thread_id) + val_i * val_i; + } + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(*workspace.at(thread_id)); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 64, workspace); + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + let variance = sum / data_len as f32; + #[cfg(not(feature = "subgroup_ops"))] + let variance = *workspace.at(0) / data_len as f32; + + *scale = 1.0 / (variance + NUDGE_FACTOR).sqrt(); + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Apply the scale. + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = in_shape.it(wid.z, wid.y, i, wid.x) as usize; + let iout = out_shape.it(wid.z, wid.y, i, wid.x) as usize; + *output.at_mut(iout) = (*input.at(ii) - *the_mean) * *scale; + } +} + +/// Layer normalization on rows. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn layernorm_rows( + #[spirv(workgroup_id)] wid: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] the_mean: &mut f32, + #[spirv(workgroup)] scale: &mut f32, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + in_shape: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + out_shape: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] input: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] output: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (in_shape, out_shape) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let in_shape = *in_shape.at(0); + #[cfg(not(feature = "push_constants"))] + let out_shape = *out_shape.at(0); + + let thread_id = local_id.x as usize; + + // Compute the MEAN + let data_len = in_shape.w; + *workspace.at_mut(thread_id) = 0.0; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *input.at(in_shape.it(wid.z, wid.y, wid.x, i) as usize); + *workspace.at_mut(thread_id) = *workspace.at(thread_id) + val_i; + } + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(*workspace.at(thread_id)); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 64, workspace); + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *the_mean = sum / data_len as f32; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *the_mean = *workspace.at(0) / data_len as f32; + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Compute the SQUARED NORM + *workspace.at_mut(thread_id) = 0.0; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *input.at(in_shape.it(wid.z, wid.y, wid.x, i) as usize) - *the_mean; + *workspace.at_mut(thread_id) = *workspace.at(thread_id) + val_i * val_i; + } + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(*workspace.at(thread_id)); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 64, workspace); + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + let variance = sum / data_len as f32; + #[cfg(not(feature = "subgroup_ops"))] + let variance = *workspace.at(0) / data_len as f32; + + *scale = 1.0 / (variance + NUDGE_FACTOR).sqrt(); + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Apply the scale. + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = in_shape.it(wid.z, wid.y, wid.x, i) as usize; + let iout = out_shape.it(wid.z, wid.y, wid.x, i) as usize; + *output.at_mut(iout) = (*input.at(ii) - *the_mean) * *scale; + } +} diff --git a/vortx-shaders/src/ml/mod.rs b/vortx-shaders/src/ml/mod.rs new file mode 100644 index 0000000..703f007 --- /dev/null +++ b/vortx-shaders/src/ml/mod.rs @@ -0,0 +1,55 @@ +// #![allow(clippy::too_many_arguments)] +// // `spirv_bindgen` generates host-side dispatch code that performs `% workgroup_size`, +// // which triggers this lint when a workgroup dimension is 1. +// #![allow(clippy::modulo_one)] +// #![allow(unexpected_cfgs)] +// // Shader entry points and their constants appear dead on host but are used on GPU. +// #![allow(dead_code, non_snake_case)] + +// TODO: keep the modules private? +pub mod batched_multiquery_attention; +pub mod concat; +pub mod conv2d; +pub mod conv_transpose_2d; +pub mod fused_attention; +pub mod gather; +pub mod gemv_quant_q4_0x2; +pub mod gemv_quant_q4_1x2; +pub mod gemv_quant_q4_k; +pub mod gemv_quant_q5_0x2; +pub mod gemv_quant_q5_1x2; +pub mod gemv_quant_q5_k; +pub mod gemv_quant_q6_kx2; +pub mod gemv_quant_q8_0x2; +pub mod gemv_quant_q8_k; +pub mod get_rel_pos; +pub mod im2col; +pub mod layernorm; +pub mod pool2d; +pub mod reduce_axis; +pub mod rms_norm; +pub mod rope; +pub mod select; +pub mod silu; +pub mod softmax; +pub mod unary; +pub mod win_part; + +pub use batched_multiquery_attention::*; +pub use concat::*; +pub use conv2d::*; +pub use conv_transpose_2d::*; +pub use fused_attention::*; +pub use gather::*; +pub use get_rel_pos::*; +pub use im2col::*; +pub use layernorm::*; +pub use pool2d::*; +pub use reduce_axis::*; +pub use rms_norm::*; +pub use rope::*; +pub use select::*; +pub use silu::*; +pub use softmax::*; +pub use unary::*; +pub use win_part::*; diff --git a/vortx-shaders/src/ml/pool2d.rs b/vortx-shaders/src/ml/pool2d.rs new file mode 100644 index 0000000..5868b76 --- /dev/null +++ b/vortx-shaders/src/ml/pool2d.rs @@ -0,0 +1,283 @@ +//! 2D Pooling operations (MaxPool2d, AvgPool2d). +//! +//! Input shape: [N, C, H, W] (NCHW format) +//! Output shape: [N, C, H_out, W_out] +//! +//! Parameters in params buffer: +//! \[0\] input_height +//! \[1\] input_width +//! \[2\] output_height +//! \[3\] output_width +//! \[4\] kernel_h +//! \[5\] kernel_w +//! \[6\] stride_h +//! \[7\] stride_w +//! \[8\] pad_h +//! \[9\] pad_w +//! \[10\] channels +//! \[11\] batch_size + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::utils::limits::MAX_NUM_WORKGROUPS; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// MaxPool2d - compute max over a 2D window. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn max_pool_2d( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] params: &[u32], +) { + let input_h = *params.at(0); + let input_w = *params.at(1); + let output_h = *params.at(2); + let output_w = *params.at(3); + let kernel_h = *params.at(4); + let kernel_w = *params.at(5); + let stride_h = *params.at(6); + let stride_w = *params.at(7); + let pad_h = *params.at(8); + let pad_w = *params.at(9); + let channels = *params.at(10); + let batch_size = *params.at(11); + + let output_len = batch_size * channels * output_h * output_w; + + for thread_id in (invocation_id.x..output_len).step_by(MAX_NUM_THREADS as usize) { + // Decompose output index into [n, c, oh, ow] + let ow = thread_id % output_w; + let oh = (thread_id / output_w) % output_h; + let c = (thread_id / (output_w * output_h)) % channels; + let n = thread_id / (output_w * output_h * channels); + + // Find the input window + let h_start_signed = (oh * stride_h) as i32 - pad_h as i32; + let w_start_signed = (ow * stride_w) as i32 - pad_w as i32; + + // Clamp to valid input range + let h_start = if h_start_signed < 0 { + 0u32 + } else { + h_start_signed as u32 + }; + let w_start = if w_start_signed < 0 { + 0u32 + } else { + w_start_signed as u32 + }; + let h_end_unclamped = (h_start_signed + kernel_h as i32) as u32; + let w_end_unclamped = (w_start_signed + kernel_w as i32) as u32; + let h_end = if h_end_unclamped > input_h { + input_h + } else { + h_end_unclamped + }; + let w_end = if w_end_unclamped > input_w { + input_w + } else { + w_end_unclamped + }; + + // Initialize max with a very small value (SPIR-V doesn't support infinity literals) + let mut max_val = -3.4028235e+38_f32; // Close to f32::MIN + + // Iterate over the pooling window + for ih in h_start..h_end { + for iw in w_start..w_end { + // Input index: n * (C * H * W) + c * (H * W) + ih * W + iw + let i_src = + (n * channels * input_h * input_w + c * input_h * input_w + ih * input_w + iw) + as usize; + let val = *src.at(i_src); + if val > max_val { + max_val = val; + } + } + } + + // Handle edge case where window was entirely in padding + if max_val < -3.4e+38_f32 { + max_val = 0.0; + } + + *dest.at_mut(thread_id as usize) = max_val; + } +} + +/// AvgPool2d - compute average over a 2D window. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn avg_pool_2d( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] params: &[u32], +) { + let input_h = *params.at(0); + let input_w = *params.at(1); + let output_h = *params.at(2); + let output_w = *params.at(3); + let kernel_h = *params.at(4); + let kernel_w = *params.at(5); + let stride_h = *params.at(6); + let stride_w = *params.at(7); + let pad_h = *params.at(8); + let pad_w = *params.at(9); + let channels = *params.at(10); + let batch_size = *params.at(11); + // params[12] = count_include_pad (0 or 1) + let count_include_pad = *params.at(12); + + let output_len = batch_size * channels * output_h * output_w; + + for thread_id in (invocation_id.x..output_len).step_by(MAX_NUM_THREADS as usize) { + // Decompose output index into [n, c, oh, ow] + let ow = thread_id % output_w; + let oh = (thread_id / output_w) % output_h; + let c = (thread_id / (output_w * output_h)) % channels; + let n = thread_id / (output_w * output_h * channels); + + // Find the input window + let h_start_signed = (oh * stride_h) as i32 - pad_h as i32; + let w_start_signed = (ow * stride_w) as i32 - pad_w as i32; + + // Clamp to valid input range + let h_start = if h_start_signed < 0 { + 0u32 + } else { + h_start_signed as u32 + }; + let w_start = if w_start_signed < 0 { + 0u32 + } else { + w_start_signed as u32 + }; + let h_end_unclamped = (h_start_signed + kernel_h as i32) as u32; + let w_end_unclamped = (w_start_signed + kernel_w as i32) as u32; + let h_end = if h_end_unclamped > input_h { + input_h + } else { + h_end_unclamped + }; + let w_end = if w_end_unclamped > input_w { + input_w + } else { + w_end_unclamped + }; + + // Sum over the pooling window + let mut sum: f32 = 0.0; + let mut count: u32 = 0; + + for ih in h_start..h_end { + for iw in w_start..w_end { + // Input index: n * (C * H * W) + c * (H * W) + ih * W + iw + let i_src = + (n * channels * input_h * input_w + c * input_h * input_w + ih * input_w + iw) + as usize; + sum += *src.at(i_src); + count += 1; + } + } + + // Compute average + let divisor = if count_include_pad != 0 { + kernel_h * kernel_w + } else { + count + }; + + let avg = if divisor > 0 { + sum / (divisor as f32) + } else { + 0.0 + }; + *dest.at_mut(thread_id as usize) = avg; + } +} + +/// GlobalAvgPool2d - average over entire spatial dimensions. +/// Input: [N, C, H, W], Output: [N, C, 1, 1] +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn global_avg_pool_2d( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] params: &[u32], // [input_h, input_w, channels, batch_size] +) { + let input_h = *params.at(0); + let input_w = *params.at(1); + let channels = *params.at(2); + let batch_size = *params.at(3); + + let output_len = batch_size * channels; + let spatial_size = input_h * input_w; + + for thread_id in (invocation_id.x..output_len).step_by(MAX_NUM_THREADS as usize) { + // Output index: [n, c] + let c = thread_id % channels; + let n = thread_id / channels; + + // Sum over all spatial elements + let mut sum: f32 = 0.0; + for ih in 0..input_h { + for iw in 0..input_w { + let i_src = + (n * channels * input_h * input_w + c * input_h * input_w + ih * input_w + iw) + as usize; + sum += *src.at(i_src); + } + } + + *dest.at_mut(thread_id as usize) = sum / (spatial_size as f32); + } +} + +/// GlobalMaxPool2d - max over entire spatial dimensions. +/// Input: [N, C, H, W], Output: [N, C, 1, 1] +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn global_max_pool_2d( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] params: &[u32], // [input_h, input_w, channels, batch_size] +) { + let input_h = *params.at(0); + let input_w = *params.at(1); + let channels = *params.at(2); + let batch_size = *params.at(3); + + let output_len = batch_size * channels; + + for thread_id in (invocation_id.x..output_len).step_by(MAX_NUM_THREADS as usize) { + // Output index: [n, c] + let c = thread_id % channels; + let n = thread_id / channels; + + // Find max over all spatial elements + let i_src_init = (n * channels * input_h * input_w + c * input_h * input_w) as usize; + let mut max_val = *src.at(i_src_init); + + for ih in 0..input_h { + for iw in 0..input_w { + let i_src = + (n * channels * input_h * input_w + c * input_h * input_w + ih * input_w + iw) + as usize; + let val = *src.at(i_src); + if val > max_val { + max_val = val; + } + } + } + + *dest.at_mut(thread_id as usize) = max_val; + } +} diff --git a/vortx-shaders/src/ml/reduce_axis.rs b/vortx-shaders/src/ml/reduce_axis.rs new file mode 100644 index 0000000..d167279 --- /dev/null +++ b/vortx-shaders/src/ml/reduce_axis.rs @@ -0,0 +1,296 @@ +//! Axis-based reduction operations (ReduceSum, ReduceMean, etc.) + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; +use crate::utils::limits::MAX_NUM_WORKGROUPS; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Reduce sum along a single axis. +/// +/// Each thread handles one element in the output. It iterates over all +/// elements along the reduce axis in the input and sums them. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reduce_sum_axis( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] params: &[u32], // [axis, reduce_size] +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *params.at(0); + let reduce_size = *params.at(1); + + for thread_id in (invocation_id.x..shape_dest.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index in output + let id_dest = shape_dest.decompose(thread_id); + + // Build source coordinates - start with output coords + let mut id_src = id_dest; + + // Sum over all elements along the reduce axis + let mut sum: f32 = 0.0; + for i in 0..reduce_size { + match axis { + 0 => id_src.x = i, + 1 => id_src.y = i, + 2 => id_src.z = i, + _ => id_src.w = i, + } + let i_src = shape_src.it_vec(id_src) as usize; + sum += *src.at(i_src); + } + + let i_dest = shape_dest.it_vec(id_dest) as usize; + *dest.at_mut(i_dest) = sum; + } +} + +/// Reduce mean along a single axis. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reduce_mean_axis( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] params: &[u32], // [axis, reduce_size] +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *params.at(0); + let reduce_size = *params.at(1); + + for thread_id in (invocation_id.x..shape_dest.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index in output + let id_dest = shape_dest.decompose(thread_id); + + // Build source coordinates - start with output coords + let mut id_src = id_dest; + + // Sum over all elements along the reduce axis + let mut sum: f32 = 0.0; + for i in 0..reduce_size { + match axis { + 0 => id_src.x = i, + 1 => id_src.y = i, + 2 => id_src.z = i, + _ => id_src.w = i, + } + let i_src = shape_src.it_vec(id_src) as usize; + sum += *src.at(i_src); + } + + let i_dest = shape_dest.it_vec(id_dest) as usize; + *dest.at_mut(i_dest) = sum / (reduce_size as f32); + } +} + +/// Reduce max along a single axis. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reduce_max_axis( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] params: &[u32], // [axis, reduce_size] +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *params.at(0); + let reduce_size = *params.at(1); + + for thread_id in (invocation_id.x..shape_dest.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index in output + let id_dest = shape_dest.decompose(thread_id); + + // Build source coordinates - start with output coords + let mut id_src = id_dest; + + // Find max over all elements along the reduce axis + id_src.x = 0; + id_src.y = 0; + id_src.z = 0; + id_src.w = 0; + match axis { + 0 => id_src.x = 0, + 1 => id_src.y = 0, + 2 => id_src.z = 0, + _ => id_src.w = 0, + } + // Restore non-axis coordinates from dest + match axis { + 0 => { + id_src.y = id_dest.y; + id_src.z = id_dest.z; + id_src.w = id_dest.w; + } + 1 => { + id_src.x = id_dest.x; + id_src.z = id_dest.z; + id_src.w = id_dest.w; + } + 2 => { + id_src.x = id_dest.x; + id_src.y = id_dest.y; + id_src.w = id_dest.w; + } + _ => { + id_src.x = id_dest.x; + id_src.y = id_dest.y; + id_src.z = id_dest.z; + } + } + + let i_src_init = shape_src.it_vec(id_src) as usize; + let mut max_val = *src.at(i_src_init); + + for i in 1..reduce_size { + match axis { + 0 => id_src.x = i, + 1 => id_src.y = i, + 2 => id_src.z = i, + _ => id_src.w = i, + } + let i_src = shape_src.it_vec(id_src) as usize; + let val = *src.at(i_src); + if val > max_val { + max_val = val; + } + } + + let i_dest = shape_dest.it_vec(id_dest) as usize; + *dest.at_mut(i_dest) = max_val; + } +} + +/// Reduce min along a single axis. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reduce_min_axis( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] params: &[u32], // [axis, reduce_size] +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let axis = *params.at(0); + let reduce_size = *params.at(1); + + for thread_id in (invocation_id.x..shape_dest.len()).step_by(MAX_NUM_THREADS as usize) { + // Decompose linear index in output + let id_dest = shape_dest.decompose(thread_id); + + // Build source coordinates + let mut id_src = id_dest; + match axis { + 0 => { + id_src.y = id_dest.y; + id_src.z = id_dest.z; + id_src.w = id_dest.w; + id_src.x = 0; + } + 1 => { + id_src.x = id_dest.x; + id_src.z = id_dest.z; + id_src.w = id_dest.w; + id_src.y = 0; + } + 2 => { + id_src.x = id_dest.x; + id_src.y = id_dest.y; + id_src.w = id_dest.w; + id_src.z = 0; + } + _ => { + id_src.x = id_dest.x; + id_src.y = id_dest.y; + id_src.z = id_dest.z; + id_src.w = 0; + } + } + + let i_src_init = shape_src.it_vec(id_src) as usize; + let mut min_val = *src.at(i_src_init); + + for i in 1..reduce_size { + match axis { + 0 => id_src.x = i, + 1 => id_src.y = i, + 2 => id_src.z = i, + _ => id_src.w = i, + } + let i_src = shape_src.it_vec(id_src) as usize; + let val = *src.at(i_src); + if val < min_val { + min_val = val; + } + } + + let i_dest = shape_dest.it_vec(id_dest) as usize; + *dest.at_mut(i_dest) = min_val; + } +} diff --git a/vortx-shaders/src/ml/rms_norm.rs b/vortx-shaders/src/ml/rms_norm.rs new file mode 100644 index 0000000..f70f2c5 --- /dev/null +++ b/vortx-shaders/src/ml/rms_norm.rs @@ -0,0 +1,121 @@ +//! RMS normalization kernel. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes3; + +#[cfg(feature = "subgroup_ops")] +const WORKGROUP_SIZE: usize = 32; +#[cfg(not(feature = "subgroup_ops"))] +const WORKGROUP_SIZE: usize = 128; + +/// RMS normalization configuration. +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct RmsNormConfig { + pub nudge_factor: f32, +} + +#[inline] +fn reduce_sum(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + if index < stride { + *workspace.at_mut(index) += *workspace.at(index + stride); + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); +} + +fn magnitude_squared( + thread_id: u32, + shape_v: Shape, + v: &[f32], + workspace: &mut [f32; WORKGROUP_SIZE], +) -> f32 { + let thread_id_usize = thread_id as usize; + *workspace.at_mut(thread_id_usize) = 0.0; + + let mut i = thread_id; + while i < shape_v.w { + let val_i = v.at(shape_v.it(0, 0, 0, i) as usize); + *workspace.at_mut(thread_id_usize) += val_i * val_i; + i += WORKGROUP_SIZE as u32; + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(*workspace.at(thread_id_usize)); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id_usize, 64, workspace); + reduce_sum(thread_id_usize, 32, workspace); + reduce_sum(thread_id_usize, 16, workspace); + reduce_sum(thread_id_usize, 8, workspace); + reduce_sum(thread_id_usize, 4, workspace); + reduce_sum(thread_id_usize, 2, workspace); + reduce_sum(thread_id_usize, 1, workspace); + } + + #[cfg(feature = "subgroup_ops")] + if thread_id_usize == 0 { + *workspace.at_mut(0) = sum; + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + *workspace.at(0) +} + +/// RMS normalization. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] +pub fn rms_norm( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes3, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_v: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_w: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_out: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] v: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] w: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] out: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] config: &[RmsNormConfig], +) { + #[cfg(feature = "push_constants")] + let (shape_v, shape_w, shape_out) = (shapes.shape_out, shapes.shape_lhs, shapes.shape_rhs); + // Load shapes and config from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_v = *shape_v.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_w = *shape_w.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_out = *shape_out.at(0); + let config = *config.at(0); + + let magnitude_sq = magnitude_squared(invocation_id.x, shape_v, v, workspace); + + let len = shape_v.w; + let rms = 1.0 / ((magnitude_sq / len as f32) + config.nudge_factor).sqrt(); + + for i in (invocation_id.x..len).step_by(WORKGROUP_SIZE) { + *out.at_mut(shape_out.it(0, 0, 0, i) as usize) = + (*v.at(shape_v.it(0, 0, 0, i) as usize) * rms) * *w.at(shape_w.it(0, 0, 0, i) as usize); + } +} diff --git a/vortx-shaders/src/ml/rope.rs b/vortx-shaders/src/ml/rope.rs new file mode 100644 index 0000000..111407e --- /dev/null +++ b/vortx-shaders/src/ml/rope.rs @@ -0,0 +1,141 @@ +//! Rotary Positional Encoding (RoPE). + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; + +const WORKGROUP_SIZE: u32 = 64; + +/// RoPE configuration. +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct RoPEConfig { + pub head_size: u32, + pub kv_dim: u32, + pub pos: u32, + pub base_freq: f32, +} + +/// 2D rotation. +#[derive(Clone, Copy)] +struct Rotation2 { + cos: f32, + sin: f32, +} + +#[inline] +fn rot2(angle: f32) -> Rotation2 { + Rotation2 { + cos: angle.cos(), + sin: angle.sin(), + } +} + +#[inline] +fn rotate2(r: Rotation2, vx: f32, vy: f32) -> (f32, f32) { + (r.cos * vx - r.sin * vy, r.sin * vx + r.cos * vy) +} + +/// Standard RoPE. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn rope( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_q: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_k: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] config: &[RoPEConfig], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] in_out_q: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] in_out_k: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_q, shape_k) = (shapes.shape_a, shapes.shape_b); + // Load shapes and config from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_q = *shape_q.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_k = *shape_k.at(0); + let config = *config.at(0); + + let i = invocation_id.x; + let head_dim = ((i * 2) % config.head_size) as f32; + let theta = config.base_freq.powf(-head_dim / config.head_size as f32); + let m_theta = config.pos as f32 * theta; + let rot = rot2(m_theta); + + let iq = shape_q.it(0, 0, i * 2, 0) as usize; + let q_rotated = rotate2(rot, *in_out_q.at(iq), *in_out_q.at(iq + 1)); + *in_out_q.at_mut(iq) = q_rotated.0; + *in_out_q.at_mut(iq + 1) = q_rotated.1; + + if i * 2 < config.kv_dim { + let ik = shape_k.it(0, 0, i * 2, 0) as usize; + let k_rotated = rotate2(rot, *in_out_k.at(ik), *in_out_k.at(ik + 1)); + *in_out_k.at_mut(ik) = k_rotated.0; + *in_out_k.at_mut(ik + 1) = k_rotated.1; + } +} + +/// NeoX-style RoPE. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn rope_neox( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_q: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_k: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] config: &[RoPEConfig], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] in_out_q: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] in_out_k: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_q, shape_k) = (shapes.shape_a, shapes.shape_b); + // Load shapes and config from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_q = *shape_q.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_k = *shape_k.at(0); + let config = *config.at(0); + + let i = invocation_id.x; + let head_dim = ((i * 2) % config.head_size) as f32; + let theta = config.base_freq.powf(-head_dim / config.head_size as f32); + let m_theta = config.pos as f32 * theta; + let rot = rot2(m_theta); + + let head_id = (i * 2) / config.head_size; + let shift = config.head_size / 2; + + let iq = shape_q.it(0, 0, i + head_id * config.head_size / 2, 0) as usize; + let q_rotated = rotate2(rot, *in_out_q.at(iq), *in_out_q.at(iq + shift as usize)); + *in_out_q.at_mut(iq) = q_rotated.0; + *in_out_q.at_mut(iq + shift as usize) = q_rotated.1; + + if i * 2 < config.kv_dim { + let ik = shape_k.it(0, 0, i + head_id * config.head_size / 2, 0) as usize; + let k_rotated = rotate2(rot, *in_out_k.at(ik), *in_out_k.at(ik + shift as usize)); + *in_out_k.at_mut(ik) = k_rotated.0; + *in_out_k.at_mut(ik + shift as usize) = k_rotated.1; + } +} diff --git a/vortx-shaders/src/ml/select.rs b/vortx-shaders/src/ml/select.rs new file mode 100644 index 0000000..2a1423f --- /dev/null +++ b/vortx-shaders/src/ml/select.rs @@ -0,0 +1,56 @@ +//! Select operation: selects elements from a source tensor based on indices. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; + +/// Select elements from source based on indices and write to destination. +/// +/// For each element in dest at position (i, j, k, l), this reads the index +/// from idx\[i\] and then copies src\[idx\[i\], j, k, l\] to dest\[i, j, k, l\]. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn select( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_dest: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dest: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] src: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] idx: &[u32], +) { + #[cfg(feature = "push_constants")] + let (shape_dest, shape_src) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dest = *shape_dest.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + let thread_id = invocation_id.x; + if thread_id < shape_dest.len() { + // Decompose linear index to 4D coordinates + let id = shape_dest.decompose(thread_id); + + // Compute destination index + let i_dest = shape_dest.it_vec(id) as usize; + + // Replace x coordinate with the index lookup + let mut id_src = id; + id_src.x = *idx.at(id.x as usize); + + // Compute source index with wrapping + let i_src = shape_src.it_repeating_vec(id_src) as usize; + + *dest.at_mut(i_dest) = *src.at(i_src); + } +} diff --git a/vortx-shaders/src/ml/silu.rs b/vortx-shaders/src/ml/silu.rs new file mode 100644 index 0000000..ef07e23 --- /dev/null +++ b/vortx-shaders/src/ml/silu.rs @@ -0,0 +1,53 @@ +//! SiLU (Swish) activation function. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; + +const WORKGROUP_SIZE: u32 = 64; + +/// SwiGLU non-linearity. +#[inline] +fn swish(x: f32, beta: f32) -> f32 { + // This is the swiglu function from https://youtu.be/Mn_9W1nCFLo?si=LT6puSAfzgpP6ydz&t=3973 + x / (1.0 + (-beta * x).exp()) +} + +/// SiLU activation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn silu( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_a: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_b: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] in_out_a: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] in_b: &[f32], +) { + #[cfg(feature = "push_constants")] + let (shape_a, shape_b) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_a = *shape_a.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_b = *shape_b.at(0); + + if invocation_id.x < shape_a.w { + let ia = shape_a.it(0, 0, 0, invocation_id.x) as usize; + let ib = shape_b.it(0, 0, 0, invocation_id.x) as usize; + let lhs = *in_out_a.at(ia); + let rhs = *in_b.at(ib); + *in_out_a.at_mut(ia) = rhs * swish(lhs, 1.0); + } +} diff --git a/vortx-shaders/src/ml/softmax.rs b/vortx-shaders/src/ml/softmax.rs new file mode 100644 index 0000000..1adcee5 --- /dev/null +++ b/vortx-shaders/src/ml/softmax.rs @@ -0,0 +1,264 @@ +//! Softmax and log-softmax kernels. + +use crate::utils::iterators::StepRng; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; + +#[cfg(feature = "subgroup_ops")] +const WORKGROUP_SIZE: usize = 32; +#[cfg(not(feature = "subgroup_ops"))] +const WORKGROUP_SIZE: usize = 64; + +#[inline] +fn reduce_max(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + if index < stride { + *workspace.at_mut(index) = workspace.at(index).max(*workspace.at(index + stride)); + } +} + +#[inline] +fn reduce_sum(index: usize, stride: usize, workspace: &mut [f32; WORKGROUP_SIZE]) { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + if index < stride { + *workspace.at_mut(index) = *workspace.at(index) + *workspace.at(index + stride); + } +} + +/// Softmax on columns. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(64, 1, 1))))] +pub fn softmax( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] the_max: &mut f32, + #[spirv(workgroup)] denominator: &mut f32, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] in_out_mat: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape = *shape.at(0); + + let j = workgroup_id.x; + let k = workgroup_id.y; + let l = workgroup_id.z; + let thread_id = local_id.x as usize; + + // Compute the MAX + let data_len = shape.w; + let mut my_max = [-1.0e38f32]; + + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *in_out_mat.at(shape.it(l, k, j, i) as usize); + my_max[0] = my_max[0].max(val_i); + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *workspace.at_mut(thread_id) = my_max[0]; + } + + #[cfg(feature = "subgroup_ops")] + let max_val = khal_std::sync::subgroup_f_max(my_max[0]); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_max(thread_id, 32, workspace); + reduce_max(thread_id, 16, workspace); + reduce_max(thread_id, 8, workspace); + reduce_max(thread_id, 4, workspace); + reduce_max(thread_id, 2, workspace); + reduce_max(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *the_max = max_val; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *the_max = *workspace.at(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Compute the denominator (sum of exponential). + let mut my_denominator = [0.0f32]; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = shape.it(l, k, j, i) as usize; + let val_i = *in_out_mat.at(ii); + let centered_val = val_i - *the_max; + let exp_i = centered_val.exp(); + my_denominator[0] += exp_i; + *in_out_mat.at_mut(ii) = exp_i; + } + + #[cfg(not(feature = "subgroup_ops"))] + { + *workspace.at_mut(thread_id) = my_denominator[0]; + } + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(my_denominator[0]); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *denominator = sum; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *denominator = *workspace.at(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Divide by the denominator. + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = shape.it(l, k, j, i) as usize; + let val_i = *in_out_mat.at(ii); + *in_out_mat.at_mut(ii) = val_i / *denominator; + } +} + +/// Log-softmax on columns. +#[spirv_bindgen] +#[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] +#[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(64, 1, 1))))] +pub fn log_softmax( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup)] workspace: &mut [f32; WORKGROUP_SIZE], + #[spirv(workgroup)] the_max: &mut f32, + #[spirv(workgroup)] denominator: &mut f32, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] in_out_mat: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape = *shape.at(0); + + let j = workgroup_id.x; + let k = workgroup_id.y; + let l = workgroup_id.z; + let thread_id = local_id.x as usize; + + // Compute the MAX + let data_len = shape.w; + let mut my_max = [-1.0e38f32]; + + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let val_i = *in_out_mat.at(shape.it(l, k, j, i) as usize); + my_max[0] = my_max[0].max(val_i); + } + + *workspace.at_mut(thread_id) = my_max[0]; + + #[cfg(feature = "subgroup_ops")] + let max_val = khal_std::sync::subgroup_f_max(my_max[0]); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_max(thread_id, 32, workspace); + reduce_max(thread_id, 16, workspace); + reduce_max(thread_id, 8, workspace); + reduce_max(thread_id, 4, workspace); + reduce_max(thread_id, 2, workspace); + reduce_max(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *the_max = max_val; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *the_max = *workspace.at(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Compute the denominator (sum of exponentials). + let mut my_denominator = [0.0f32]; + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = shape.it(l, k, j, i) as usize; + let val_i = *in_out_mat.at(ii); + let centered_val_i = val_i - *the_max; + let exp_i = centered_val_i.exp(); + my_denominator[0] += exp_i; + *in_out_mat.at_mut(ii) = centered_val_i; + } + + *workspace.at_mut(thread_id) = my_denominator[0]; + + #[cfg(feature = "subgroup_ops")] + let sum = khal_std::sync::subgroup_f_add(my_denominator[0]); + + #[cfg(not(feature = "subgroup_ops"))] + { + reduce_sum(thread_id, 32, workspace); + reduce_sum(thread_id, 16, workspace); + reduce_sum(thread_id, 8, workspace); + reduce_sum(thread_id, 4, workspace); + reduce_sum(thread_id, 2, workspace); + reduce_sum(thread_id, 1, workspace); + } + + if thread_id == 0 { + #[cfg(feature = "subgroup_ops")] + { + *denominator = sum; + } + #[cfg(not(feature = "subgroup_ops"))] + { + *denominator = *workspace.at(0); + } + } + + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + // Subtract log(denominator). + for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { + let ii = shape.it(l, k, j, i) as usize; + let val_i = *in_out_mat.at(ii); + *in_out_mat.at_mut(ii) = val_i - *denominator; + } +} diff --git a/vortx-shaders/src/ml/unary.rs b/vortx-shaders/src/ml/unary.rs new file mode 100644 index 0000000..dd29da5 --- /dev/null +++ b/vortx-shaders/src/ml/unary.rs @@ -0,0 +1,1671 @@ +//! Unary operations for tensors. + +use khal_std::glamx::{UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::{Shapes1, Shapes2}; +use crate::utils::limits::MAX_NUM_WORKGROUPS; +use crate::utils::trig::stable_tanh; + +const WORKGROUP_SIZE: u32 = 64; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +// // GELU constants +const GELU_COEF_A: f32 = 0.044715; +const SQRT_2_OVER_PI: f32 = 0.79788456080286535587989211986876; +const GELU_QUICK_COEF: f32 = -1.702; + +// Unary operations without arguments + +#[inline] +fn abs_op_fn(x: f32) -> f32 { + x.abs() +} + +#[inline] +fn sgn_op_fn(x: f32) -> f32 { + if x >= 0.0 { + 1.0 + } else { + -1.0 + } +} + +#[inline] +fn neg_op_fn(x: f32) -> f32 { + -x +} + +#[inline] +fn step_op_fn(x: f32) -> f32 { + if x > 0.0 { + 1.0 + } else { + 0.0 + } +} + +#[inline] +fn elu_op_fn(x: f32) -> f32 { + if x > 0.0 { + x + } else { + x.exp() - 1.0 + } +} + +#[inline] +fn gelu_op_fn(x: f32) -> f32 { + 0.5 * x * (1.0 + stable_tanh(SQRT_2_OVER_PI * x * (1.0 + GELU_COEF_A * x * x))) +} + +#[inline] +fn gelu_quick_op_fn(x: f32) -> f32 { + x * (1.0 / (1.0 + (GELU_QUICK_COEF * x).exp())) +} + +#[inline] +fn silu_op_fn(x: f32) -> f32 { + x / (1.0 + (-x).exp()) +} + +#[inline] +fn tanh_op_fn(x: f32) -> f32 { + stable_tanh(x) +} + +#[inline] +fn relu_op_fn(x: f32) -> f32 { + x.max(0.0) +} + +#[inline] +fn sigmoid_op_fn(x: f32) -> f32 { + 1.0 / (1.0 + (-x).exp()) +} + +#[inline] +fn hard_sigmoid_op_fn(x: f32) -> f32 { + 1.0f32.min(0.0f32.max((x + 3.0) / 6.0)) +} + +#[inline] +fn hard_swish_op_fn(x: f32) -> f32 { + x * 1.0f32.min(0.0f32.max((x + 3.0) / 6.0)) +} + +#[inline] +fn sqr_op_fn(x: f32) -> f32 { + x * x +} + +#[inline] +fn sqrt_op_fn(x: f32) -> f32 { + x.sqrt() +} + +#[inline] +fn sin_op_fn(x: f32) -> f32 { + x.sin() +} + +#[inline] +fn cos_op_fn(x: f32) -> f32 { + x.cos() +} + +#[inline] +fn log_op_fn(x: f32) -> f32 { + x.ln() +} + +#[inline] +fn exp_op_fn(x: f32) -> f32 { + x.exp() +} + +#[inline] +fn reciprocal_op_fn(x: f32) -> f32 { + 1.0 / x +} + +/// Erf (error function) approximation using Abramowitz and Stegun formula. +#[inline] +#[allow(clippy::excessive_precision)] +fn erf_op_fn(x: f32) -> f32 { + let a1: f32 = 0.254829592; + let a2: f32 = -0.284496736; + let a3: f32 = 1.421413741; + let a4: f32 = -1.453152027; + let a5: f32 = 1.061405429; + let p: f32 = 0.3275911; + + let sign = if x >= 0.0 { 1.0 } else { -1.0 }; + let x = x.abs(); + let t = 1.0 / (1.0 + p * x); + let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp(); + sign * y +} + +// Unary operations with arguments + +#[inline] +fn leaky_relu_op_fn(x: f32, args: Vec4) -> f32 { + x.max(0.0) + x.min(0.0) * args.x +} + +#[inline] +fn clamp_op_fn(x: f32, args: Vec4) -> f32 { + args.x.max(x).min(args.y) +} + +#[inline] +fn scale_op_fn(x: f32, args: Vec4) -> f32 { + x * args.x +} + +#[inline] +fn add_scalar_op_fn(x: f32, args: Vec4) -> f32 { + x + args.x +} + +#[inline] +fn pow_op_fn(x: f32, args: Vec4) -> f32 { + x.powf(args.x) +} + +// Macro-like helper for generating shader entry points +// Since we can't use actual macros in no_std easily, we'll define each manually + +/// Abs operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn abs_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = abs_op_fn(*src.at(isrc)); + } +} + +/// Abs operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn abs_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = abs_op_fn(*src.at(isrc)); + } +} + +/// Sign operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sgn_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = sgn_op_fn(*src.at(isrc)); + } +} + +/// Sign operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sgn_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = sgn_op_fn(*src.at(isrc)); + } +} + +/// Negation operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn neg_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = neg_op_fn(*src.at(isrc)); + } +} + +/// Negation operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn neg_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = neg_op_fn(*src.at(isrc)); + } +} + +/// Step operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn step_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = step_op_fn(*src.at(isrc)); + } +} + +/// Step operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn step_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = step_op_fn(*src.at(isrc)); + } +} + +/// ELU operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn elu_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = elu_op_fn(*src.at(isrc)); + } +} + +/// ELU operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn elu_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = elu_op_fn(*src.at(isrc)); + } +} + +/// GELU operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gelu_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = gelu_op_fn(*src.at(isrc)); + } +} + +/// GELU operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gelu_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = gelu_op_fn(*src.at(isrc)); + } +} + +/// GELU Quick operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gelu_quick_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = gelu_quick_op_fn(*src.at(isrc)); + } +} + +/// GELU Quick operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn gelu_quick_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = gelu_quick_op_fn(*src.at(isrc)); + } +} + +/// SiLU operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn silu_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = silu_op_fn(*src.at(isrc)); + } +} + +/// SiLU operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn silu_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = silu_op_fn(*src.at(isrc)); + } +} + +/// Tanh operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn tanh_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = tanh_op_fn(*src.at(isrc)); + } +} + +/// Tanh operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn tanh_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = tanh_op_fn(*src.at(isrc)); + } +} + +/// ReLU operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn relu_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = relu_op_fn(*src.at(isrc)); + } +} + +/// ReLU operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn relu_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = relu_op_fn(*src.at(isrc)); + } +} + +/// Sigmoid operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sigmoid_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = sigmoid_op_fn(*src.at(isrc)); + } +} + +/// Sigmoid operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sigmoid_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = sigmoid_op_fn(*src.at(isrc)); + } +} + +/// Hard sigmoid operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn hard_sigmoid_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = hard_sigmoid_op_fn(*src.at(isrc)); + } +} + +/// Hard sigmoid operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn hard_sigmoid_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = hard_sigmoid_op_fn(*src.at(isrc)); + } +} + +/// Square operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sqr_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = sqr_op_fn(*src.at(isrc)); + } +} + +/// Square operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sqr_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = sqr_op_fn(*src.at(isrc)); + } +} + +/// Square root operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sqrt_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = sqrt_op_fn(*src.at(isrc)); + } +} + +/// Square root operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sqrt_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = sqrt_op_fn(*src.at(isrc)); + } +} + +/// Sine operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sin_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = sin_op_fn(*src.at(isrc)); + } +} + +/// Sine operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn sin_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = sin_op_fn(*src.at(isrc)); + } +} + +/// Cosine operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn cos_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = cos_op_fn(*src.at(isrc)); + } +} + +/// Cosine operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn cos_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = cos_op_fn(*src.at(isrc)); + } +} + +/// Log operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn log_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = log_op_fn(*src.at(isrc)); + } +} + +/// Log operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn log_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load shape from storage buffer to local variable (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = log_op_fn(*src.at(isrc)); + } +} + +/// Exp operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn exp_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = exp_op_fn(*src.at(isrc)); + } +} + +/// Exp operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn exp_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = exp_op_fn(*src.at(isrc)); + } +} + +/// Reciprocal operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reciprocal_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = reciprocal_op_fn(*src.at(isrc)); + } +} + +/// Reciprocal operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn reciprocal_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = reciprocal_op_fn(*src.at(isrc)); + } +} + +// Operations with arguments + +/// Leaky ReLU operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn leaky_relu_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = leaky_relu_op_fn(*src.at(isrc), args); + } +} + +/// Leaky ReLU operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn leaky_relu_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = leaky_relu_op_fn(*src.at(isrc), args); + } +} + +/// Clamp operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn clamp_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = clamp_op_fn(*src.at(isrc), args); + } +} + +/// Clamp operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn clamp_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = clamp_op_fn(*src.at(isrc), args); + } +} + +/// Scale operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn scale_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = scale_op_fn(*src.at(isrc), args); + } +} + +/// Scale operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn scale_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = scale_op_fn(*src.at(isrc), args); + } +} + +/// Add scalar operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn add_scalar_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = add_scalar_op_fn(*src.at(isrc), args); + } +} + +/// Add scalar operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn add_scalar_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = add_scalar_op_fn(*src.at(isrc), args); + } +} + +/// Erf (error function) operation. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn erf_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = erf_op_fn(*src.at(isrc)); + } +} + +/// Erf (error function) operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn erf_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = erf_op_fn(*src.at(isrc)); + } +} + +/// Pow operation (x raised to power in args.x). +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn pow_op( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[f32], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + shape_dst: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dst: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let (shape_src, shape_dst) = (shapes.shape_a, shapes.shape_b); + #[cfg(not(feature = "push_constants"))] + let shape_dst = *shape_dst.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + let idst = shape_dst.it_vec(id) as usize; + *dst.at_mut(idst) = pow_op_fn(*src.at(isrc), args); + } +} + +/// Pow operation inplace. +#[spirv_bindgen] +#[spirv(compute(threads(64, 1, 1)))] +pub fn pow_inplace( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes1, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_src: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] args: &[Vec4], +) { + #[cfg(feature = "push_constants")] + let shape_src = shapes.shape; + #[cfg(not(feature = "push_constants"))] + let shape_src = *shape_src.at(0); + let args = *args.at(0); + + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { + let id = shape_src.decompose(thread_id); + let isrc = shape_src.it_vec(id) as usize; + *src.at_mut(isrc) = pow_op_fn(*src.at(isrc), args); + } +} diff --git a/vortx-shaders/src/ml/win_part.rs b/vortx-shaders/src/ml/win_part.rs new file mode 100644 index 0000000..f0e1bd6 --- /dev/null +++ b/vortx-shaders/src/ml/win_part.rs @@ -0,0 +1,109 @@ +//! Window partitioning. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; + +const WORKGROUP_SIZE: u32 = 128; + +/// Window partition. +/// source: [R1, C2, M1, T1] +/// result: [R1, W, W, _] +#[spirv_bindgen] +#[spirv(compute(threads(128, 1, 1)))] +pub fn win_part( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_result: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_source: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] result: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] source: &[f32], +) { + #[cfg(feature = "push_constants")] + let (shape_result, shape_source) = (shapes.shape_a, shapes.shape_b); + // Load shapes from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_result = *shape_result.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_source = *shape_source.at(0); + + if invocation_id.x >= shape_result.len() { + return; + } + + let id = shape_result.decompose(invocation_id.x); + let i = shape_result.it_vec(id) as usize; + + let w = shape_result.c; + let pad_x = (w - shape_source.h % w) % w; + // NOTE: notation nep0 === npx + let nep0 = (pad_x + shape_source.h) / w; + + // NOTE: id[3] spans [0..nep0*nep1[ by definition of the result tensor. + let py = id.w / nep0; + let px = id.w - py * nep0; + let i02 = py * w + id.z; + let i01 = px * w + id.x; + let i00 = id.y; + + if py * w + id.z >= shape_source.c || px * w + id.x >= shape_source.h { + *result.at_mut(i) = 0.0; + } else { + let j = shape_source.it(0, i02, i01, i00) as usize; + *result.at_mut(i) = *source.at(j); + } +} + +/// Window unpartition. +/// source: [R1, C2, M1, T1] +/// result: [R1, W, W, _] +#[spirv_bindgen] +#[spirv(compute(threads(128, 1, 1)))] +pub fn win_unpart( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[cfg(feature = "push_constants")] + #[spirv(push_constant)] + shapes: &Shapes2, + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + shape_result: &[Shape], + #[cfg(not(feature = "push_constants"))] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + shape_source: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] w: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] result: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] source: &[f32], +) { + #[cfg(feature = "push_constants")] + let (shape_result, shape_source) = (shapes.shape_a, shapes.shape_b); + // Load from storage buffer to local variables (enables LICM) + #[cfg(not(feature = "push_constants"))] + let shape_result = *shape_result.at(0); + #[cfg(not(feature = "push_constants"))] + let shape_source = *shape_source.at(0); + let w = *w.at(0); + + if invocation_id.x >= shape_result.len() { + return; + } + + let id = shape_result.decompose(invocation_id.x); + let j = shape_result.it_vec(id) as usize; + + let px = (w - shape_result.h % w) % w; + let npx = (px + shape_result.h) / w; + let ip2 = id.z / w; + let ip1 = id.x / w; + let i = shape_source.it(ip2 * npx + ip1, id.z % w, id.x % w, id.y) as usize; + + *result.at_mut(j) = *source.at(i); +} diff --git a/vortx-shaders/src/utils/half.rs b/vortx-shaders/src/utils/half.rs new file mode 100644 index 0000000..006a3ff --- /dev/null +++ b/vortx-shaders/src/utils/half.rs @@ -0,0 +1,32 @@ +//! Half-precision floating point utilities. + +use khal_std::glamx::{IVec4, UVec4, Vec2}; + +/// Unpack a u32 containing two f16 values into a Vec2 of f32. +/// The low 16 bits are the first half, high 16 bits are the second half. +#[inline] +pub fn unpack_half2x16(v: u32) -> Vec2 { + khal_std::float::unpack_half2x16(v) +} + +/// Unpack a u32 containing 4 signed 8-bit integers into an IVec4. +/// Extracts bytes and sign-extends them to i32. +#[inline] +pub fn unpack_int4x8(v: u32) -> IVec4 { + // Extract each byte and sign-extend from i8 to i32 + let b0 = ((v & 0xFF) as i32) << 24 >> 24; + let b1 = (((v >> 8) & 0xFF) as i32) << 24 >> 24; + let b2 = (((v >> 16) & 0xFF) as i32) << 24 >> 24; + let b3 = (((v >> 24) & 0xFF) as i32) << 24 >> 24; + IVec4::new(b0, b1, b2, b3) +} + +/// Unpack a u32 containing 4 unsigned 8-bit integers into a UVec4. +#[inline] +pub fn unpack_uint4x8(v: u32) -> UVec4 { + let b0 = v & 0xFF; + let b1 = (v >> 8) & 0xFF; + let b2 = (v >> 16) & 0xFF; + let b3 = (v >> 24) & 0xFF; + UVec4::new(b0, b1, b2, b3) +} diff --git a/vortx-shaders/src/utils/mod.rs b/vortx-shaders/src/utils/mod.rs index b8d83e8..a58d81b 100644 --- a/vortx-shaders/src/utils/mod.rs +++ b/vortx-shaders/src/utils/mod.rs @@ -4,3 +4,4 @@ pub mod iterators; pub mod limits; pub mod mat; pub mod trig; +pub mod half; \ No newline at end of file From 514eb35dcefad73da8075737876bb3a43bc4315e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 02/13] chore: cargo fmt --- src/lib.rs | 4 ++-- src/ml/batched_multiquery_attention.rs | 8 ++++---- src/ml/concat.rs | 4 ++-- src/ml/conv2d_nchw.rs | 4 ++-- src/ml/conv_transpose_2d.rs | 4 ++-- src/ml/gather.rs | 4 ++-- src/ml/gemv_quant.rs | 12 ++++++------ src/ml/get_rel_pos.rs | 4 ++-- src/ml/im2col.rs | 5 ++--- src/ml/layernorm.rs | 10 +++++----- src/ml/mod.rs | 18 ++++++++---------- src/ml/pool2d.rs | 4 ++-- src/ml/quantization.rs | 6 +----- src/ml/quantized_matrix.rs | 10 +++++----- src/ml/reduce_axis.rs | 8 ++++---- src/ml/rms_norm.rs | 10 +++++----- src/ml/rope.rs | 10 +++++----- src/ml/select.rs | 8 ++++---- src/ml/silu.rs | 10 +++++----- src/ml/softmax.rs | 10 +++++----- src/ml/unary.rs | 14 +++++++------- src/ml/win_part.rs | 4 ++-- vortx-shaders/src/lib.rs | 4 ++-- vortx-shaders/src/ml/concat.rs | 6 +++--- vortx-shaders/src/ml/conv2d.rs | 2 +- vortx-shaders/src/ml/conv_transpose_2d.rs | 6 +++--- vortx-shaders/src/ml/gather.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q4_0x2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q4_1x2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q4_k.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q5_0x2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q5_1x2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q5_k.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q6_kx2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q8_0x2.rs | 6 +++--- vortx-shaders/src/ml/gemv_quant_q8_k.rs | 6 +++--- vortx-shaders/src/ml/get_rel_pos.rs | 6 +++--- vortx-shaders/src/ml/layernorm.rs | 6 +++--- vortx-shaders/src/ml/pool2d.rs | 2 +- vortx-shaders/src/ml/reduce_axis.rs | 6 +++--- vortx-shaders/src/ml/rms_norm.rs | 6 +++--- vortx-shaders/src/ml/rope.rs | 6 +++--- vortx-shaders/src/ml/select.rs | 6 +++--- vortx-shaders/src/ml/silu.rs | 6 +++--- vortx-shaders/src/ml/softmax.rs | 6 +++--- vortx-shaders/src/ml/unary.rs | 10 +++++----- vortx-shaders/src/ml/win_part.rs | 6 +++--- vortx-shaders/src/utils/mod.rs | 2 +- 48 files changed, 155 insertions(+), 162 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 9957bc2..b92f6da 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,7 +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; -#[cfg(feature = "ml")] -pub mod ml; \ No newline at end of file diff --git a/src/ml/batched_multiquery_attention.rs b/src/ml/batched_multiquery_attention.rs index 053ee75..d02ba14 100644 --- a/src/ml/batched_multiquery_attention.rs +++ b/src/ml/batched_multiquery_attention.rs @@ -1,9 +1,9 @@ -use vortx_shaders::ml::AttentionParams; -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use crate::ml::SoftMax; +use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; use nalgebra::{DMatrix, DVector}; -use crate::tensor::{AsTensorMut, AsTensorRef}; -use crate::ml::SoftMax; +use vortx_shaders::ml::AttentionParams; #[derive(Shader)] /// Fused attention shader - combines Q*K^T, scale, mask, softmax, and *V into one kernel. diff --git a/src/ml/concat.rs b/src/ml/concat.rs index da4049e..45114d5 100644 --- a/src/ml/concat.rs +++ b/src/ml/concat.rs @@ -1,9 +1,9 @@ //! Concat operation: concatenates tensors along a given axis. -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::{BufferUsages, Shader}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; /// Shader for the Concat operation. #[derive(Shader)] diff --git a/src/ml/conv2d_nchw.rs b/src/ml/conv2d_nchw.rs index b2d7767..4e73219 100644 --- a/src/ml/conv2d_nchw.rs +++ b/src/ml/conv2d_nchw.rs @@ -2,9 +2,9 @@ //! //! This implementation works with ONNX tensor format directly. -use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; -use khal::Shader; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; #[derive(Shader)] pub struct Conv2dNchw { diff --git a/src/ml/conv_transpose_2d.rs b/src/ml/conv_transpose_2d.rs index e6a91e1..aef6bfe 100644 --- a/src/ml/conv_transpose_2d.rs +++ b/src/ml/conv_transpose_2d.rs @@ -1,7 +1,7 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; #[derive(Shader)] pub struct ConvTranspose2d { diff --git a/src/ml/gather.rs b/src/ml/gather.rs index eb04562..b05f92c 100644 --- a/src/ml/gather.rs +++ b/src/ml/gather.rs @@ -1,9 +1,9 @@ //! Gather operation: gathers elements from a tensor based on indices along an axis. -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::{BufferUsages, Shader}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; /// Shader for the Gather operation. #[derive(Shader)] diff --git a/src/ml/gemv_quant.rs b/src/ml/gemv_quant.rs index 401ad08..647bc49 100644 --- a/src/ml/gemv_quant.rs +++ b/src/ml/gemv_quant.rs @@ -1,10 +1,10 @@ +use crate::Gemm; use crate::ml::quantization::{BlockQ4K, BlockQ5K, BlockQ8K}; use crate::ml::quantized_matrix::GpuQuantTensor; -use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; -use crate::Gemm; +use khal::Shader; +use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; #[cfg(feature = "rand")] use rand::distr::{Distribution, StandardUniform}; @@ -351,17 +351,17 @@ mod test { use super::*; use crate::ml::quantization::*; use crate::ml::quantized_matrix::GpuQuantTensor; - use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; - use khal::BufferUsages; use crate::shapes::TensorLayoutBuffers; use crate::tensor::Tensor; + use khal::BufferUsages; + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; use wgpu::{Features, Limits}; /// Generate a random valid f16 scale (avoids NaN/Inf). /// Includes subnormal values to cover the full f16 range. fn rand_f16_scale() -> u16 { let val: f32 = rand::random::() * 2.0 - 1.0; // [-1, 1] - // Scale down ~25% of values into the subnormal f16 range (< 2^-14). + // Scale down ~25% of values into the subnormal f16 range (< 2^-14). let val = if rand::random::() < 64 { val * 1e-5 } else { diff --git a/src/ml/get_rel_pos.rs b/src/ml/get_rel_pos.rs index a264d8c..c9061ac 100644 --- a/src/ml/get_rel_pos.rs +++ b/src/ml/get_rel_pos.rs @@ -1,7 +1,7 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; #[derive(Shader)] pub struct GetRelPos { diff --git a/src/ml/im2col.rs b/src/ml/im2col.rs index 613ffcd..d0d82df 100644 --- a/src/ml/im2col.rs +++ b/src/ml/im2col.rs @@ -1,7 +1,6 @@ -use khal::backend::{Backend, DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; - +use khal::Shader; +use khal::backend::{Backend, DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; pub type Im2ColConfig = vortx_shaders::ml::im2col::Im2ColParams; diff --git a/src/ml/layernorm.rs b/src/ml/layernorm.rs index ff8d666..9cb5752 100644 --- a/src/ml/layernorm.rs +++ b/src/ml/layernorm.rs @@ -1,8 +1,8 @@ -use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::DVector; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; +use nalgebra::DVector; #[derive(Shader)] /// Shader implementing the layer normalization kernel. @@ -144,12 +144,12 @@ impl LayerNorm { #[cfg(test)] mod test { use crate::ml::LayerNorm; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/mod.rs b/src/ml/mod.rs index 63bf926..ee40714 100644 --- a/src/ml/mod.rs +++ b/src/ml/mod.rs @@ -10,6 +10,8 @@ mod get_rel_pos; mod im2col; mod layernorm; mod pool2d; +pub mod quantization; +mod quantized_matrix; mod reduce_axis; mod rms_norm; mod rope; @@ -18,25 +20,21 @@ mod silu; mod softmax; mod unary; mod win_part; -mod quantized_matrix; -pub mod quantization; -pub use quantized_matrix::*; -pub use batched_multiquery_attention::{ - FusedAttention, -}; +pub use batched_multiquery_attention::FusedAttention; pub use concat::Concat; -pub use conv2d_nchw::{conv_output_size, Conv2dNchw}; pub use conv_transpose_2d::ConvTranspose2d; +pub use conv2d_nchw::{Conv2dNchw, conv_output_size}; pub use gather::Gather; pub use gemv_quant::{ - GemvQuant, GpuBlockQ4K, GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ5K, GpuBlockQ5_0x2, - GpuBlockQ5_1x2, GpuBlockQ6Kx2, GpuBlockQ8K, GpuBlockQ8_0x2, QuantizedValue, + GemvQuant, GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ4K, GpuBlockQ5_0x2, GpuBlockQ5_1x2, + GpuBlockQ5K, GpuBlockQ6Kx2, GpuBlockQ8_0x2, GpuBlockQ8K, QuantizedValue, }; pub use get_rel_pos::GetRelPos; pub use im2col::{Im2Col, Im2ColConfig}; pub use layernorm::LayerNorm; -pub use pool2d::{pool_output_size, GlobalPool2dConfig, Pool2d, Pool2dConfig}; +pub use pool2d::{GlobalPool2dConfig, Pool2d, Pool2dConfig, pool_output_size}; +pub use quantized_matrix::*; pub use reduce_axis::{ReduceAxis, ReduceOp}; pub use rms_norm::{RmsNorm, RmsNormConfig}; pub use rope::{RoPE, RoPEConfig, RoPEVariant}; diff --git a/src/ml/pool2d.rs b/src/ml/pool2d.rs index 65ee3c8..2c6ce42 100644 --- a/src/ml/pool2d.rs +++ b/src/ml/pool2d.rs @@ -1,8 +1,8 @@ //! 2D Pooling operations (MaxPool2d, AvgPool2d, GlobalAvgPool2d, GlobalMaxPool2d). -use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; -use khal::Shader; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuBuffer, GpuPass}; /// Pool2d configuration parameters. #[derive(Copy, Clone, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable, Debug, Default)] diff --git a/src/ml/quantization.rs b/src/ml/quantization.rs index 80d4ca9..dd089bd 100644 --- a/src/ml/quantization.rs +++ b/src/ml/quantization.rs @@ -421,11 +421,7 @@ pub fn decode_f16(half: u16) -> f32 { } else { f32::NAN }; - if half & 0x8000 != 0 { - -val - } else { - val - } + if half & 0x8000 != 0 { -val } else { val } } pub fn decode_bf16(half: u16) -> f32 { diff --git a/src/ml/quantized_matrix.rs b/src/ml/quantized_matrix.rs index aa845cd..6e73017 100644 --- a/src/ml/quantized_matrix.rs +++ b/src/ml/quantized_matrix.rs @@ -1,12 +1,12 @@ use crate::ml::{ - GpuBlockQ4K, GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ5K, GpuBlockQ5_0x2, GpuBlockQ5_1x2, - GpuBlockQ6Kx2, GpuBlockQ8K, GpuBlockQ8_0x2, + GpuBlockQ4_0x2, GpuBlockQ4_1x2, GpuBlockQ4K, GpuBlockQ5_0x2, GpuBlockQ5_1x2, GpuBlockQ5K, + GpuBlockQ6Kx2, GpuBlockQ8_0x2, GpuBlockQ8K, }; -use khal::backend::{GpuDispatch, ShaderBinding}; -use khal::shader::ShaderArgsError; -use khal::ShaderArgs; use crate::shapes::TensorLayout; use crate::tensor::Tensor; +use khal::ShaderArgs; +use khal::backend::{GpuDispatch, ShaderBinding}; +use khal::shader::ShaderArgsError; pub enum GpuQuantTensor { F32(Tensor), diff --git a/src/ml/reduce_axis.rs b/src/ml/reduce_axis.rs index 069ef36..1a4eb60 100644 --- a/src/ml/reduce_axis.rs +++ b/src/ml/reduce_axis.rs @@ -1,9 +1,9 @@ //! Axis-based reduction operations (ReduceSum, ReduceMean, etc.) -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::{BufferUsages, Shader}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, TensorBuilder}; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; /// Type of reduction operation. #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -103,10 +103,10 @@ impl ReduceAxis { #[cfg(test)] mod test { use super::*; - use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; - use khal::{BufferUsages, Shader}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::Tensor; + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; + use khal::{BufferUsages, Shader}; use wgpu::{Features, Limits}; async fn test_reduce_sum_axis_generic(backend: &GpuBackend) { diff --git a/src/ml/rms_norm.rs b/src/ml/rms_norm.rs index 77a9044..f4104a3 100644 --- a/src/ml/rms_norm.rs +++ b/src/ml/rms_norm.rs @@ -1,8 +1,8 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::{DVector, Dyn, Storage, Vector}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use nalgebra::{DVector, Dyn, Storage, Vector}; #[derive(Shader)] /// Shader implementing the RMS norm kernel. @@ -82,12 +82,12 @@ impl RmsNorm { #[cfg(test)] mod test { use crate::ml::{RmsNorm, RmsNormConfig}; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/rope.rs b/src/ml/rope.rs index a4d5573..ab8312b 100644 --- a/src/ml/rope.rs +++ b/src/ml/rope.rs @@ -1,8 +1,8 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::{vector, DVector, DVectorViewMut, Rotation2}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, Tensor}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use nalgebra::{DVector, DVectorViewMut, Rotation2, vector}; #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum RoPEVariant { @@ -139,12 +139,12 @@ impl RoPE { mod test { use super::RoPEConfig; use crate::ml::{RoPE, RoPEVariant}; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/select.rs b/src/ml/select.rs index a687f50..620da34 100644 --- a/src/ml/select.rs +++ b/src/ml/select.rs @@ -1,7 +1,7 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; #[derive(Shader)] pub struct Select { @@ -66,10 +66,10 @@ impl Select { #[cfg(test)] mod test { - use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; - use khal::{BufferUsages, Shader}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::Tensor; + use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; + use khal::{BufferUsages, Shader}; use wgpu::{Features, Limits}; /// Select rows from a matrix by index: dest[i] = src[idx[i], :]. diff --git a/src/ml/silu.rs b/src/ml/silu.rs index 34bcfbf..00aed8e 100644 --- a/src/ml/silu.rs +++ b/src/ml/silu.rs @@ -1,8 +1,8 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::DVector; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use nalgebra::DVector; #[derive(Shader)] /// Shader implementing the Silu activation function. @@ -65,12 +65,12 @@ impl Silu { #[cfg(test)] mod test { + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/softmax.rs b/src/ml/softmax.rs index b3e7ea2..e28649f 100644 --- a/src/ml/softmax.rs +++ b/src/ml/softmax.rs @@ -1,8 +1,8 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::{Dyn, StorageMut, Vector}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::AsTensorMut; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use nalgebra::{Dyn, StorageMut, Vector}; /* layout (push_constant) uniform parameter @@ -146,12 +146,12 @@ impl SoftMax { #[cfg(test)] mod test { use crate::ml::SoftMax; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/unary.rs b/src/ml/unary.rs index 1e3ac77..2e0231e 100644 --- a/src/ml/unary.rs +++ b/src/ml/unary.rs @@ -1,9 +1,9 @@ -use khal_std::glamx::Vec4; -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; -use nalgebra::{Dyn, StorageMut, Vector}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal_std::glamx::Vec4; +use nalgebra::{Dyn, StorageMut, Vector}; #[derive(Copy, Clone, PartialEq, Eq, Debug)] #[non_exhaustive] @@ -930,13 +930,13 @@ impl Unary { #[cfg(feature = "rand")] mod test { use crate::ml::UnaryOp; - use khal_std::glamx::Vec4; + use crate::shapes::TensorLayoutBuffers; + use crate::tensor::Tensor; use khal::backend::WebGpu; use khal::backend::{Backend, Encoder, GpuBackend}; use khal::{BufferUsages, Shader}; + use khal_std::glamx::Vec4; use nalgebra::DVector; - use crate::shapes::TensorLayoutBuffers; - use crate::tensor::Tensor; use wgpu::{Features, Limits}; #[futures_test::test] diff --git a/src/ml/win_part.rs b/src/ml/win_part.rs index 7503308..337a262 100644 --- a/src/ml/win_part.rs +++ b/src/ml/win_part.rs @@ -1,7 +1,7 @@ -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -use khal::Shader; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef, Tensor}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; #[derive(Shader)] pub struct WinPart { diff --git a/vortx-shaders/src/lib.rs b/vortx-shaders/src/lib.rs index f7128b1..a097b46 100644 --- a/vortx-shaders/src/lib.rs +++ b/vortx-shaders/src/lib.rs @@ -12,6 +12,6 @@ extern crate std; pub mod linalg; -pub mod utils; #[cfg(feature = "ml")] -pub mod ml; \ No newline at end of file +pub mod ml; +pub mod utils; diff --git a/vortx-shaders/src/ml/concat.rs b/vortx-shaders/src/ml/concat.rs index 4ab9c58..49e37b8 100644 --- a/vortx-shaders/src/ml/concat.rs +++ b/vortx-shaders/src/ml/concat.rs @@ -1,12 +1,12 @@ //! Concat operation: concatenates tensors along a given axis. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::Shapes2; use crate::utils::limits::MAX_NUM_WORKGROUPS; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/conv2d.rs b/vortx-shaders/src/ml/conv2d.rs index 6d5081e..29b0559 100644 --- a/vortx-shaders/src/ml/conv2d.rs +++ b/vortx-shaders/src/ml/conv2d.rs @@ -22,10 +22,10 @@ //! \[14\] batch_size //! \[15\] groups (must be 1 for now) +use crate::utils::limits::MAX_NUM_WORKGROUPS; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::utils::limits::MAX_NUM_WORKGROUPS; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/conv_transpose_2d.rs b/vortx-shaders/src/ml/conv_transpose_2d.rs index 81d0e16..7ce86ec 100644 --- a/vortx-shaders/src/ml/conv_transpose_2d.rs +++ b/vortx-shaders/src/ml/conv_transpose_2d.rs @@ -1,11 +1,11 @@ //! Transposed 2D convolution. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::{Shapes1, Shapes2, Shapes3}; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/gather.rs b/vortx-shaders/src/ml/gather.rs index 4d2fb03..6adaf0a 100644 --- a/vortx-shaders/src/ml/gather.rs +++ b/vortx-shaders/src/ml/gather.rs @@ -1,12 +1,12 @@ //! Gather operation: gathers elements from source tensor based on indices along a given axis. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::Shapes2; use crate::utils::limits::MAX_NUM_WORKGROUPS; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs index 06e03ca..d21b800 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs @@ -2,13 +2,13 @@ //! //! BlockQ4_0x2 contains two BlockQ4_0 blocks (f16 scale + 16 x 4-bit quants each). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{Mat4, UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: usize = 32; const COLS_STEP: u32 = 4; diff --git a/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs index 78fb824..0de5339 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs @@ -2,13 +2,13 @@ //! //! BlockQ4_1x2 contains two BlockQ4_1 blocks (f16 scale + f16 min + 16 x 4-bit quants each). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{UVec3, Vec2, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/gemv_quant_q4_k.rs b/vortx-shaders/src/ml/gemv_quant_q4_k.rs index cc0d5c2..731b797 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_k.rs @@ -2,13 +2,13 @@ //! //! BlockQ4K: super-block scale, super-block min, 12 bytes scales/mins, 128 bytes quants (256 4-bit values). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{UVec2, UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: usize = 32; // BlockQ4K size in u32s: 1 (d_dmin) + 3 (scales) + 32 (qs) = 36 diff --git a/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs index 5e06dca..695f71e 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs @@ -2,13 +2,13 @@ //! //! BlockQ5_0x2 contains two BlockQ5_0 blocks (f16 scale + u32 high bits + 16 x 4-bit quants each). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs index bab4ae2..b15f8f6 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs @@ -2,13 +2,13 @@ //! //! BlockQ5_1x2 contains two BlockQ5_1 blocks (f16 scale + f16 min + u32 high bits + 16 x 4-bit quants each). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{UVec3, Vec2, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/gemv_quant_q5_k.rs b/vortx-shaders/src/ml/gemv_quant_q5_k.rs index 1de769b..a736077 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_k.rs @@ -2,13 +2,13 @@ //! //! BlockQ5K: super-block scale, super-block min, 12 bytes scales/mins, 32 bytes high bits, 128 bytes quants. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_half2x16; use khal_std::glamx::{UVec2, UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: usize = 32; // BlockQ5K size in u32s: 1 (d_dmin) + 3 (scales) + 8 (qh) + 32 (qs) = 44 diff --git a/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs index cec6c19..73f204c 100644 --- a/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs @@ -3,13 +3,13 @@ //! BlockQ6Kx2 contains two BlockQ6K blocks packed together (105 u32s total). //! Each BlockQ6K: f16 scale, 64 bytes ql (low 4 bits), 32 bytes qh (high 2 bits), 16 bytes scales. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::{unpack_half2x16, unpack_int4x8, unpack_uint4x8}; use khal_std::glamx::{IVec4, Mat4, UVec3, UVec4, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: usize = 32; // BlockQ6Kx2 size in u32s: 105 diff --git a/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs index ddbce38..13383b3 100644 --- a/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs @@ -2,13 +2,13 @@ //! //! BlockQ8_0x2 contains two BlockQ8_0 blocks (f16 scale + 32 x 8-bit signed quants each). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::{unpack_half2x16, unpack_int4x8}; use khal_std::glamx::{UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: usize = 32; const BLOCK_Q8_0X2_SIZE: u32 = 17; // 17 u32s diff --git a/vortx-shaders/src/ml/gemv_quant_q8_k.rs b/vortx-shaders/src/ml/gemv_quant_q8_k.rs index 777babb..4a89682 100644 --- a/vortx-shaders/src/ml/gemv_quant_q8_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q8_k.rs @@ -2,13 +2,13 @@ //! //! BlockQ8K: f32 delta, 256 x 8-bit signed quants, 16 x 16-bit bsums. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::half::unpack_int4x8; use khal_std::glamx::{UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; const WORKGROUP_SIZE: u32 = 32; diff --git a/vortx-shaders/src/ml/get_rel_pos.rs b/vortx-shaders/src/ml/get_rel_pos.rs index eab0112..2d75ba5 100644 --- a/vortx-shaders/src/ml/get_rel_pos.rs +++ b/vortx-shaders/src/ml/get_rel_pos.rs @@ -1,11 +1,11 @@ //! Relative position computation. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::{Shapes1, Shapes2}; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 128; diff --git a/vortx-shaders/src/ml/layernorm.rs b/vortx-shaders/src/ml/layernorm.rs index f3c8628..25517e6 100644 --- a/vortx-shaders/src/ml/layernorm.rs +++ b/vortx-shaders/src/ml/layernorm.rs @@ -1,14 +1,14 @@ //! Layer normalization kernels. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; use crate::utils::iterators::StepRng; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes2; #[cfg(feature = "subgroup_ops")] const WORKGROUP_SIZE: usize = 32; diff --git a/vortx-shaders/src/ml/pool2d.rs b/vortx-shaders/src/ml/pool2d.rs index 5868b76..4588ffa 100644 --- a/vortx-shaders/src/ml/pool2d.rs +++ b/vortx-shaders/src/ml/pool2d.rs @@ -17,10 +17,10 @@ //! \[10\] channels //! \[11\] batch_size +use crate::utils::limits::MAX_NUM_WORKGROUPS; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use crate::utils::limits::MAX_NUM_WORKGROUPS; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/reduce_axis.rs b/vortx-shaders/src/ml/reduce_axis.rs index d167279..bd93b4a 100644 --- a/vortx-shaders/src/ml/reduce_axis.rs +++ b/vortx-shaders/src/ml/reduce_axis.rs @@ -1,12 +1,12 @@ //! Axis-based reduction operations (ReduceSum, ReduceMean, etc.) -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::Shapes2; use crate::utils::limits::MAX_NUM_WORKGROUPS; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/rms_norm.rs b/vortx-shaders/src/ml/rms_norm.rs index f70f2c5..beeebfb 100644 --- a/vortx-shaders/src/ml/rms_norm.rs +++ b/vortx-shaders/src/ml/rms_norm.rs @@ -1,13 +1,13 @@ //! RMS normalization kernel. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes3; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes3; #[cfg(feature = "subgroup_ops")] const WORKGROUP_SIZE: usize = 32; diff --git a/vortx-shaders/src/ml/rope.rs b/vortx-shaders/src/ml/rope.rs index 111407e..08997c0 100644 --- a/vortx-shaders/src/ml/rope.rs +++ b/vortx-shaders/src/ml/rope.rs @@ -1,13 +1,13 @@ //! Rotary Positional Encoding (RoPE). +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes2; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/select.rs b/vortx-shaders/src/ml/select.rs index 2a1423f..5d88cea 100644 --- a/vortx-shaders/src/ml/select.rs +++ b/vortx-shaders/src/ml/select.rs @@ -1,11 +1,11 @@ //! Select operation: selects elements from a source tensor based on indices. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::Shapes2; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; /// Select elements from source based on indices and write to destination. /// diff --git a/vortx-shaders/src/ml/silu.rs b/vortx-shaders/src/ml/silu.rs index ef07e23..395a54e 100644 --- a/vortx-shaders/src/ml/silu.rs +++ b/vortx-shaders/src/ml/silu.rs @@ -1,13 +1,13 @@ //! SiLU (Swish) activation function. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes2; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes2; const WORKGROUP_SIZE: u32 = 64; diff --git a/vortx-shaders/src/ml/softmax.rs b/vortx-shaders/src/ml/softmax.rs index 1adcee5..dc09b45 100644 --- a/vortx-shaders/src/ml/softmax.rs +++ b/vortx-shaders/src/ml/softmax.rs @@ -1,14 +1,14 @@ //! Softmax and log-softmax kernels. +use crate::linalg::Shape; +#[cfg(feature = "push_constants")] +use crate::linalg::Shapes1; use crate::utils::iterators::StepRng; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -use crate::linalg::Shape; -#[cfg(feature = "push_constants")] -use crate::linalg::Shapes1; #[cfg(feature = "subgroup_ops")] const WORKGROUP_SIZE: usize = 32; diff --git a/vortx-shaders/src/ml/unary.rs b/vortx-shaders/src/ml/unary.rs index dd29da5..4f091ef 100644 --- a/vortx-shaders/src/ml/unary.rs +++ b/vortx-shaders/src/ml/unary.rs @@ -1,15 +1,15 @@ //! Unary operations for tensors. -use khal_std::glamx::{UVec3, Vec4}; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; -#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] -use khal_std::num_traits::Float; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::{Shapes1, Shapes2}; use crate::utils::limits::MAX_NUM_WORKGROUPS; use crate::utils::trig::stable_tanh; +use khal_std::glamx::{UVec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; diff --git a/vortx-shaders/src/ml/win_part.rs b/vortx-shaders/src/ml/win_part.rs index f0e1bd6..b3541ab 100644 --- a/vortx-shaders/src/ml/win_part.rs +++ b/vortx-shaders/src/ml/win_part.rs @@ -1,11 +1,11 @@ //! Window partitioning. -use khal_std::glamx::UVec3; -use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use crate::linalg::Shape; #[cfg(feature = "push_constants")] use crate::linalg::Shapes2; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 128; diff --git a/vortx-shaders/src/utils/mod.rs b/vortx-shaders/src/utils/mod.rs index a58d81b..03ea89b 100644 --- a/vortx-shaders/src/utils/mod.rs +++ b/vortx-shaders/src/utils/mod.rs @@ -1,7 +1,7 @@ //! Utility modules for shaders. +pub mod half; pub mod iterators; pub mod limits; pub mod mat; pub mod trig; -pub mod half; \ No newline at end of file From 115504c0325eee85a227fa57b80fe1d1d25cb301 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 03/13] feat(ml): add tanh activation + Adam optimizer GPU ops Replaces #5 Co-Authored-By: Haixuan Xavier Tao --- src/linalg/activation.rs | 72 ++++++++++++++++++++++++++ src/linalg/mod.rs | 4 ++ src/linalg/optim.rs | 61 ++++++++++++++++++++++ vortx-shaders/src/linalg/activation.rs | 54 +++++++++++++++++++ vortx-shaders/src/linalg/mod.rs | 6 +++ vortx-shaders/src/linalg/optim.rs | 63 ++++++++++++++++++++++ 6 files changed, 260 insertions(+) create mode 100644 src/linalg/activation.rs create mode 100644 src/linalg/optim.rs create mode 100644 vortx-shaders/src/linalg/activation.rs create mode 100644 vortx-shaders/src/linalg/optim.rs diff --git a/src/linalg/activation.rs b/src/linalg/activation.rs new file mode 100644 index 0000000..91867db --- /dev/null +++ b/src/linalg/activation.rs @@ -0,0 +1,72 @@ +//! Element-wise activation functions (host dispatch). +//! +//! Added for zealot's MLP policy — vortx upstream has no activations. + +use crate::shaders::linalg::{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, +} + +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(), + ) + } +} 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/src/linalg/activation.rs b/vortx-shaders/src/linalg/activation.rs new file mode 100644 index 0000000..04c880d --- /dev/null +++ b/vortx-shaders/src/linalg/activation.rs @@ -0,0 +1,54 @@ +//! 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; +use glamx::UVec3; +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; + } +} diff --git a/vortx-shaders/src/linalg/mod.rs b/vortx-shaders/src/linalg/mod.rs index 0ecdaff..1a2d8c5 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(target_arch_is_gpu))] +pub use activation::{GpuTanh, GpuTanhBackward}; +#[cfg(not(target_arch_is_gpu))] pub use contiguous::{Contiguous, ContiguousWithOffset}; #[cfg(not(target_arch_is_gpu))] pub use gemm::{GemmNaive, GemmTiled}; #[cfg(not(target_arch_is_gpu))] pub use op_assign::{GpuAdd, GpuCopy, GpuCopyWithOffsets, GpuDiv, GpuMul, GpuSub}; #[cfg(not(target_arch_is_gpu))] +pub use optim::GpuAdam; +#[cfg(not(target_arch_is_gpu))] pub use reduce::*; #[cfg(not(target_arch_is_gpu))] 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); + } +} From fd94b788b8e4413a8d35ae8fd05a009702923cd7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 04/13] refactor(ml): move the tanh-backward + Adam kernels into vortx::ml and drop the duplicated tanh forward Completes #5 --- src/linalg/mod.rs | 4 -- src/{linalg => ml}/activation.rs | 35 +++------------ src/ml/mod.rs | 4 ++ src/{linalg => ml}/optim.rs | 6 +-- vortx-shaders/src/linalg/activation.rs | 54 ----------------------- vortx-shaders/src/linalg/mod.rs | 6 --- vortx-shaders/src/ml/activation.rs | 35 +++++++++++++++ vortx-shaders/src/ml/mod.rs | 4 ++ vortx-shaders/src/{linalg => ml}/optim.rs | 18 ++++---- 9 files changed, 61 insertions(+), 105 deletions(-) rename src/{linalg => ml}/activation.rs (59%) rename src/{linalg => ml}/optim.rs (92%) delete mode 100644 vortx-shaders/src/linalg/activation.rs create mode 100644 vortx-shaders/src/ml/activation.rs rename vortx-shaders/src/{linalg => ml}/optim.rs (86%) diff --git a/src/linalg/mod.rs b/src/linalg/mod.rs index c13ed33..7a65987 100644 --- a/src/linalg/mod.rs +++ b/src/linalg/mod.rs @@ -1,18 +1,14 @@ //! 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/activation.rs b/src/ml/activation.rs similarity index 59% rename from src/linalg/activation.rs rename to src/ml/activation.rs index 91867db..00a93ed 100644 --- a/src/linalg/activation.rs +++ b/src/ml/activation.rs @@ -1,43 +1,22 @@ -//! Element-wise activation functions (host dispatch). +//! Backward passes of the element-wise activations (host dispatch). //! -//! Added for zealot's MLP policy — vortx upstream has no activations. +//! The forward directions are provided by [`crate::ml::Unary`] +//! (`UnaryOp::Tanh`, `UnaryOp::Elu`); only the gradients live here. -use crate::shaders::linalg::{GpuTanh, GpuTanhBackward}; +use crate::shaders::ml::GpuTanhBackward; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; -/// Element-wise activation kernels. +/// Element-wise activation gradient kernels. #[derive(Shader)] -pub struct Activation { - /// In-place tanh. - pub tanh: GpuTanh, +pub struct ActivationBackward { /// In-place tanh backward (`g *= 1 - y^2`). pub tanh_backward: GpuTanhBackward, } -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) - } - +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( diff --git a/src/ml/mod.rs b/src/ml/mod.rs index ee40714..f92da28 100644 --- a/src/ml/mod.rs +++ b/src/ml/mod.rs @@ -1,5 +1,6 @@ //! Primitives for building LLM inferences. +mod activation; mod batched_multiquery_attention; mod concat; mod conv2d_nchw; @@ -9,6 +10,7 @@ mod gemv_quant; mod get_rel_pos; mod im2col; mod layernorm; +mod optim; mod pool2d; pub mod quantization; mod quantized_matrix; @@ -21,6 +23,7 @@ mod softmax; mod unary; mod win_part; +pub use activation::ActivationBackward; pub use batched_multiquery_attention::FusedAttention; pub use concat::Concat; pub use conv_transpose_2d::ConvTranspose2d; @@ -33,6 +36,7 @@ pub use gemv_quant::{ pub use get_rel_pos::GetRelPos; pub use im2col::{Im2Col, Im2ColConfig}; pub use layernorm::LayerNorm; +pub use optim::{Adam, AdamParams}; pub use pool2d::{GlobalPool2dConfig, Pool2d, Pool2dConfig, pool_output_size}; pub use quantized_matrix::*; pub use reduce_axis::{ReduceAxis, ReduceOp}; diff --git a/src/linalg/optim.rs b/src/ml/optim.rs similarity index 92% rename from src/linalg/optim.rs rename to src/ml/optim.rs index 3c9b98e..4c7d759 100644 --- a/src/linalg/optim.rs +++ b/src/ml/optim.rs @@ -1,13 +1,13 @@ -//! Optimizer host dispatch (Adam). Added for zealot. +//! Optimizer host dispatch (Adam). -use crate::shaders::linalg::GpuAdam; +use crate::shaders::ml::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; +pub use vortx_shaders::ml::optim::AdamParams; /// The Adam optimizer kernel. #[derive(Shader)] diff --git a/vortx-shaders/src/linalg/activation.rs b/vortx-shaders/src/linalg/activation.rs deleted file mode 100644 index 04c880d..0000000 --- a/vortx-shaders/src/linalg/activation.rs +++ /dev/null @@ -1,54 +0,0 @@ -//! 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; -use glamx::UVec3; -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; - } -} diff --git a/vortx-shaders/src/linalg/mod.rs b/vortx-shaders/src/linalg/mod.rs index 1a2d8c5..0ecdaff 100644 --- a/vortx-shaders/src/linalg/mod.rs +++ b/vortx-shaders/src/linalg/mod.rs @@ -1,11 +1,9 @@ //! 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; @@ -16,16 +14,12 @@ pub use shape::{Shapes1, Shapes2, Shapes3}; // Re-export generated ShaderArgs structs (only available on host) #[cfg(not(target_arch_is_gpu))] -pub use activation::{GpuTanh, GpuTanhBackward}; -#[cfg(not(target_arch_is_gpu))] pub use contiguous::{Contiguous, ContiguousWithOffset}; #[cfg(not(target_arch_is_gpu))] pub use gemm::{GemmNaive, GemmTiled}; #[cfg(not(target_arch_is_gpu))] pub use op_assign::{GpuAdd, GpuCopy, GpuCopyWithOffsets, GpuDiv, GpuMul, GpuSub}; #[cfg(not(target_arch_is_gpu))] -pub use optim::GpuAdam; -#[cfg(not(target_arch_is_gpu))] pub use reduce::*; #[cfg(not(target_arch_is_gpu))] pub use repeat::Repeat; diff --git a/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs new file mode 100644 index 0000000..c503d7f --- /dev/null +++ b/vortx-shaders/src/ml/activation.rs @@ -0,0 +1,35 @@ +//! Backward passes of the element-wise activations. +//! +//! The forward directions live in [`crate::ml::unary`] (`UnaryOp::Tanh`, +//! `UnaryOp::Elu`); only the gradients are provided here, for training. + +use crate::linalg::Shape; +use crate::utils::iterators::StepRng; +use crate::utils::limits::MAX_NUM_WORKGROUPS; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +const WORKGROUP_SIZE: u32 = 256; +const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; + +/// Backward of tanh, in place: `g *= 1 - y*y`, where `y = tanh(x)` is the forward output. +/// +/// `g` and `y` must have the same shape. +#[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 StepRng::new(invocation_id.x..shape_g.len(), MAX_NUM_THREADS) { + 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; + } +} diff --git a/vortx-shaders/src/ml/mod.rs b/vortx-shaders/src/ml/mod.rs index 703f007..8984774 100644 --- a/vortx-shaders/src/ml/mod.rs +++ b/vortx-shaders/src/ml/mod.rs @@ -7,6 +7,7 @@ // #![allow(dead_code, non_snake_case)] // TODO: keep the modules private? +pub mod activation; pub mod batched_multiquery_attention; pub mod concat; pub mod conv2d; @@ -25,6 +26,7 @@ pub mod gemv_quant_q8_k; pub mod get_rel_pos; pub mod im2col; pub mod layernorm; +pub mod optim; pub mod pool2d; pub mod reduce_axis; pub mod rms_norm; @@ -35,6 +37,7 @@ pub mod softmax; pub mod unary; pub mod win_part; +pub use activation::*; pub use batched_multiquery_attention::*; pub use concat::*; pub use conv2d::*; @@ -44,6 +47,7 @@ pub use gather::*; pub use get_rel_pos::*; pub use im2col::*; pub use layernorm::*; +pub use optim::*; pub use pool2d::*; pub use reduce_axis::*; pub use rms_norm::*; diff --git a/vortx-shaders/src/linalg/optim.rs b/vortx-shaders/src/ml/optim.rs similarity index 86% rename from vortx-shaders/src/linalg/optim.rs rename to vortx-shaders/src/ml/optim.rs index da6b6df..4cbeec6 100644 --- a/vortx-shaders/src/linalg/optim.rs +++ b/vortx-shaders/src/ml/optim.rs @@ -1,14 +1,15 @@ -//! Optimizer kernels (Adam). Added for zealot; vortx upstream has no optimizers. +//! Optimizer kernels (Adam). -use super::shape::Shape; +use crate::linalg::Shape; +use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; -use glamx::UVec3; +use khal_std::glamx::UVec3; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; 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; @@ -16,10 +17,7 @@ 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) -)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] pub struct AdamParams { pub lr: f32, pub beta1: f32, @@ -46,7 +44,7 @@ pub fn gpu_adam( #[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) { + for thread_id in StepRng::new(invocation_id.x..shape.len(), MAX_NUM_THREADS) { let id = shape.decompose(thread_id); let i = shape.it_vec(id) as usize; let g = grad.read(i); From e9f89f89cc5e3ecd8420b4efac1555a522c81891 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 05/13] feat(ml): add ELU activation GPU ops (fwd/backward/vec4) Replaces #6 Co-Authored-By: Haixuan Xavier Tao --- Cargo.toml | 1 + src/ml/activation.rs | 83 +++++++++++++++++++++++++++++- vortx-shaders/src/ml/activation.rs | 71 ++++++++++++++++++++++++- 3 files changed, 153 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2caeccb..a33c412 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ khal = { version = "0.3", features = ["derive"]} [dependencies] bytemuck = "1" +glamx = { version = "0.3", default-features = false, features = ["bytemuck"] } include_dir = "0.7" nalgebra = "0.35" khal = { workspace = true } diff --git a/src/ml/activation.rs b/src/ml/activation.rs index 00a93ed..72c4427 100644 --- a/src/ml/activation.rs +++ b/src/ml/activation.rs @@ -3,7 +3,7 @@ //! The forward directions are provided by [`crate::ml::Unary`] //! (`UnaryOp::Tanh`, `UnaryOp::Elu`); only the gradients live here. -use crate::shaders::ml::GpuTanhBackward; +use crate::shaders::ml::{GpuElu, GpuEluBackward, GpuEluVec4, GpuTanhBackward}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; @@ -14,6 +14,12 @@ use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; pub struct ActivationBackward { /// 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 ActivationBackward { @@ -48,4 +54,79 @@ impl ActivationBackward { &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/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs index c503d7f..890927d 100644 --- a/vortx-shaders/src/ml/activation.rs +++ b/vortx-shaders/src/ml/activation.rs @@ -6,9 +6,11 @@ use crate::linalg::Shape; use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; -use khal_std::glamx::UVec3; +use khal_std::glamx::{UVec3, Vec4}; use khal_std::index::MaybeIndexUnchecked; use khal_std::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; @@ -33,3 +35,70 @@ pub fn gpu_tanh_backward( *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 }; + } +} From d074535e8a789a450c0605778455300e8a64dd8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 06/13] refactor(ml): keep only the ELU backward pass, the forward is already UnaryOp::Elu Completes #6 --- Cargo.toml | 1 - src/ml/activation.rs | 49 +--------------------------- vortx-shaders/src/ml/activation.rs | 52 ++---------------------------- 3 files changed, 4 insertions(+), 98 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index a33c412..2caeccb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,6 @@ khal = { version = "0.3", features = ["derive"]} [dependencies] bytemuck = "1" -glamx = { version = "0.3", default-features = false, features = ["bytemuck"] } include_dir = "0.7" nalgebra = "0.35" khal = { workspace = true } diff --git a/src/ml/activation.rs b/src/ml/activation.rs index 72c4427..dc91cad 100644 --- a/src/ml/activation.rs +++ b/src/ml/activation.rs @@ -3,7 +3,7 @@ //! The forward directions are provided by [`crate::ml::Unary`] //! (`UnaryOp::Tanh`, `UnaryOp::Elu`); only the gradients live here. -use crate::shaders::ml::{GpuElu, GpuEluBackward, GpuEluVec4, GpuTanhBackward}; +use crate::shaders::ml::{GpuEluBackward, GpuTanhBackward}; use crate::shapes::TensorLayoutBuffers; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; @@ -14,12 +14,8 @@ use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; pub struct ActivationBackward { /// 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 ActivationBackward { @@ -55,49 +51,6 @@ impl ActivationBackward { ) } - /// 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( diff --git a/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs index 890927d..cbb1b21 100644 --- a/vortx-shaders/src/ml/activation.rs +++ b/vortx-shaders/src/ml/activation.rs @@ -6,11 +6,9 @@ use crate::linalg::Shape; use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; -use khal_std::glamx::{UVec3, Vec4}; +use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::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; @@ -36,55 +34,11 @@ pub fn gpu_tanh_backward( } } -/// 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`. +/// and `y > 0 <=> x > 0`. #[spirv_bindgen] #[spirv(compute(threads(256, 1, 1)))] pub fn gpu_elu_backward( @@ -94,7 +48,7 @@ pub fn gpu_elu_backward( #[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) { + for thread_id in StepRng::new(invocation_id.x..shape_g.len(), MAX_NUM_THREADS) { 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; From 3b8b06b6b836428f0e654d4b49f1c7e3ce87c953 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 07/13] feat(ml): add PPO loss-gradient GPU kernels Replaces #5 Co-Authored-By: Haixuan Xavier Tao --- Cargo.toml | 1 + src/ml/mod.rs | 2 + src/ml/ppo.rs | 99 +++++++++++++++++++++ vortx-shaders/src/ml/mod.rs | 2 + vortx-shaders/src/ml/ppo.rs | 167 ++++++++++++++++++++++++++++++++++++ 5 files changed, 271 insertions(+) create mode 100644 src/ml/ppo.rs create mode 100644 vortx-shaders/src/ml/ppo.rs diff --git a/Cargo.toml b/Cargo.toml index 2caeccb..ef314a9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ khal = { version = "0.3", features = ["derive"]} [dependencies] bytemuck = "1" +glamx = { version = "0.2", default-features = false, features = ["bytemuck"] } include_dir = "0.7" nalgebra = "0.35" khal = { workspace = true } diff --git a/src/ml/mod.rs b/src/ml/mod.rs index f92da28..4c555c1 100644 --- a/src/ml/mod.rs +++ b/src/ml/mod.rs @@ -12,6 +12,7 @@ mod im2col; mod layernorm; mod optim; mod pool2d; +mod ppo; pub mod quantization; mod quantized_matrix; mod reduce_axis; @@ -38,6 +39,7 @@ pub use im2col::{Im2Col, Im2ColConfig}; pub use layernorm::LayerNorm; pub use optim::{Adam, AdamParams}; pub use pool2d::{GlobalPool2dConfig, Pool2d, Pool2dConfig, pool_output_size}; +pub use ppo::{Ppo, PpoActorParams, PpoValueParams}; pub use quantized_matrix::*; pub use reduce_axis::{ReduceAxis, ReduceOp}; pub use rms_norm::{RmsNorm, RmsNormConfig}; diff --git a/src/ml/ppo.rs b/src/ml/ppo.rs new file mode 100644 index 0000000..069555f --- /dev/null +++ b/src/ml/ppo.rs @@ -0,0 +1,99 @@ +//! PPO loss-gradient host dispatch. Added for zealot's GPU policy update. +//! +//! Wraps the two PPO output-gradient kernels (clipped-surrogate actor gradient + +//! log_std contribution, and clipped value-loss gradient). These have no `Shape` +//! uniform — dimensions ride in the params struct and indexing is row-major — so +//! no `TensorLayoutBuffers` is needed. + +use crate::shaders::linalg::{GpuPpoActorGrad, GpuPpoValueGrad}; +use crate::tensor::{AsTensorMut, AsTensorRef}; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; + +// Re-export the params structs from the shader crate. +pub use vortx_shaders::linalg::ppo::{PpoActorParams, PpoValueParams}; + +/// PPO loss-gradient kernels. +#[derive(Shader)] +pub struct Ppo { + /// Clipped-surrogate actor gradient + log_std contribution. + pub actor_grad: GpuPpoActorGrad, + /// Clipped value-loss gradient. + pub value_grad: GpuPpoValueGrad, +} + +impl Ppo { + /// Actor PPO gradient. All per-sample tensors are row-major `[action_dim x M]` + /// except `log_std` (`[action_dim]`), `adv` / `logp_old` (`[M]`). Writes + /// `g_mean` and `g_logstd` (`[action_dim x M]`). `params.num_cols` must equal `M`. + #[allow(clippy::too_many_arguments)] + pub fn actor_grad( + &self, + pass: &mut GpuPass, + params: impl AsTensorRef, + mean: impl AsTensorRef, + action: impl AsTensorRef, + log_std: impl AsTensorRef, + adv: impl AsTensorRef, + logp_old: impl AsTensorRef, + mut g_mean: impl AsTensorMut, + mut g_logstd: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let params = params.as_tensor_ref(); + let mean = mean.as_tensor_ref(); + let action = action.as_tensor_ref(); + let log_std = log_std.as_tensor_ref(); + let adv = adv.as_tensor_ref(); + let logp_old = logp_old.as_tensor_ref(); + let mut g_mean = g_mean.as_tensor_mut(); + let mut g_logstd = g_logstd.as_tensor_mut(); + + let num_threads = adv.len() as u32; // one thread per sample column + let mut buf_g_mean = g_mean.buffer_mut(); + let mut buf_g_logstd = g_logstd.buffer_mut(); + + self.actor_grad.call( + pass, + num_threads, + ¶ms.buffer(), + &mean.buffer(), + &action.buffer(), + &log_std.buffer(), + &adv.buffer(), + &logp_old.buffer(), + &mut buf_g_mean, + &mut buf_g_logstd, + ) + } + + /// Clipped value-loss gradient. `v_pred` / `value_old` / `ret` are `[M]`; + /// writes `g_v` (`[M]`). `params.num_cols` must equal `M`. + pub fn value_grad( + &self, + pass: &mut GpuPass, + params: impl AsTensorRef, + v_pred: impl AsTensorRef, + value_old: impl AsTensorRef, + ret: impl AsTensorRef, + mut g_v: impl AsTensorMut, + ) -> Result<(), GpuBackendError> { + let params = params.as_tensor_ref(); + let v_pred = v_pred.as_tensor_ref(); + let value_old = value_old.as_tensor_ref(); + let ret = ret.as_tensor_ref(); + let mut g_v = g_v.as_tensor_mut(); + + let num_threads = v_pred.len() as u32; + let mut buf_g_v = g_v.buffer_mut(); + + self.value_grad.call( + pass, + num_threads, + ¶ms.buffer(), + &v_pred.buffer(), + &value_old.buffer(), + &ret.buffer(), + &mut buf_g_v, + ) + } +} diff --git a/vortx-shaders/src/ml/mod.rs b/vortx-shaders/src/ml/mod.rs index 8984774..136ec0c 100644 --- a/vortx-shaders/src/ml/mod.rs +++ b/vortx-shaders/src/ml/mod.rs @@ -28,6 +28,7 @@ pub mod im2col; pub mod layernorm; pub mod optim; pub mod pool2d; +pub mod ppo; pub mod reduce_axis; pub mod rms_norm; pub mod rope; @@ -49,6 +50,7 @@ pub use im2col::*; pub use layernorm::*; pub use optim::*; pub use pool2d::*; +pub use ppo::*; pub use reduce_axis::*; pub use rms_norm::*; pub use rope::*; diff --git a/vortx-shaders/src/ml/ppo.rs b/vortx-shaders/src/ml/ppo.rs new file mode 100644 index 0000000..f538f82 --- /dev/null +++ b/vortx-shaders/src/ml/ppo.rs @@ -0,0 +1,167 @@ +//! PPO loss-gradient kernels (added for zealot's GPU policy update). +//! +//! These produce the per-sample OUTPUT gradients that feed the generic +//! GEMM/`elu_backward` backward backbone: the clipped-surrogate actor gradient +//! `g_mean` plus the state-independent `log_std` gradient contribution, and the +//! clipped value-loss gradient. An exact port of `zealot-rl`'s `minibatch_step` +//! (ppo.rs). Every per-sample tensor is row-major `[rows x M]` (M = minibatch +//! columns); one GPU thread handles one sample column `m`, looping over the +//! (small) action dimension internally. + +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 the actor PPO gradient (uniform buffer; 32 bytes). +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct PpoActorParams { + /// PPO clip epsilon. + pub clip: f32, + /// Entropy bonus coefficient (subtracted from the log_std gradient). + pub entropy_coef: f32, + /// Per-sample averaging factor `1 / minibatch_size`. + pub scale: f32, + /// `0.5·ln(2π)` — the Gaussian log-prob normalisation constant. + pub log_sqrt_2pi: f32, + /// Action dimensionality (rows). + pub action_dim: u32, + /// Number of sample columns `M`. + pub num_cols: u32, + pub pad0: u32, + pub pad1: u32, +} + +/// Scalar parameters for the clipped value-loss gradient (uniform; 32 bytes). +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] +pub struct PpoValueParams { + /// PPO clip epsilon (value clipping range). + pub clip: f32, + /// Value-loss coefficient. + pub value_coef: f32, + /// Per-sample averaging factor `1 / minibatch_size`. + pub scale: f32, + /// Number of sample columns `M`. + pub num_cols: u32, + pub pad0: u32, + pub pad1: u32, + pub pad2: u32, + pub pad3: u32, +} + +/// Clipped-surrogate actor gradient + log_std gradient contribution, per sample. +/// +/// For sample column `m` (one thread): compute the new diagonal-Gaussian +/// log-prob over the `action_dim` rows, the importance ratio +/// `exp(logp − logp_old)`, the PPO clip mask, then write `g_mean[k,m]` and +/// `g_logstd[k,m]` for every action dim `k`. Matches `minibatch_step`: +/// if !clipped: g_mean = −(adv·ratio·d/σ²)·scale, +/// g_logstd += −adv·ratio·(d²/σ² − 1)·scale, +/// always: g_logstd += −entropy_coef·scale. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_ppo_actor_grad( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &PpoActorParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mean: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] action: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] log_std: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] adv: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] logp_old: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] g_mean: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] g_logstd: &mut [f32], +) { + let a = params.action_dim as usize; + let m_cols = params.num_cols as usize; + let clip = params.clip; + let scale = params.scale; + let ent = params.entropy_coef; + for m in (invocation_id.x as usize..m_cols).step_by(MAX_NUM_THREADS as usize) { + // New log-prob over the action dims (matches ActorCritic::logp). + let mut logp = 0.0f32; + for k in 0..a { + let idx = k * m_cols + m; + let ls = log_std.read(k); + let std = ls.exp(); + let d = (action.read(idx) - mean.read(idx)) / std; + logp += -0.5 * d * d - ls - params.log_sqrt_2pi; + } + let ratio = (logp - logp_old.read(m)).exp(); + let av = adv.read(m); + let clipped = + (av >= 0.0 && ratio > 1.0 + clip) || (av < 0.0 && ratio < 1.0 - clip); + for k in 0..a { + let idx = k * m_cols + m; + let ls = log_std.read(k); + let inv_var = (-2.0 * ls).exp(); // 1/σ² + if clipped { + *g_mean.at_mut(idx) = 0.0; + *g_logstd.at_mut(idx) = -ent * scale; + } else { + let d = action.read(idx) - mean.read(idx); + *g_mean.at_mut(idx) = -(av * ratio * d * inv_var) * scale; + let dls = av * ratio * (d * d * inv_var - 1.0); + *g_logstd.at_mut(idx) = (-dls - ent) * scale; + } + } + } +} + +/// Clipped value-loss gradient, per sample. +/// +/// For sample column `m`: `v_clipped = value_old + clamp(v − value_old, ±clip)`, +/// and `dv = 2·(v_clipped − ret)` if the clipped squared error is larger else +/// `2·(v − ret)`; writes `g_v[m] = value_coef·dv·scale`. Matches `minibatch_step`. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_ppo_value_grad( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &PpoValueParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] v_pred: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] value_old: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] ret: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] g_v: &mut [f32], +) { + let m_cols = params.num_cols as usize; + let clip = params.clip; + let scale = params.scale; + for m in (invocation_id.x as usize..m_cols).step_by(MAX_NUM_THREADS as usize) { + let v = v_pred.read(m); + let vo = value_old.read(m); + let r = ret.read(m); + let diff = v - vo; + let clamped = if diff > clip { + clip + } else if diff < -clip { + -clip + } else { + diff + }; + let v_clipped = vo + clamped; + let l_un = (v - r) * (v - r); + let l_cl = (v_clipped - r) * (v_clipped - r); + let dv = if l_cl > l_un { + 2.0 * (v_clipped - r) + } else { + 2.0 * (v - r) + }; + *g_v.at_mut(m) = params.value_coef * dv * scale; + } +} From e3b30c16a2284ecde8819bbd78fc60d5e9c66250 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 08/13] refactor(ml): wire the PPO kernels to vortx::ml and use StepRng for uniform control flow Completes #5 --- Cargo.toml | 1 - src/ml/ppo.rs | 14 ++++++------ vortx-shaders/src/ml/ppo.rs | 45 +++++++++++++++++-------------------- 3 files changed, 27 insertions(+), 33 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ef314a9..2caeccb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,6 @@ khal = { version = "0.3", features = ["derive"]} [dependencies] bytemuck = "1" -glamx = { version = "0.2", default-features = false, features = ["bytemuck"] } include_dir = "0.7" nalgebra = "0.35" khal = { workspace = true } diff --git a/src/ml/ppo.rs b/src/ml/ppo.rs index 069555f..bec4e68 100644 --- a/src/ml/ppo.rs +++ b/src/ml/ppo.rs @@ -1,17 +1,17 @@ -//! PPO loss-gradient host dispatch. Added for zealot's GPU policy update. +//! PPO loss-gradient host dispatch. //! -//! Wraps the two PPO output-gradient kernels (clipped-surrogate actor gradient + -//! log_std contribution, and clipped value-loss gradient). These have no `Shape` -//! uniform — dimensions ride in the params struct and indexing is row-major — so -//! no `TensorLayoutBuffers` is needed. +//! Wraps the two PPO output-gradient kernels (clipped-surrogate actor gradient +//! plus log_std contribution, and clipped value-loss gradient). They need no +//! `Shape` uniform or `TensorLayoutBuffers`: dimensions ride in the params +//! struct and indexing is row-major. -use crate::shaders::linalg::{GpuPpoActorGrad, GpuPpoValueGrad}; +use crate::shaders::ml::{GpuPpoActorGrad, GpuPpoValueGrad}; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; // Re-export the params structs from the shader crate. -pub use vortx_shaders::linalg::ppo::{PpoActorParams, PpoValueParams}; +pub use vortx_shaders::ml::ppo::{PpoActorParams, PpoValueParams}; /// PPO loss-gradient kernels. #[derive(Shader)] diff --git a/vortx-shaders/src/ml/ppo.rs b/vortx-shaders/src/ml/ppo.rs index f538f82..5b943ff 100644 --- a/vortx-shaders/src/ml/ppo.rs +++ b/vortx-shaders/src/ml/ppo.rs @@ -1,21 +1,21 @@ -//! PPO loss-gradient kernels (added for zealot's GPU policy update). +//! PPO loss-gradient kernels. //! -//! These produce the per-sample OUTPUT gradients that feed the generic +//! These produce the per-sample output gradients that feed the generic //! GEMM/`elu_backward` backward backbone: the clipped-surrogate actor gradient //! `g_mean` plus the state-independent `log_std` gradient contribution, and the -//! clipped value-loss gradient. An exact port of `zealot-rl`'s `minibatch_step` -//! (ppo.rs). Every per-sample tensor is row-major `[rows x M]` (M = minibatch -//! columns); one GPU thread handles one sample column `m`, looping over the -//! (small) action dimension internally. +//! clipped value-loss gradient. Every per-sample tensor is row-major +//! `[rows x M]` (M = minibatch columns); one GPU thread handles one sample +//! column `m`, looping over the (small) action dimension internally. +use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; -use glamx::UVec3; +use khal_std::glamx::UVec3; +#[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] +use khal_std::num_traits::Float; 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; @@ -23,10 +23,7 @@ const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Scalar parameters for the actor PPO gradient (uniform buffer; 32 bytes). #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr( - not(any(target_arch = "spirv", target_arch = "nvptx64")), - derive(bytemuck::Pod, bytemuck::Zeroable) -)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] pub struct PpoActorParams { /// PPO clip epsilon. pub clip: f32, @@ -34,7 +31,7 @@ pub struct PpoActorParams { pub entropy_coef: f32, /// Per-sample averaging factor `1 / minibatch_size`. pub scale: f32, - /// `0.5·ln(2π)` — the Gaussian log-prob normalisation constant. + /// `0.5·ln(2π)`: the Gaussian log-prob normalisation constant. pub log_sqrt_2pi: f32, /// Action dimensionality (rows). pub action_dim: u32, @@ -47,10 +44,7 @@ pub struct PpoActorParams { /// Scalar parameters for the clipped value-loss gradient (uniform; 32 bytes). #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr( - not(any(target_arch = "spirv", target_arch = "nvptx64")), - derive(bytemuck::Pod, bytemuck::Zeroable) -)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] pub struct PpoValueParams { /// PPO clip epsilon (value clipping range). pub clip: f32, @@ -71,7 +65,7 @@ pub struct PpoValueParams { /// For sample column `m` (one thread): compute the new diagonal-Gaussian /// log-prob over the `action_dim` rows, the importance ratio /// `exp(logp − logp_old)`, the PPO clip mask, then write `g_mean[k,m]` and -/// `g_logstd[k,m]` for every action dim `k`. Matches `minibatch_step`: +/// `g_logstd[k,m]` for every action dim `k`: /// if !clipped: g_mean = −(adv·ratio·d/σ²)·scale, /// g_logstd += −adv·ratio·(d²/σ² − 1)·scale, /// always: g_logstd += −entropy_coef·scale. @@ -93,8 +87,9 @@ pub fn gpu_ppo_actor_grad( let clip = params.clip; let scale = params.scale; let ent = params.entropy_coef; - for m in (invocation_id.x as usize..m_cols).step_by(MAX_NUM_THREADS as usize) { - // New log-prob over the action dims (matches ActorCritic::logp). + for m in StepRng::new(invocation_id.x..m_cols as u32, MAX_NUM_THREADS) { + let m = m as usize; + // New log-prob over the action dims. let mut logp = 0.0f32; for k in 0..a { let idx = k * m_cols + m; @@ -105,8 +100,7 @@ pub fn gpu_ppo_actor_grad( } let ratio = (logp - logp_old.read(m)).exp(); let av = adv.read(m); - let clipped = - (av >= 0.0 && ratio > 1.0 + clip) || (av < 0.0 && ratio < 1.0 - clip); + let clipped = (av >= 0.0 && ratio > 1.0 + clip) || (av < 0.0 && ratio < 1.0 - clip); for k in 0..a { let idx = k * m_cols + m; let ls = log_std.read(k); @@ -128,7 +122,7 @@ pub fn gpu_ppo_actor_grad( /// /// For sample column `m`: `v_clipped = value_old + clamp(v − value_old, ±clip)`, /// and `dv = 2·(v_clipped − ret)` if the clipped squared error is larger else -/// `2·(v − ret)`; writes `g_v[m] = value_coef·dv·scale`. Matches `minibatch_step`. +/// `2·(v − ret)`; writes `g_v[m] = value_coef·dv·scale`. #[spirv_bindgen] #[spirv(compute(threads(256, 1, 1)))] pub fn gpu_ppo_value_grad( @@ -142,7 +136,8 @@ pub fn gpu_ppo_value_grad( let m_cols = params.num_cols as usize; let clip = params.clip; let scale = params.scale; - for m in (invocation_id.x as usize..m_cols).step_by(MAX_NUM_THREADS as usize) { + for m in StepRng::new(invocation_id.x..m_cols as u32, MAX_NUM_THREADS) { + let m = m as usize; let v = v_pred.read(m); let vo = value_old.read(m); let r = ret.read(m); From 2a5ad2d787b3334637a332000076489aef853092 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 09/13] perf(linalg): vec4 GEMM (compute-FMA inner loop + vec4 global-load variant) Replaces #7 Co-Authored-By: Haixuan Xavier Tao --- src/linalg/gemm.rs | 45 +++++++++- vortx-shaders/src/linalg/gemm.rs | 143 +++++++++++++++++++++++++------ vortx-shaders/src/linalg/mod.rs | 2 +- 3 files changed, 162 insertions(+), 28 deletions(-) diff --git a/src/linalg/gemm.rs b/src/linalg/gemm.rs index cd43867..b69d2cc 100644 --- a/src/linalg/gemm.rs +++ b/src/linalg/gemm.rs @@ -4,7 +4,7 @@ use khal::Shader; use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; // Use generated ShaderArgs from spirv_bindgen -use crate::shaders::linalg::{GemmNaive, GemmTiled}; +use crate::shaders::linalg::{GemmNaive, GemmTiled, GemmTiledVec4}; /// Indicates if a matrix needs to be considered as-is or as its transpose when running a matrix /// multiplication operation. @@ -38,6 +38,8 @@ pub struct Gemm { pub gemm_naive: GemmNaive, /// Optimized tiled GEMM using shared memory (64x64 tiles, 16x16 workgroups). pub gemm_tiled: GemmTiled, + /// Tiled GEMM with 128-bit vec4 global loads (contiguous, tile-aligned only). + pub gemm_tiled_vec4: GemmTiledVec4, } // GemmArgs is now generated by spirv_bindgen from vortx_shaders::linalg::gemm @@ -304,6 +306,47 @@ impl Gemm { } } + /// vec4-global-load tiled GEMM. The caller MUST guarantee: CONTIGUOUS + /// row-major `lhs`/`rhs`, single batch, and `M%64==0 && N%64==0 && K%16==0` + /// (full tiles). Used only for the qualifying forward hidden GEMMs; all other + /// GEMMs (transposed backward, K=43/49 input layer) use `dispatch_tiled`. + pub fn dispatch_tiled_vec4( + &self, + backend: &GpuBackend, + shapes: &mut TensorLayoutBuffers, + pass: &mut GpuPass, + mut out: impl AsTensorMut, + lhs: impl AsTensorRef, + rhs: impl AsTensorRef, + ) -> Result<(), GpuBackendError> { + let mut out = out.as_tensor_mut(); + let lhs = lhs.as_tensor_ref(); + let rhs = rhs.as_tensor_ref(); + + let shape_out = out.layout().canonicalize(); + let shape_lhs = lhs.layout().canonicalize(); + let shape_rhs = rhs.layout().canonicalize(); + let grid = Self::tiled_grid(&shape_out); + + shapes.insert(backend, shape_out)?; + shapes.insert(backend, shape_lhs)?; + shapes.insert(backend, shape_rhs)?; + let mut buf_out = out.buffer_mut(); + let lhs_v4 = lhs.buffer().reinterpret::(); + let rhs_v4 = rhs.buffer().reinterpret::(); + + self.gemm_tiled_vec4.call( + pass, + grid, + &shapes.get(shape_out).unwrap().as_slice(), + &shapes.get(shape_lhs).unwrap().as_slice(), + &shapes.get(shape_rhs).unwrap().as_slice(), + &mut buf_out, + &lhs_v4, + &rhs_v4, + ) + } + fn naive_grid(shape_out: &crate::shapes::TensorLayout) -> [u32; 3] { [shape_out.size[3], shape_out.size[2], shape_out.size[1]] } diff --git a/vortx-shaders/src/linalg/gemm.rs b/vortx-shaders/src/linalg/gemm.rs index 17e1a08..c697e19 100644 --- a/vortx-shaders/src/linalg/gemm.rs +++ b/vortx-shaders/src/linalg/gemm.rs @@ -5,7 +5,7 @@ use super::shape::Shape; #[cfg(feature = "push_constants")] use super::shape::Shapes3; -use glamx::UVec3; +use glamx::{UVec3, Vec4}; use khal_std::{ index::MaybeIndexUnchecked, macros::{spirv, spirv_bindgen}, @@ -78,14 +78,15 @@ pub fn gemm_tiled( let m = shape_out.h; let n = shape_out.w; let k = shape_lhs.w; - let mut acc: [f32; 16]; + // Register accumulator: 4 rows × a vec4 of the 4 output columns per thread. + // The inner loop does vec4 FMAs (4-wide) instead of 16 scalar MACs. + let mut acc: [Vec4; 4]; // Process batch dimension for batch in 0..shape_out.n { let batch_c = wg_id.z % shape_out.c; - // Register accumulator for 4x4 outputs per thread - acc = [0.0; 16]; + acc = [Vec4::ZERO; 4]; // Loop over K dimension in tiles let mut k_tile: u32 = 0; @@ -141,27 +142,17 @@ pub fn gemm_tiled( let a2 = smem_a.read(((a_row_base + 2) * SMEM_A_STRIDE + kk) as usize); let a3 = smem_a.read(((a_row_base + 3) * SMEM_A_STRIDE + kk) as usize); - let b0 = smem_b.read((kk * SMEM_B_STRIDE + b_col_base) as usize); - let b1 = smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 1) as usize); - let b2 = smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 2) as usize); - let b3 = smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 3) as usize); - - acc[0] += a0 * b0; - acc[1] += a0 * b1; - acc[2] += a0 * b2; - acc[3] += a0 * b3; - acc[4] += a1 * b0; - acc[5] += a1 * b1; - acc[6] += a1 * b2; - acc[7] += a1 * b3; - acc[8] += a2 * b0; - acc[9] += a2 * b1; - acc[10] += a2 * b2; - acc[11] += a2 * b3; - acc[12] += a3 * b0; - acc[13] += a3 * b1; - acc[14] += a3 * b2; - acc[15] += a3 * b3; + // 4 contiguous B columns as a vec4; 4-wide FMA per A row. + let bvec = Vec4::new( + smem_b.read((kk * SMEM_B_STRIDE + b_col_base) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 1) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 2) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 3) as usize), + ); + acc[0] = bvec.mul_add(Vec4::splat(a0), acc[0]); + acc[1] = bvec.mul_add(Vec4::splat(a1), acc[1]); + acc[2] = bvec.mul_add(Vec4::splat(a2), acc[2]); + acc[3] = bvec.mul_add(Vec4::splat(a3), acc[3]); kk += 1; } @@ -178,12 +169,13 @@ pub fn gemm_tiled( while i < THREAD_M { let row = out_row + i; if row < m { + let arr = acc[i as usize].to_array(); let mut j: u32 = 0; while j < THREAD_N { let col = out_col + j; if col < n { let idx = shape_out.it(batch, batch_c, row, col) as usize; - out.write(idx, acc[(i * THREAD_N + j) as usize]); + out.write(idx, arr[j as usize]); } j += 1; } @@ -193,6 +185,105 @@ pub fn gemm_tiled( } } +/// vec4 tiled GEMM: identical math to `gemm_tiled` but loads the A/B tiles from +/// global memory with **128-bit vec4 transactions** (4 contiguous f32 each). It +/// assumes CONTIGUOUS row-major `lhs`/`rhs`, a single batch, and dims that fill +/// the tiles exactly (`M%TILE_M==0`, `N%TILE_N==0`, `K%TILE_K==0`) so no boundary +/// handling is needed. The host (`dispatch_tiled_vec4`) enforces this and falls +/// back to `gemm_tiled` (scalar) for transposed/odd-dim cases (backward GEMMs, +/// input layer with K=43/49). +#[spirv_bindgen] +#[spirv(compute(threads(16, 16, 1)))] +pub fn gemm_tiled_vec4( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] wg_id: UVec3, + #[spirv(workgroup)] smem_a: &mut [f32; SMEM_A_SIZE], + #[spirv(workgroup)] smem_b: &mut [f32; SMEM_B_SIZE], + #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_out: &Shape, + #[spirv(uniform, descriptor_set = 0, binding = 1)] shape_lhs: &Shape, + #[spirv(uniform, descriptor_set = 0, binding = 2)] shape_rhs: &Shape, + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] out: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] lhs: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] rhs: &[Vec4], +) { + let _ = shape_rhs; + let tid_x = local_id.x; + let tid_y = local_id.y; + let linear_tid = tid_y * WG_N + tid_x; + let tile_row = wg_id.y * TILE_M; + let tile_col = wg_id.x * TILE_N; + let n = shape_out.w; + let k = shape_lhs.w; + + let mut acc: [Vec4; 4] = [Vec4::ZERO; 4]; + + let mut k_tile: u32 = 0; + while k_tile < k { + // Load A tile: one vec4 (4 contiguous f32 of a row) per thread. + { + let lin = linear_tid * 4; + let row = lin / TILE_K; + let col0 = lin % TILE_K; + let v = lhs.read((((tile_row + row) * k + k_tile + col0) / 4) as usize); + let base = (row * SMEM_A_STRIDE + col0) as usize; + smem_a.write(base, v.x); + smem_a.write(base + 1, v.y); + smem_a.write(base + 2, v.z); + smem_a.write(base + 3, v.w); + } + // Load B tile: one vec4 per thread. + { + let lin = linear_tid * 4; + let row = lin / TILE_N; + let col0 = lin % TILE_N; + let v = rhs.read((((k_tile + row) * n + tile_col + col0) / 4) as usize); + let base = (row * SMEM_B_STRIDE + col0) as usize; + smem_b.write(base, v.x); + smem_b.write(base + 1, v.y); + smem_b.write(base + 2, v.z); + smem_b.write(base + 3, v.w); + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + + let a_row_base = tid_y * THREAD_M; + let b_col_base = tid_x * THREAD_N; + let mut kk: u32 = 0; + while kk < TILE_K { + let a0 = smem_a.read((a_row_base * SMEM_A_STRIDE + kk) as usize); + let a1 = smem_a.read(((a_row_base + 1) * SMEM_A_STRIDE + kk) as usize); + let a2 = smem_a.read(((a_row_base + 2) * SMEM_A_STRIDE + kk) as usize); + let a3 = smem_a.read(((a_row_base + 3) * SMEM_A_STRIDE + kk) as usize); + let bvec = Vec4::new( + smem_b.read((kk * SMEM_B_STRIDE + b_col_base) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 1) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 2) as usize), + smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 3) as usize), + ); + acc[0] = bvec.mul_add(Vec4::splat(a0), acc[0]); + acc[1] = bvec.mul_add(Vec4::splat(a1), acc[1]); + acc[2] = bvec.mul_add(Vec4::splat(a2), acc[2]); + acc[3] = bvec.mul_add(Vec4::splat(a3), acc[3]); + kk += 1; + } + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + k_tile += TILE_K; + } + + // Store (contiguous out, full tile -> no bounds checks). + let out_row = tile_row + tid_y * THREAD_M; + let out_col = tile_col + tid_x * THREAD_N; + let mut i: u32 = 0; + while i < THREAD_M { + let arr = acc[i as usize].to_array(); + let mut j: u32 = 0; + while j < THREAD_N { + out.write((((out_row + i) * n) + out_col + j) as usize, arr[j as usize]); + j += 1; + } + i += 1; + } +} + /// Naive GEMM (kept for reference and small matrices) #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] diff --git a/vortx-shaders/src/linalg/mod.rs b/vortx-shaders/src/linalg/mod.rs index 0ecdaff..f0d61ae 100644 --- a/vortx-shaders/src/linalg/mod.rs +++ b/vortx-shaders/src/linalg/mod.rs @@ -16,7 +16,7 @@ pub use shape::{Shapes1, Shapes2, Shapes3}; #[cfg(not(target_arch_is_gpu))] pub use contiguous::{Contiguous, ContiguousWithOffset}; #[cfg(not(target_arch_is_gpu))] -pub use gemm::{GemmNaive, GemmTiled}; +pub use gemm::{GemmNaive, GemmTiled, GemmTiledVec4}; #[cfg(not(target_arch_is_gpu))] pub use op_assign::{GpuAdd, GpuCopy, GpuCopyWithOffsets, GpuDiv, GpuMul, GpuSub}; #[cfg(not(target_arch_is_gpu))] From 7f46b2d788be2ef89f6b61731e38af5b27b107a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 10/13] perf(linalg): keep the vec4 FMA inner loop, drop the unmeasured vec4 global-load GEMM variant Completes #7 --- src/linalg/gemm.rs | 45 +------------- src/ml/ppo.rs | 2 +- vortx-shaders/src/linalg/gemm.rs | 102 +------------------------------ vortx-shaders/src/linalg/mod.rs | 2 +- 4 files changed, 4 insertions(+), 147 deletions(-) diff --git a/src/linalg/gemm.rs b/src/linalg/gemm.rs index b69d2cc..cd43867 100644 --- a/src/linalg/gemm.rs +++ b/src/linalg/gemm.rs @@ -4,7 +4,7 @@ use khal::Shader; use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; // Use generated ShaderArgs from spirv_bindgen -use crate::shaders::linalg::{GemmNaive, GemmTiled, GemmTiledVec4}; +use crate::shaders::linalg::{GemmNaive, GemmTiled}; /// Indicates if a matrix needs to be considered as-is or as its transpose when running a matrix /// multiplication operation. @@ -38,8 +38,6 @@ pub struct Gemm { pub gemm_naive: GemmNaive, /// Optimized tiled GEMM using shared memory (64x64 tiles, 16x16 workgroups). pub gemm_tiled: GemmTiled, - /// Tiled GEMM with 128-bit vec4 global loads (contiguous, tile-aligned only). - pub gemm_tiled_vec4: GemmTiledVec4, } // GemmArgs is now generated by spirv_bindgen from vortx_shaders::linalg::gemm @@ -306,47 +304,6 @@ impl Gemm { } } - /// vec4-global-load tiled GEMM. The caller MUST guarantee: CONTIGUOUS - /// row-major `lhs`/`rhs`, single batch, and `M%64==0 && N%64==0 && K%16==0` - /// (full tiles). Used only for the qualifying forward hidden GEMMs; all other - /// GEMMs (transposed backward, K=43/49 input layer) use `dispatch_tiled`. - pub fn dispatch_tiled_vec4( - &self, - backend: &GpuBackend, - shapes: &mut TensorLayoutBuffers, - pass: &mut GpuPass, - mut out: impl AsTensorMut, - lhs: impl AsTensorRef, - rhs: impl AsTensorRef, - ) -> Result<(), GpuBackendError> { - let mut out = out.as_tensor_mut(); - let lhs = lhs.as_tensor_ref(); - let rhs = rhs.as_tensor_ref(); - - let shape_out = out.layout().canonicalize(); - let shape_lhs = lhs.layout().canonicalize(); - let shape_rhs = rhs.layout().canonicalize(); - let grid = Self::tiled_grid(&shape_out); - - shapes.insert(backend, shape_out)?; - shapes.insert(backend, shape_lhs)?; - shapes.insert(backend, shape_rhs)?; - let mut buf_out = out.buffer_mut(); - let lhs_v4 = lhs.buffer().reinterpret::(); - let rhs_v4 = rhs.buffer().reinterpret::(); - - self.gemm_tiled_vec4.call( - pass, - grid, - &shapes.get(shape_out).unwrap().as_slice(), - &shapes.get(shape_lhs).unwrap().as_slice(), - &shapes.get(shape_rhs).unwrap().as_slice(), - &mut buf_out, - &lhs_v4, - &rhs_v4, - ) - } - fn naive_grid(shape_out: &crate::shapes::TensorLayout) -> [u32; 3] { [shape_out.size[3], shape_out.size[2], shape_out.size[1]] } diff --git a/src/ml/ppo.rs b/src/ml/ppo.rs index bec4e68..6da232d 100644 --- a/src/ml/ppo.rs +++ b/src/ml/ppo.rs @@ -8,7 +8,7 @@ use crate::shaders::ml::{GpuPpoActorGrad, GpuPpoValueGrad}; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::backend::{GpuBackendError, GpuPass}; // Re-export the params structs from the shader crate. pub use vortx_shaders::ml::ppo::{PpoActorParams, PpoValueParams}; diff --git a/vortx-shaders/src/linalg/gemm.rs b/vortx-shaders/src/linalg/gemm.rs index c697e19..490cec6 100644 --- a/vortx-shaders/src/linalg/gemm.rs +++ b/vortx-shaders/src/linalg/gemm.rs @@ -78,8 +78,7 @@ pub fn gemm_tiled( let m = shape_out.h; let n = shape_out.w; let k = shape_lhs.w; - // Register accumulator: 4 rows × a vec4 of the 4 output columns per thread. - // The inner loop does vec4 FMAs (4-wide) instead of 16 scalar MACs. + // Register accumulator: 4 rows, each a vec4 of the thread's 4 output columns. let mut acc: [Vec4; 4]; // Process batch dimension @@ -185,105 +184,6 @@ pub fn gemm_tiled( } } -/// vec4 tiled GEMM: identical math to `gemm_tiled` but loads the A/B tiles from -/// global memory with **128-bit vec4 transactions** (4 contiguous f32 each). It -/// assumes CONTIGUOUS row-major `lhs`/`rhs`, a single batch, and dims that fill -/// the tiles exactly (`M%TILE_M==0`, `N%TILE_N==0`, `K%TILE_K==0`) so no boundary -/// handling is needed. The host (`dispatch_tiled_vec4`) enforces this and falls -/// back to `gemm_tiled` (scalar) for transposed/odd-dim cases (backward GEMMs, -/// input layer with K=43/49). -#[spirv_bindgen] -#[spirv(compute(threads(16, 16, 1)))] -pub fn gemm_tiled_vec4( - #[spirv(local_invocation_id)] local_id: UVec3, - #[spirv(workgroup_id)] wg_id: UVec3, - #[spirv(workgroup)] smem_a: &mut [f32; SMEM_A_SIZE], - #[spirv(workgroup)] smem_b: &mut [f32; SMEM_B_SIZE], - #[spirv(uniform, descriptor_set = 0, binding = 0)] shape_out: &Shape, - #[spirv(uniform, descriptor_set = 0, binding = 1)] shape_lhs: &Shape, - #[spirv(uniform, descriptor_set = 0, binding = 2)] shape_rhs: &Shape, - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] out: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] lhs: &[Vec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] rhs: &[Vec4], -) { - let _ = shape_rhs; - let tid_x = local_id.x; - let tid_y = local_id.y; - let linear_tid = tid_y * WG_N + tid_x; - let tile_row = wg_id.y * TILE_M; - let tile_col = wg_id.x * TILE_N; - let n = shape_out.w; - let k = shape_lhs.w; - - let mut acc: [Vec4; 4] = [Vec4::ZERO; 4]; - - let mut k_tile: u32 = 0; - while k_tile < k { - // Load A tile: one vec4 (4 contiguous f32 of a row) per thread. - { - let lin = linear_tid * 4; - let row = lin / TILE_K; - let col0 = lin % TILE_K; - let v = lhs.read((((tile_row + row) * k + k_tile + col0) / 4) as usize); - let base = (row * SMEM_A_STRIDE + col0) as usize; - smem_a.write(base, v.x); - smem_a.write(base + 1, v.y); - smem_a.write(base + 2, v.z); - smem_a.write(base + 3, v.w); - } - // Load B tile: one vec4 per thread. - { - let lin = linear_tid * 4; - let row = lin / TILE_N; - let col0 = lin % TILE_N; - let v = rhs.read((((k_tile + row) * n + tile_col + col0) / 4) as usize); - let base = (row * SMEM_B_STRIDE + col0) as usize; - smem_b.write(base, v.x); - smem_b.write(base + 1, v.y); - smem_b.write(base + 2, v.z); - smem_b.write(base + 3, v.w); - } - khal_std::sync::workgroup_memory_barrier_with_group_sync(); - - let a_row_base = tid_y * THREAD_M; - let b_col_base = tid_x * THREAD_N; - let mut kk: u32 = 0; - while kk < TILE_K { - let a0 = smem_a.read((a_row_base * SMEM_A_STRIDE + kk) as usize); - let a1 = smem_a.read(((a_row_base + 1) * SMEM_A_STRIDE + kk) as usize); - let a2 = smem_a.read(((a_row_base + 2) * SMEM_A_STRIDE + kk) as usize); - let a3 = smem_a.read(((a_row_base + 3) * SMEM_A_STRIDE + kk) as usize); - let bvec = Vec4::new( - smem_b.read((kk * SMEM_B_STRIDE + b_col_base) as usize), - smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 1) as usize), - smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 2) as usize), - smem_b.read((kk * SMEM_B_STRIDE + b_col_base + 3) as usize), - ); - acc[0] = bvec.mul_add(Vec4::splat(a0), acc[0]); - acc[1] = bvec.mul_add(Vec4::splat(a1), acc[1]); - acc[2] = bvec.mul_add(Vec4::splat(a2), acc[2]); - acc[3] = bvec.mul_add(Vec4::splat(a3), acc[3]); - kk += 1; - } - khal_std::sync::workgroup_memory_barrier_with_group_sync(); - k_tile += TILE_K; - } - - // Store (contiguous out, full tile -> no bounds checks). - let out_row = tile_row + tid_y * THREAD_M; - let out_col = tile_col + tid_x * THREAD_N; - let mut i: u32 = 0; - while i < THREAD_M { - let arr = acc[i as usize].to_array(); - let mut j: u32 = 0; - while j < THREAD_N { - out.write((((out_row + i) * n) + out_col + j) as usize, arr[j as usize]); - j += 1; - } - i += 1; - } -} - /// Naive GEMM (kept for reference and small matrices) #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] diff --git a/vortx-shaders/src/linalg/mod.rs b/vortx-shaders/src/linalg/mod.rs index f0d61ae..0ecdaff 100644 --- a/vortx-shaders/src/linalg/mod.rs +++ b/vortx-shaders/src/linalg/mod.rs @@ -16,7 +16,7 @@ pub use shape::{Shapes1, Shapes2, Shapes3}; #[cfg(not(target_arch_is_gpu))] pub use contiguous::{Contiguous, ContiguousWithOffset}; #[cfg(not(target_arch_is_gpu))] -pub use gemm::{GemmNaive, GemmTiled, GemmTiledVec4}; +pub use gemm::{GemmNaive, GemmTiled}; #[cfg(not(target_arch_is_gpu))] pub use op_assign::{GpuAdd, GpuCopy, GpuCopyWithOffsets, GpuDiv, GpuMul, GpuSub}; #[cfg(not(target_arch_is_gpu))] From 71f7cb1c309c141399160b1f6125c82ef5b3da21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 11/13] feat(ml): gpu_ppo_stage_batch, building the PPO minibatch on device --- src/linalg/contiguous.rs | 2 +- src/linalg/op_assign.rs | 4 +- src/ml/activation.rs | 4 +- src/ml/mod.rs | 2 +- src/ml/optim.rs | 2 +- src/ml/ppo.rs | 57 ++++++++++++++++-- src/ml/unary.rs | 4 +- vortx-shaders/src/linalg/contiguous.rs | 6 +- vortx-shaders/src/linalg/op_assign.rs | 6 +- vortx-shaders/src/ml/activation.rs | 6 +- vortx-shaders/src/ml/optim.rs | 6 +- vortx-shaders/src/ml/ppo.rs | 83 +++++++++++++++++++++++++- vortx-shaders/src/ml/unary.rs | 6 +- 13 files changed, 169 insertions(+), 19 deletions(-) diff --git a/src/linalg/contiguous.rs b/src/linalg/contiguous.rs index 8c5289f..af716b5 100644 --- a/src/linalg/contiguous.rs +++ b/src/linalg/contiguous.rs @@ -51,7 +51,7 @@ 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"))] diff --git a/src/linalg/op_assign.rs b/src/linalg/op_assign.rs index 4f45e48..6fbc4cd 100644 --- a/src/linalg/op_assign.rs +++ b/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)?; diff --git a/src/ml/activation.rs b/src/ml/activation.rs index dc91cad..4e70e6b 100644 --- a/src/ml/activation.rs +++ b/src/ml/activation.rs @@ -33,7 +33,7 @@ impl ActivationBackward { 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; + 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)?; @@ -65,7 +65,7 @@ impl ActivationBackward { 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; + 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)?; diff --git a/src/ml/mod.rs b/src/ml/mod.rs index 4c555c1..d9d2195 100644 --- a/src/ml/mod.rs +++ b/src/ml/mod.rs @@ -39,7 +39,7 @@ pub use im2col::{Im2Col, Im2ColConfig}; pub use layernorm::LayerNorm; pub use optim::{Adam, AdamParams}; pub use pool2d::{GlobalPool2dConfig, Pool2d, Pool2dConfig, pool_output_size}; -pub use ppo::{Ppo, PpoActorParams, PpoValueParams}; +pub use ppo::{Ppo, PpoActorParams, PpoStageParams, PpoValueParams}; pub use quantized_matrix::*; pub use reduce_axis::{ReduceAxis, ReduceOp}; pub use rms_norm::{RmsNorm, RmsNormConfig}; diff --git a/src/ml/optim.rs b/src/ml/optim.rs index 4c7d759..b176b6b 100644 --- a/src/ml/optim.rs +++ b/src/ml/optim.rs @@ -39,7 +39,7 @@ impl Adam { let mut v = v.as_tensor_mut(); let shape = theta.layout().canonicalize(); - let num_threads = theta.len() as u32; + let num_threads = (theta.len() as u32).min(vortx_shaders::ml::optim::MAX_NUM_THREADS); shapes.insert(backend, shape)?; let shape_buf = shapes.get(shape).unwrap(); diff --git a/src/ml/ppo.rs b/src/ml/ppo.rs index 6da232d..e265be0 100644 --- a/src/ml/ppo.rs +++ b/src/ml/ppo.rs @@ -5,13 +5,13 @@ //! `Shape` uniform or `TensorLayoutBuffers`: dimensions ride in the params //! struct and indexing is row-major. -use crate::shaders::ml::{GpuPpoActorGrad, GpuPpoValueGrad}; +use crate::shaders::ml::{GpuPpoActorGrad, GpuPpoStageBatch, GpuPpoValueGrad}; use crate::tensor::{AsTensorMut, AsTensorRef}; use khal::Shader; use khal::backend::{GpuBackendError, GpuPass}; // Re-export the params structs from the shader crate. -pub use vortx_shaders::ml::ppo::{PpoActorParams, PpoValueParams}; +pub use vortx_shaders::ml::ppo::{PpoActorParams, PpoStageParams, PpoValueParams}; /// PPO loss-gradient kernels. #[derive(Shader)] @@ -20,6 +20,8 @@ pub struct Ppo { pub actor_grad: GpuPpoActorGrad, /// Clipped value-loss gradient. pub value_grad: GpuPpoValueGrad, + /// On-device staging of the PPO minibatch from raw rollout observations. + pub stage_batch: GpuPpoStageBatch, } impl Ppo { @@ -48,7 +50,8 @@ impl Ppo { let mut g_mean = g_mean.as_tensor_mut(); let mut g_logstd = g_logstd.as_tensor_mut(); - let num_threads = adv.len() as u32; // one thread per sample column + // One thread per sample column, clamped to the kernel's stride. + let num_threads = (adv.len() as u32).min(vortx_shaders::ml::ppo::MAX_NUM_THREADS); let mut buf_g_mean = g_mean.buffer_mut(); let mut buf_g_logstd = g_logstd.buffer_mut(); @@ -83,7 +86,7 @@ impl Ppo { let ret = ret.as_tensor_ref(); let mut g_v = g_v.as_tensor_mut(); - let num_threads = v_pred.len() as u32; + let num_threads = (v_pred.len() as u32).min(vortx_shaders::ml::ppo::MAX_NUM_THREADS); let mut buf_g_v = g_v.buffer_mut(); self.value_grad.call( @@ -96,4 +99,50 @@ impl Ppo { &mut buf_g_v, ) } + + /// Stages (one half of) the PPO batch on device: reads the step-blocked raw + /// rollout observations, applies the signed-perm mirror, the normalizer + /// affine and the ±5 clamp, and writes row-major `[dim x total_cols]` + /// columns starting at `params.col_offset`. + /// + /// `mean` / `inv_std` / `perm` / `sign` are all `[dim]`; pass identity + /// tables in `perm` / `sign` for the un-mirrored half. See + /// [`gpu_ppo_stage_batch`](vortx_shaders::ml::ppo::gpu_ppo_stage_batch) for + /// the layout contract. + #[allow(clippy::too_many_arguments)] + pub fn stage_batch( + &self, + pass: &mut GpuPass, + params: impl AsTensorRef, + raw: impl AsTensorRef, + mean: impl AsTensorRef, + inv_std: impl AsTensorRef, + perm: impl AsTensorRef, + sign: impl AsTensorRef, + mut out: impl AsTensorMut, + cols: u32, + dim: u32, + ) -> Result<(), GpuBackendError> { + let params = params.as_tensor_ref(); + let raw = raw.as_tensor_ref(); + let mean = mean.as_tensor_ref(); + let inv_std = inv_std.as_tensor_ref(); + let perm = perm.as_tensor_ref(); + let sign = sign.as_tensor_ref(); + let mut out = out.as_tensor_mut(); + + let mut buf_out = out.buffer_mut(); + + self.stage_batch.call( + pass, + [cols, dim, 1], + ¶ms.buffer(), + &raw.buffer(), + &mean.buffer(), + &inv_std.buffer(), + &perm.buffer(), + &sign.buffer(), + &mut buf_out, + ) + } } diff --git a/src/ml/unary.rs b/src/ml/unary.rs index 2e0231e..ae57b0f 100644 --- a/src/ml/unary.rs +++ b/src/ml/unary.rs @@ -342,7 +342,7 @@ impl Unary { args: Option<&Tensor>, ) -> Result<(), GpuBackendError> { let mut src = src.as_tensor_mut(); - let len = src.len() as u32; + let len = (src.len() as u32).min(vortx_shaders::ml::unary::MAX_NUM_THREADS); assert_eq!( op.has_args(), @@ -623,7 +623,7 @@ impl Unary { ) -> Result<(), GpuBackendError> { let mut dest = dest.as_tensor_mut(); let src = src.as_tensor_ref(); - let len = dest.len() as u32; + let len = (dest.len() as u32).min(vortx_shaders::ml::unary::MAX_NUM_THREADS); assert_eq!( op.has_args(), diff --git a/vortx-shaders/src/linalg/contiguous.rs b/vortx-shaders/src/linalg/contiguous.rs index cd64375..e8bcc08 100644 --- a/vortx-shaders/src/linalg/contiguous.rs +++ b/vortx-shaders/src/linalg/contiguous.rs @@ -12,7 +12,11 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 128; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Convert to contiguous row-major layout. #[spirv_bindgen] diff --git a/vortx-shaders/src/linalg/op_assign.rs b/vortx-shaders/src/linalg/op_assign.rs index f589446..92c7044 100644 --- a/vortx-shaders/src/linalg/op_assign.rs +++ b/vortx-shaders/src/linalg/op_assign.rs @@ -12,7 +12,11 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Binary operation offsets. #[repr(C)] diff --git a/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs index cbb1b21..946831e 100644 --- a/vortx-shaders/src/ml/activation.rs +++ b/vortx-shaders/src/ml/activation.rs @@ -11,7 +11,11 @@ use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 256; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Backward of tanh, in place: `g *= 1 - y*y`, where `y = tanh(x)` is the forward output. /// diff --git a/vortx-shaders/src/ml/optim.rs b/vortx-shaders/src/ml/optim.rs index 4cbeec6..1446d3b 100644 --- a/vortx-shaders/src/ml/optim.rs +++ b/vortx-shaders/src/ml/optim.rs @@ -12,7 +12,11 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Scalar parameters for one Adam step (uniform buffer; padded to 32 bytes). #[repr(C)] diff --git a/vortx-shaders/src/ml/ppo.rs b/vortx-shaders/src/ml/ppo.rs index 5b943ff..f69e53d 100644 --- a/vortx-shaders/src/ml/ppo.rs +++ b/vortx-shaders/src/ml/ppo.rs @@ -18,7 +18,11 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Scalar parameters for the actor PPO gradient (uniform buffer; 32 bytes). #[repr(C)] @@ -160,3 +164,80 @@ pub fn gpu_ppo_value_grad( *g_v.at_mut(m) = params.value_coef * dv * scale; } } + +/// Scalar parameters for the PPO batch staging (uniform buffer; 32 bytes). +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +pub struct PpoStageParams { + /// Observation dimensionality (rows). + pub dim: u32, + /// Environments per rollout step. + pub n: u32, + /// Rollout steps `T`; the raw buffer is step-blocked `[T][dim][n]`. + pub steps: u32, + /// Total batch columns of `out` (its row stride). + pub total_cols: u32, + /// First output column this dispatch writes (the mirrored or the original + /// half of the batch). + pub col_offset: u32, + /// 0 = batch mode: the columns cover all `T·n` samples, env-major. + /// Otherwise single-step mode, staging only rollout step `step_select - 1`, + /// so the columns are the `n` envs (the per-step policy-input staging). + pub step_select: u32, + pub pad1: u32, + pub pad2: u32, +} + +/// Builds (one half of) the `[dim x total_cols]` row-major PPO batch straight +/// from the step-blocked raw rollout observations, applying the signed-perm +/// mirror, the normalizer affine and the ±5 clamp in one dispatch. +/// +/// The mirror arrives as an explicit signed permutation (`perm`/`sign`, with +/// identity tables for the un-mirrored half) rather than re-derived index +/// maths, so it cannot drift from the caller's definition. Normalization +/// happens here, not before: the mirror is defined on raw observations +/// (`normalize ∘ mirror`) and the clamp is lossy, so a mirror taken from +/// already-normalized values is wrong for every saturated feature. +/// +/// Batch columns are env-major (`col = e·T + t`, the trainer's sample flatten +/// order) while the raw buffer is step-blocked, hence the +/// `raw[(t·dim + perm[d])·n + e]` gather. Dispatch `[cols, dim, 1]` threads. +#[spirv_bindgen] +#[spirv(compute(threads(256, 1, 1)))] +pub fn gpu_ppo_stage_batch( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &PpoStageParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] raw: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mean: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] inv_std: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] perm: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sign: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] out: &mut [f32], +) { + // 2-D: `x` walks the batch columns (contiguous within an `out` row, so the + // writes coalesce), `y` walks the observation dimensions. Flattening this + // to 1-D would put `cols · dim` threads on one axis, which overruns the + // 65535-workgroup limit at realistic batch sizes. + let cols = if params.step_select != 0 { + params.n + } else { + params.steps * params.n + }; + let x = invocation_id.x; + let d = invocation_id.y; + if x >= cols || d >= params.dim { + return; + } + let (t, e) = if params.step_select != 0 { + (params.step_select - 1, x) + } else { + (x % params.steps, x / params.steps) + }; + let src_d = perm.read(d as usize); + let v = raw.read(((t * params.dim + src_d) * params.n + e) as usize) * sign.read(d as usize); + let v = ((v - mean.read(d as usize)) * inv_std.read(d as usize)) + .max(-5.0) + .min(5.0); + out.write((d * params.total_cols + params.col_offset + x) as usize, v); +} diff --git a/vortx-shaders/src/ml/unary.rs b/vortx-shaders/src/ml/unary.rs index 4f091ef..1ea63aa 100644 --- a/vortx-shaders/src/ml/unary.rs +++ b/vortx-shaders/src/ml/unary.rs @@ -12,7 +12,11 @@ use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::num_traits::Float; const WORKGROUP_SIZE: u32 = 64; -const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +/// Largest dispatch these kernels accept. They stride by exactly this, +/// so a host dispatch must be clamped to it: dispatching fewer threads +/// would leave a gap in the stride, and more would overrun the +/// 65535-workgroup limit. +pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; // // GELU constants const GELU_COEF_A: f32 = 0.044715; From eb1d7d08c71dd1fb238274363fb732e5a2fb39be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:07:00 +0200 Subject: [PATCH 12/13] chore(ml): comment cleanups --- src/ml/batched_multiquery_attention.rs | 5 +-- src/ml/gemv_quant.rs | 27 ++++--------- src/ml/rope.rs | 15 +++---- vortx-shaders/src/linalg/contiguous.rs | 7 ++-- vortx-shaders/src/linalg/op_assign.rs | 7 ++-- vortx-shaders/src/ml/activation.rs | 7 ++-- vortx-shaders/src/ml/conv2d.rs | 10 ++--- vortx-shaders/src/ml/fused_attention.rs | 54 +++++-------------------- vortx-shaders/src/ml/gemv_quant_q8_k.rs | 2 +- vortx-shaders/src/ml/layernorm.rs | 8 ++-- vortx-shaders/src/ml/mod.rs | 8 ---- vortx-shaders/src/ml/optim.rs | 7 ++-- vortx-shaders/src/ml/pool2d.rs | 10 ++--- vortx-shaders/src/ml/ppo.rs | 30 ++++++-------- vortx-shaders/src/ml/reduce_axis.rs | 6 +-- vortx-shaders/src/ml/softmax.rs | 4 +- vortx-shaders/src/ml/unary.rs | 12 +++--- 17 files changed, 76 insertions(+), 143 deletions(-) diff --git a/src/ml/batched_multiquery_attention.rs b/src/ml/batched_multiquery_attention.rs index d02ba14..f2e9e72 100644 --- a/src/ml/batched_multiquery_attention.rs +++ b/src/ml/batched_multiquery_attention.rs @@ -6,7 +6,7 @@ use nalgebra::{DMatrix, DVector}; use vortx_shaders::ml::AttentionParams; #[derive(Shader)] -/// Fused attention shader - combines Q*K^T, scale, mask, softmax, and *V into one kernel. +/// Fused attention shader: Q*K^T, scale, mask, softmax and *V in one kernel. pub struct FusedAttention { pub fused_attention: vortx_shaders::ml::FusedAttention, pub fused_attention_online: vortx_shaders::ml::FusedAttentionOnline, @@ -15,9 +15,6 @@ pub struct FusedAttention { impl FusedAttention { /// Launch the fused attention kernel. - /// - /// This replaces the 4-dispatch attention (matmul -> mask -> softmax -> matmul) - /// with a single fused kernel dispatch. pub fn launch( &self, _backend: &GpuBackend, diff --git a/src/ml/gemv_quant.rs b/src/ml/gemv_quant.rs index 647bc49..dd87a72 100644 --- a/src/ml/gemv_quant.rs +++ b/src/ml/gemv_quant.rs @@ -72,8 +72,8 @@ impl QuantizedValue for GpuBlockQ6Kx2 { const DEQUANTIZED_LEN: usize = 512; } -// SAFETY: These impls are safe, they don't exist in bytemuck because they don't -// provide impls for non-power-of-two largeish arrays. +// SAFETY: sound; bytemuck just doesn't provide impls for largeish +// non-power-of-two arrays. unsafe impl bytemuck::Zeroable for GpuBlockQ6Kx2 {} unsafe impl bytemuck::Pod for GpuBlockQ6Kx2 {} @@ -245,10 +245,9 @@ impl GemvQuant { | GpuQuantTensor::Q5K(_) | GpuQuantTensor::Q4K(_) => out.layout().f32_to_vec4(), // Non-optimized shaders (Q4_1, Q5_0, Q5_1, Q8K) index output by - // global_invocation_id and write Vec4::splat per row. On WebGPU the - // OOB writes are clamped; on CUDA they corrupt memory. These shader - // paths are currently broken for CUDA and only work by accident on - // WebGPU. They are rarely used in practice (most models use Q4K/Q5K). + // global_invocation_id and write Vec4::splat per row, so they write + // out of bounds: WebGPU clamps that, CUDA corrupts memory. Rarely + // used in practice (most models are Q4K/Q5K). _ => out.layout(), }; @@ -279,9 +278,9 @@ impl GemvQuant { let grid = DispatchGrid::Grid([launch, 1, 1]); - // Dispatch to the appropriate kernel based on quantization type. - // Each variant has its own generated args type, but all share the same field names - // (shape_m, out, m, v) since the shader functions have the same signature. + // Each quantization variant has its own generated args type, but they + // share the field names (shape_m, out, m, v) since the shader signatures + // match. macro_rules! dispatch_gemv { ($kernel:expr, $tensor:expr) => {{ #[cfg(not(feature = "push_constants"))] @@ -495,9 +494,7 @@ mod test { result } - // ========================================================================= // Q8_0 - // ========================================================================= /// Dequantize a GpuBlockQ8_0x2 slice into flat f32s. Each GPU block = 2 CPU blocks = 64 f32s. fn dequantize_q8_0x2(blocks: &[GpuBlockQ8_0x2]) -> Vec { @@ -567,9 +564,7 @@ mod test { test_gemv_q8_0_generic(&backend).await; } - // ========================================================================= // Q4_0 - // ========================================================================= fn dequantize_q4_0x2(blocks: &[GpuBlockQ4_0x2]) -> Vec { let cpu_blocks: &[BlockQ4_0] = bytemuck::cast_slice(blocks); @@ -638,9 +633,7 @@ mod test { test_gemv_q4_0_generic(&backend).await; } - // ========================================================================= // Q4K - // ========================================================================= fn dequantize_q4k(blocks: &[GpuBlockQ4K]) -> Vec { blocks.iter().flat_map(|b| b.dequantize()).collect() @@ -707,9 +700,7 @@ mod test { test_gemv_q4k_generic(&backend).await; } - // ========================================================================= // Q5K - // ========================================================================= fn dequantize_q5k(blocks: &[GpuBlockQ5K]) -> Vec { blocks.iter().flat_map(|b| b.dequantize()).collect() @@ -776,9 +767,7 @@ mod test { test_gemv_q5k_generic(&backend).await; } - // ========================================================================= // Q6K (optimized path, uses shared memory + workgroup reduction) - // ========================================================================= fn dequantize_q6kx2(blocks: &[GpuBlockQ6Kx2]) -> Vec { let cpu_blocks: &[BlockQ6K] = bytemuck::cast_slice(blocks); diff --git a/src/ml/rope.rs b/src/ml/rope.rs index ab8312b..876f348 100644 --- a/src/ml/rope.rs +++ b/src/ml/rope.rs @@ -108,12 +108,10 @@ impl RoPE { // for each head. So we need to transform `i` into the corresponding index within // the head. let head_dim = (i % head_size) as f32; - // Not that the formulae from the video linked above would be: + // The formula from the video linked above would be: // 10000.0.powf(-2.0 * ((i / 2) as f32 - 1.0) / dim as f32) - // Although in the paper shown in the video, their index is 1-based which his why thy - // have to subtract 1.0 whereas we don't need to.The `i / 2` and multiplication by 2.0 - // are both accounted for by stepping only on even values for `i`. - // Therefore, the formulae below is equivalent to the RoPE paper's formulae. + // Its index is 1-based, hence the `- 1.0`; the `i / 2` and the factor 2.0 + // are both accounted for by stepping `i` on even values only. let theta = 10000.0_f32.powf(-head_dim / head_size as f32); let m_theta = pos as f32 * theta; let rot = Rotation2::new(m_theta); @@ -122,10 +120,9 @@ impl RoPE { let mut out_q = q.fixed_rows_mut::<2>(i); out_q.copy_from(&(rot * qi)); - // When i >= kv_dim, we are done rotating all the elements from the keys. That's - // because there are less key heads than query heads, but each key head sub-vector has - // the same dimension as the query head (they loose dimension when multiplied with the - // key weight matrices). + // Past kv_dim, every key element has been rotated: there are fewer key + // heads than query heads, though each key head sub-vector has the same + // dimension as a query head. if i < kv_dim { let ki = vector![k[i], k[i + 1]]; let mut out_k = k.fixed_rows_mut::<2>(i); diff --git a/vortx-shaders/src/linalg/contiguous.rs b/vortx-shaders/src/linalg/contiguous.rs index e8bcc08..6e4c1fc 100644 --- a/vortx-shaders/src/linalg/contiguous.rs +++ b/vortx-shaders/src/linalg/contiguous.rs @@ -12,10 +12,9 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 128; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Convert to contiguous row-major layout. diff --git a/vortx-shaders/src/linalg/op_assign.rs b/vortx-shaders/src/linalg/op_assign.rs index 92c7044..a300e5c 100644 --- a/vortx-shaders/src/linalg/op_assign.rs +++ b/vortx-shaders/src/linalg/op_assign.rs @@ -12,10 +12,9 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Binary operation offsets. diff --git a/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs index 946831e..3fe147a 100644 --- a/vortx-shaders/src/ml/activation.rs +++ b/vortx-shaders/src/ml/activation.rs @@ -11,10 +11,9 @@ use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 256; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Backward of tanh, in place: `g *= 1 - y*y`, where `y = tanh(x)` is the forward output. diff --git a/vortx-shaders/src/ml/conv2d.rs b/vortx-shaders/src/ml/conv2d.rs index 29b0559..2765957 100644 --- a/vortx-shaders/src/ml/conv2d.rs +++ b/vortx-shaders/src/ml/conv2d.rs @@ -30,10 +30,10 @@ use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; -/// Conv2d - compute 2D convolution. +/// Conv2d: 2D convolution. /// -/// This is a straightforward implementation, not optimized for performance. -/// For each output element, iterate over the kernel and compute the convolution. +/// Straightforward and unoptimized: iterate over the kernel for each output +/// element. #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] pub fn conv_2d_nchw( @@ -109,8 +109,8 @@ pub fn conv_2d_nchw( + ih * input_w + iw) as usize; - // Weight index: [oc, ic_local, kh, kw] - // Note: ic_local is used because weight has shape [out_channels, in_channels/groups, kh, kw] + // Weight index: [oc, ic_local, kh, kw], local because the + // weight shape is [out_channels, in_channels/groups, kh, kw]. let i_weight = (oc * in_channels_per_group * kernel_h * kernel_w + ic_local * kernel_h * kernel_w + kh * kernel_w diff --git a/vortx-shaders/src/ml/fused_attention.rs b/vortx-shaders/src/ml/fused_attention.rs index 300fd7e..9c26ca3 100644 --- a/vortx-shaders/src/ml/fused_attention.rs +++ b/vortx-shaders/src/ml/fused_attention.rs @@ -16,7 +16,7 @@ use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(any(target_arch = "spirv", target_arch = "nvptx64"))] use khal_std::num_traits::Float; -/// Workgroup size - should be >= head_size for efficient V accumulation. +/// Workgroup size; should be >= head_size for efficient V accumulation. #[cfg(feature = "subgroup_ops")] const WORKGROUP_SIZE: usize = 32; #[cfg(not(feature = "subgroup_ops"))] @@ -26,7 +26,7 @@ const WORKGROUP_SIZE: usize = 128; /// For longer sequences, we use online softmax to avoid storing all scores. const MAX_SEQ_LEN: usize = 2048; -/// Block size for Flash Attention - number of KV tokens processed per iteration. +/// Block size for Flash Attention: KV tokens processed per iteration. /// Must be tuned to fit shared memory: kv_tile uses BLOCK_KV * WORKGROUP_SIZE * 4 bytes. /// With BLOCK_KV=32 and WORKGROUP_SIZE=128: 32 * 128 * 4 = 16KB for kv_tile alone. const BLOCK_KV: usize = 32; @@ -92,9 +92,7 @@ pub fn fused_attention( let q_base = (head_idx * params.head_size) as usize; let kv_base = (kv_head * params.head_size) as usize; - // ========================================================================== // Phase 1: Compute Q · K^T for all positions, scale, and find max - // ========================================================================== // Each thread computes dot products for a subset of positions // Max iterations: ceil(MAX_SEQ_LEN / WORKGROUP_SIZE) = ceil(2048/128) = 16 @@ -153,9 +151,7 @@ pub fn fused_attention( khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ========================================================================== // Phase 2: Compute exp(score - max) and sum - // ========================================================================== let the_max = *max_score; let mut my_sum = 0.0f32; @@ -200,9 +196,7 @@ pub fn fused_attention( khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ========================================================================== // Phase 3: Normalize attention weights (divide by sum) - // ========================================================================== let the_sum = *sum_exp; let inv_sum = 1.0 / the_sum; @@ -217,9 +211,7 @@ pub fn fused_attention( khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ========================================================================== // Phase 4: Compute weighted sum of values - // ========================================================================== // Each thread computes output for a subset of head dimensions let out_base = (head_idx * params.head_size) as usize; @@ -242,8 +234,8 @@ pub fn fused_attention( /// Fused attention with online softmax for long sequences. /// -/// This variant uses online softmax to avoid storing all attention scores, -/// making it memory-efficient for arbitrarily long sequences. +/// Online softmax avoids storing all attention scores, so this stays +/// memory-efficient for arbitrarily long sequences. #[spirv_bindgen] #[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] #[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] @@ -289,7 +281,7 @@ pub fn fused_attention_online( // Compute Q · K[t] let k_base = t * (n_kv_heads * params.head_size) as usize + kv_base; - // Collaborative dot product - each thread handles part of the dimensions + // Collaborative dot product: each thread handles part of the dimensions. let mut partial_dot = 0.0f32; if tid < head_size { let q_val = q.read(q_base + tid); @@ -360,17 +352,12 @@ pub fn fused_attention_online( /// Flash Attention kernel with tiled/block-wise processing. /// -/// This kernel processes KV in blocks of BLOCK_KV tokens, using online softmax -/// to maintain O(1) memory per softmax row. This is ~100x more efficient than -/// the fused_attention_online kernel which processes one token at a time. +/// Processes the KV cache in blocks of BLOCK_KV tokens, with online softmax +/// rescaled between blocks, so softmax needs O(1) memory per row and the +/// weighted V values accumulate incrementally. /// -/// Workgroup layout: [WORKGROUP_SIZE, 1, 1] -/// Dispatch: [n_heads, 1, 1] workgroups -/// -/// Each workgroup computes attention for one query head using Flash Attention: -/// - Processes KV cache in blocks of BLOCK_KV tokens -/// - Uses online softmax with rescaling between blocks -/// - Accumulates weighted V values incrementally +/// Workgroup layout: [WORKGROUP_SIZE, 1, 1]; dispatch [n_heads, 1, 1] +/// workgroups, one query head each. #[spirv_bindgen] #[cfg_attr(feature = "subgroup_ops", spirv(compute(threads(32, 1, 1))))] #[cfg_attr(not(feature = "subgroup_ops"), spirv(compute(threads(128, 1, 1))))] @@ -407,9 +394,7 @@ pub fn flash_attention( let kv_stride = (n_kv_heads * params.head_size) as usize; let out_base = (head_idx * params.head_size) as usize; - // ========================================================================== // Phase 0: Load Q into shared memory and initialize accumulators - // ========================================================================== if tid < head_size { q_shared.write(tid, q.read(q_base + tid)); out_accum.write(tid, 0.0); @@ -420,10 +405,7 @@ pub fn flash_attention( } khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ========================================================================== - // Main loop: Process KV in blocks of BLOCK_KV tokens - // ========================================================================== - // Maximum number of blocks we might process + // Main loop: process the KV cache in blocks of BLOCK_KV tokens. let num_blocks = seq_len.div_ceil(BLOCK_KV); for block_idx in 0..num_blocks { @@ -435,9 +417,7 @@ pub fn flash_attention( }; let block_len = block_end - block_start; - // ---------------------------------------------------------------------- // Step 1: Load K block into shared memory - // ---------------------------------------------------------------------- // Layout: kv_tile[pos_in_block * head_size + dim] // Max elements = BLOCK_KV * head_size, max iterations = ceil(BLOCK_KV * head_size / WORKGROUP_SIZE) for iter in 0..BLOCK_KV { @@ -452,9 +432,7 @@ pub fn flash_attention( } khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ---------------------------------------------------------------------- // Step 2: Compute Q · K[t] for all positions in block - // ---------------------------------------------------------------------- // Each thread handles at most one position (since BLOCK_KV <= WORKGROUP_SIZE) if tid < block_len { let mut dot = 0.0f32; @@ -475,9 +453,7 @@ pub fn flash_attention( } khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ---------------------------------------------------------------------- // Step 3: Find block max via parallel reduction - // ---------------------------------------------------------------------- let my_max = if tid < block_len { scores.read(tid) } else { @@ -515,9 +491,7 @@ pub fn flash_attention( let block_max = *block_max_shared; - // ---------------------------------------------------------------------- // Step 4: Compute exp(score - block_max) and sum - // ---------------------------------------------------------------------- let my_sum = if tid < block_len { let exp_score = (scores.read(tid) - block_max).exp(); scores.write(tid, exp_score); // Overwrite with exp values @@ -556,9 +530,7 @@ pub fn flash_attention( let block_sum = *block_sum_shared; - // ---------------------------------------------------------------------- // Step 5: Update running statistics with rescaling - // ---------------------------------------------------------------------- let old_max = *running_max; let old_sum = *running_sum; let new_max = old_max.max(block_max); @@ -571,9 +543,7 @@ pub fn flash_attention( } khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // ---------------------------------------------------------------------- // Step 6: Load V block and accumulate weighted values - // ---------------------------------------------------------------------- // Reuse kv_tile for V block for iter in 0..BLOCK_KV { let load_idx = tid + iter * WORKGROUP_SIZE; @@ -606,9 +576,7 @@ pub fn flash_attention( khal_std::sync::workgroup_memory_barrier_with_group_sync(); } - // ========================================================================== // Final: Normalize by running sum and write output - // ========================================================================== let final_sum = *running_sum; if tid < head_size { let val = out_accum.read(tid) / final_sum; diff --git a/vortx-shaders/src/ml/gemv_quant_q8_k.rs b/vortx-shaders/src/ml/gemv_quant_q8_k.rs index 4a89682..ea0ad92 100644 --- a/vortx-shaders/src/ml/gemv_quant_q8_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q8_k.rs @@ -15,7 +15,7 @@ const WORKGROUP_SIZE: u32 = 32; // BlockQ8K structure (repr(C), alignment 4): // - d: f32 (1 u32) // - qs: [i8; 256] (64 u32s) -// - bsums: [i16; 16] (8 u32s, no padding — 260 is already 2-byte aligned) +// - bsums: [i16; 16] (8 u32s, no padding: 260 is already 2-byte aligned) // Total: 73 u32s = 292 bytes const BLOCK_Q8K_SIZE: u32 = 73; diff --git a/vortx-shaders/src/ml/layernorm.rs b/vortx-shaders/src/ml/layernorm.rs index 25517e6..d42bfd7 100644 --- a/vortx-shaders/src/ml/layernorm.rs +++ b/vortx-shaders/src/ml/layernorm.rs @@ -57,7 +57,7 @@ pub fn layernorm_cols( let thread_id = local_id.x as usize; - // Compute the MEAN + // Compute the mean. let data_len = in_shape.h; *workspace.at_mut(thread_id) = 0.0; for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { @@ -92,7 +92,7 @@ pub fn layernorm_cols( khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // Compute the SQUARED NORM + // Compute the squared norm. *workspace.at_mut(thread_id) = 0.0; for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { let val_i = *input.at(in_shape.it(wid.z, wid.y, i, wid.x) as usize) - *the_mean; @@ -164,7 +164,7 @@ pub fn layernorm_rows( let thread_id = local_id.x as usize; - // Compute the MEAN + // Compute the mean. let data_len = in_shape.w; *workspace.at_mut(thread_id) = 0.0; for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { @@ -199,7 +199,7 @@ pub fn layernorm_rows( khal_std::sync::workgroup_memory_barrier_with_group_sync(); - // Compute the SQUARED NORM + // Compute the squared norm. *workspace.at_mut(thread_id) = 0.0; for i in StepRng::new(thread_id as u32..data_len, WORKGROUP_SIZE as u32) { let val_i = *input.at(in_shape.it(wid.z, wid.y, wid.x, i) as usize) - *the_mean; diff --git a/vortx-shaders/src/ml/mod.rs b/vortx-shaders/src/ml/mod.rs index 136ec0c..ff269dc 100644 --- a/vortx-shaders/src/ml/mod.rs +++ b/vortx-shaders/src/ml/mod.rs @@ -1,11 +1,3 @@ -// #![allow(clippy::too_many_arguments)] -// // `spirv_bindgen` generates host-side dispatch code that performs `% workgroup_size`, -// // which triggers this lint when a workgroup dimension is 1. -// #![allow(clippy::modulo_one)] -// #![allow(unexpected_cfgs)] -// // Shader entry points and their constants appear dead on host but are used on GPU. -// #![allow(dead_code, non_snake_case)] - // TODO: keep the modules private? pub mod activation; pub mod batched_multiquery_attention; diff --git a/vortx-shaders/src/ml/optim.rs b/vortx-shaders/src/ml/optim.rs index 1446d3b..263b303 100644 --- a/vortx-shaders/src/ml/optim.rs +++ b/vortx-shaders/src/ml/optim.rs @@ -12,10 +12,9 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Scalar parameters for one Adam step (uniform buffer; padded to 32 bytes). diff --git a/vortx-shaders/src/ml/pool2d.rs b/vortx-shaders/src/ml/pool2d.rs index 4588ffa..c2d3a04 100644 --- a/vortx-shaders/src/ml/pool2d.rs +++ b/vortx-shaders/src/ml/pool2d.rs @@ -25,7 +25,7 @@ use khal_std::macros::{spirv, spirv_bindgen}; const WORKGROUP_SIZE: u32 = 64; const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; -/// MaxPool2d - compute max over a 2D window. +/// MaxPool2d: max over a 2D window. #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] pub fn max_pool_2d( @@ -84,7 +84,7 @@ pub fn max_pool_2d( w_end_unclamped }; - // Initialize max with a very small value (SPIR-V doesn't support infinity literals) + // SPIR-V has no infinity literals, so start from a very small value. let mut max_val = -3.4028235e+38_f32; // Close to f32::MIN // Iterate over the pooling window @@ -110,7 +110,7 @@ pub fn max_pool_2d( } } -/// AvgPool2d - compute average over a 2D window. +/// AvgPool2d: average over a 2D window. #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] pub fn avg_pool_2d( @@ -202,7 +202,7 @@ pub fn avg_pool_2d( } } -/// GlobalAvgPool2d - average over entire spatial dimensions. +/// GlobalAvgPool2d: average over the entire spatial dimensions. /// Input: [N, C, H, W], Output: [N, C, 1, 1] #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] @@ -240,7 +240,7 @@ pub fn global_avg_pool_2d( } } -/// GlobalMaxPool2d - max over entire spatial dimensions. +/// GlobalMaxPool2d: max over the entire spatial dimensions. /// Input: [N, C, H, W], Output: [N, C, 1, 1] #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] diff --git a/vortx-shaders/src/ml/ppo.rs b/vortx-shaders/src/ml/ppo.rs index f69e53d..f72a14b 100644 --- a/vortx-shaders/src/ml/ppo.rs +++ b/vortx-shaders/src/ml/ppo.rs @@ -18,10 +18,9 @@ use khal_std::{ }; const WORKGROUP_SIZE: u32 = 256; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Scalar parameters for the actor PPO gradient (uniform buffer; 32 bytes). @@ -193,16 +192,14 @@ pub struct PpoStageParams { /// from the step-blocked raw rollout observations, applying the signed-perm /// mirror, the normalizer affine and the ±5 clamp in one dispatch. /// -/// The mirror arrives as an explicit signed permutation (`perm`/`sign`, with -/// identity tables for the un-mirrored half) rather than re-derived index -/// maths, so it cannot drift from the caller's definition. Normalization -/// happens here, not before: the mirror is defined on raw observations -/// (`normalize ∘ mirror`) and the clamp is lossy, so a mirror taken from -/// already-normalized values is wrong for every saturated feature. +/// The mirror is an explicit signed permutation (`perm`/`sign`, identity tables +/// for the un-mirrored half). It has to be applied to raw observations, ahead of +/// the lossy clamp: a mirror of already-normalized values is wrong for every +/// saturated feature. /// -/// Batch columns are env-major (`col = e·T + t`, the trainer's sample flatten -/// order) while the raw buffer is step-blocked, hence the -/// `raw[(t·dim + perm[d])·n + e]` gather. Dispatch `[cols, dim, 1]` threads. +/// Batch columns are env-major (`col = e·T + t`) while the raw buffer is +/// step-blocked, hence the `raw[(t·dim + perm[d])·n + e]` gather. Dispatch +/// `[cols, dim, 1]` threads. #[spirv_bindgen] #[spirv(compute(threads(256, 1, 1)))] pub fn gpu_ppo_stage_batch( @@ -215,10 +212,9 @@ pub fn gpu_ppo_stage_batch( #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sign: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] out: &mut [f32], ) { - // 2-D: `x` walks the batch columns (contiguous within an `out` row, so the - // writes coalesce), `y` walks the observation dimensions. Flattening this - // to 1-D would put `cols · dim` threads on one axis, which overruns the - // 65535-workgroup limit at realistic batch sizes. + // `x` walks the batch columns (contiguous within an `out` row, so the writes + // coalesce), `y` the observation dimensions. A flat 1-D dispatch of + // `cols · dim` threads would overrun the 65535-workgroup limit. let cols = if params.step_select != 0 { params.n } else { diff --git a/vortx-shaders/src/ml/reduce_axis.rs b/vortx-shaders/src/ml/reduce_axis.rs index bd93b4a..06b10e5 100644 --- a/vortx-shaders/src/ml/reduce_axis.rs +++ b/vortx-shaders/src/ml/reduce_axis.rs @@ -46,7 +46,7 @@ pub fn reduce_sum_axis( // Decompose linear index in output let id_dest = shape_dest.decompose(thread_id); - // Build source coordinates - start with output coords + // Build the source coordinates, starting from the output coords. let mut id_src = id_dest; // Sum over all elements along the reduce axis @@ -99,7 +99,7 @@ pub fn reduce_mean_axis( // Decompose linear index in output let id_dest = shape_dest.decompose(thread_id); - // Build source coordinates - start with output coords + // Build the source coordinates, starting from the output coords. let mut id_src = id_dest; // Sum over all elements along the reduce axis @@ -152,7 +152,7 @@ pub fn reduce_max_axis( // Decompose linear index in output let id_dest = shape_dest.decompose(thread_id); - // Build source coordinates - start with output coords + // Build the source coordinates, starting from the output coords. let mut id_src = id_dest; // Find max over all elements along the reduce axis diff --git a/vortx-shaders/src/ml/softmax.rs b/vortx-shaders/src/ml/softmax.rs index dc09b45..4ae8940 100644 --- a/vortx-shaders/src/ml/softmax.rs +++ b/vortx-shaders/src/ml/softmax.rs @@ -60,7 +60,7 @@ pub fn softmax( let l = workgroup_id.z; let thread_id = local_id.x as usize; - // Compute the MAX + // Compute the max. let data_len = shape.w; let mut my_max = [-1.0e38f32]; @@ -179,7 +179,7 @@ pub fn log_softmax( let l = workgroup_id.z; let thread_id = local_id.x as usize; - // Compute the MAX + // Compute the max. let data_len = shape.w; let mut my_max = [-1.0e38f32]; diff --git a/vortx-shaders/src/ml/unary.rs b/vortx-shaders/src/ml/unary.rs index 1ea63aa..50085c3 100644 --- a/vortx-shaders/src/ml/unary.rs +++ b/vortx-shaders/src/ml/unary.rs @@ -12,13 +12,12 @@ use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::num_traits::Float; const WORKGROUP_SIZE: u32 = 64; -/// Largest dispatch these kernels accept. They stride by exactly this, -/// so a host dispatch must be clamped to it: dispatching fewer threads -/// would leave a gap in the stride, and more would overrun the -/// 65535-workgroup limit. +/// Largest dispatch these kernels accept: they stride by exactly this, so a host +/// dispatch must be clamped to it (fewer threads leave a gap in the stride, more +/// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; -// // GELU constants +// GELU constants const GELU_COEF_A: f32 = 0.044715; const SQRT_2_OVER_PI: f32 = 0.79788456080286535587989211986876; const GELU_QUICK_COEF: f32 = -1.702; @@ -182,8 +181,7 @@ fn pow_op_fn(x: f32, args: Vec4) -> f32 { x.powf(args.x) } -// Macro-like helper for generating shader entry points -// Since we can't use actual macros in no_std easily, we'll define each manually +// Entry points are written out by hand: macros are awkward in no_std. /// Abs operation. #[spirv_bindgen] From bb38a3de74541b71f0d7173ac74ec6fcfb4efbf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 19:49:06 +0200 Subject: [PATCH 13/13] chore: CI fixes --- Cargo.toml | 1 + src/linalg/contiguous.rs | 3 +- src/ml/gemv_quant.rs | 18 ++-- vortx-shaders/Cargo.toml | 1 + vortx-shaders/src/linalg/contiguous.rs | 7 ++ vortx-shaders/src/linalg/op_assign.rs | 7 ++ vortx-shaders/src/ml/activation.rs | 7 ++ vortx-shaders/src/ml/gemv_quant_q4_0x2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q4_1x2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q4_k.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q5_0x2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q5_1x2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q5_k.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q6_kx2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q8_0x2.rs | 2 +- vortx-shaders/src/ml/gemv_quant_q8_k.rs | 2 +- vortx-shaders/src/ml/mod.rs | 106 +++++++++++++++++----- vortx-shaders/src/ml/optim.rs | 7 ++ vortx-shaders/src/ml/ppo.rs | 7 ++ vortx-shaders/src/ml/unary.rs | 7 ++ 20 files changed, 149 insertions(+), 40 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2caeccb..69d852d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,6 +46,7 @@ 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" diff --git a/src/linalg/contiguous.rs b/src/linalg/contiguous.rs index af716b5..978a158 100644 --- a/src/linalg/contiguous.rs +++ b/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).min(vortx_shaders::linalg::contiguous::MAX_NUM_THREADS); + 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"))] diff --git a/src/ml/gemv_quant.rs b/src/ml/gemv_quant.rs index dd87a72..9cfca65 100644 --- a/src/ml/gemv_quant.rs +++ b/src/ml/gemv_quant.rs @@ -132,55 +132,55 @@ impl GemvQuant { #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ8_0x2 { - pub gemv: vortx_shaders::ml::gemv_quant_q8_0x2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q8_0x2::GemvQ80x2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ5_0x2 { - pub gemv: vortx_shaders::ml::gemv_quant_q5_0x2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q5_0x2::GemvQ50x2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ5_1x2 { - pub gemv: vortx_shaders::ml::gemv_quant_q5_1x2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q5_1x2::GemvQ51x2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ4_0x2 { - pub gemv: vortx_shaders::ml::gemv_quant_q4_0x2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q4_0x2::GemvQ40x2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ4_1x2 { - pub gemv: vortx_shaders::ml::gemv_quant_q4_1x2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q4_1x2::GemvQ41x2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ8K { - pub gemv: vortx_shaders::ml::gemv_quant_q8_k::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q8_k::GemvQ8K, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ6Kx2 { - pub gemv: vortx_shaders::ml::gemv_quant_q6_kx2::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q6_kx2::GemvQ6Kx2, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ5K { - pub gemv: vortx_shaders::ml::gemv_quant_q5_k::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q5_k::GemvQ5K, } #[derive(Shader)] /// Shader for computing the product of a matrix and a vector. pub struct GemvQ4K { - pub gemv: vortx_shaders::ml::gemv_quant_q4_k::Gemv, + pub gemv: vortx_shaders::ml::gemv_quant_q4_k::GemvQ4K, } impl GemvQuant { diff --git a/vortx-shaders/Cargo.toml b/vortx-shaders/Cargo.toml index 26a55ba..4891264 100644 --- a/vortx-shaders/Cargo.toml +++ b/vortx-shaders/Cargo.toml @@ -37,3 +37,4 @@ khal-std = { workspace = true } [target.'cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))'.dependencies] khal = { workspace = true } bytemuck = { version = "1", features = ["derive"] } +static_assertions = "1" diff --git a/vortx-shaders/src/linalg/contiguous.rs b/vortx-shaders/src/linalg/contiguous.rs index 6e4c1fc..c57a47a 100644 --- a/vortx-shaders/src/linalg/contiguous.rs +++ b/vortx-shaders/src/linalg/contiguous.rs @@ -17,6 +17,13 @@ const WORKGROUP_SIZE: u32 = 128; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + /// Convert to contiguous row-major layout. #[spirv_bindgen] #[spirv(compute(threads(128, 1, 1)))] diff --git a/vortx-shaders/src/linalg/op_assign.rs b/vortx-shaders/src/linalg/op_assign.rs index a300e5c..8a17f69 100644 --- a/vortx-shaders/src/linalg/op_assign.rs +++ b/vortx-shaders/src/linalg/op_assign.rs @@ -17,6 +17,13 @@ const WORKGROUP_SIZE: u32 = 256; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + /// Binary operation offsets. #[repr(C)] #[derive(Clone, Copy)] diff --git a/vortx-shaders/src/ml/activation.rs b/vortx-shaders/src/ml/activation.rs index 3fe147a..b30ebcb 100644 --- a/vortx-shaders/src/ml/activation.rs +++ b/vortx-shaders/src/ml/activation.rs @@ -16,6 +16,13 @@ const WORKGROUP_SIZE: u32 = 256; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + /// Backward of tanh, in place: `g *= 1 - y*y`, where `y = tanh(x)` is the forward output. /// /// `g` and `y` must have the same shape. diff --git a/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs index d21b800..727b663 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_0x2.rs @@ -44,7 +44,7 @@ fn dequantize_part(data: u32, scale: f32) -> [Vec4; 2] { #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q4_0x2( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], diff --git a/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs index 0de5339..96a030f 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_1x2.rs @@ -57,7 +57,7 @@ fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] -pub fn gemv( +pub fn gemv_q4_1x2( #[spirv(global_invocation_id)] invocation_id: UVec3, #[cfg(feature = "push_constants")] #[spirv(push_constant)] diff --git a/vortx-shaders/src/ml/gemv_quant_q4_k.rs b/vortx-shaders/src/ml/gemv_quant_q4_k.rs index 731b797..95f65d6 100644 --- a/vortx-shaders/src/ml/gemv_quant_q4_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q4_k.rs @@ -94,7 +94,7 @@ fn dequantize_q4_k_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 2] { #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q4_k( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], diff --git a/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs index 695f71e..c9cf6fd 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_0x2.rs @@ -70,7 +70,7 @@ fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] -pub fn gemv( +pub fn gemv_q5_0x2( #[spirv(global_invocation_id)] invocation_id: UVec3, #[cfg(feature = "push_constants")] #[spirv(push_constant)] diff --git a/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs index b15f8f6..6452072 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_1x2.rs @@ -70,7 +70,7 @@ fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 16] { #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] -pub fn gemv( +pub fn gemv_q5_1x2( #[spirv(global_invocation_id)] invocation_id: UVec3, #[cfg(feature = "push_constants")] #[spirv(push_constant)] diff --git a/vortx-shaders/src/ml/gemv_quant_q5_k.rs b/vortx-shaders/src/ml/gemv_quant_q5_k.rs index a736077..32994af 100644 --- a/vortx-shaders/src/ml/gemv_quant_q5_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q5_k.rs @@ -108,7 +108,7 @@ fn dequantize_q5_k_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 2] { #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q5_k( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], diff --git a/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs index 73f204c..8c70fde 100644 --- a/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q6_kx2.rs @@ -141,7 +141,7 @@ fn dequantize_q6_kx2_workgroup(m: &[u32], block_id: u32, k: u32) -> [Vec4; 4] { #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q6_kx2( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], diff --git a/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs index 13383b3..4b96a4a 100644 --- a/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs +++ b/vortx-shaders/src/ml/gemv_quant_q8_0x2.rs @@ -24,7 +24,7 @@ fn reduce_sum(index: usize, stride: usize, sketch: &mut [Vec4; WORKGROUP_SIZE]) #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q8_0x2( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup)] sketch: &mut [Vec4; WORKGROUP_SIZE], diff --git a/vortx-shaders/src/ml/gemv_quant_q8_k.rs b/vortx-shaders/src/ml/gemv_quant_q8_k.rs index ea0ad92..b0162e0 100644 --- a/vortx-shaders/src/ml/gemv_quant_q8_k.rs +++ b/vortx-shaders/src/ml/gemv_quant_q8_k.rs @@ -38,7 +38,7 @@ fn dequantize_block(data: &[u32], base: usize) -> [Vec4; 64] { #[spirv_bindgen] #[spirv(compute(threads(32, 1, 1)))] -pub fn gemv( +pub fn gemv_q8_k( #[spirv(global_invocation_id)] invocation_id: UVec3, #[cfg(feature = "push_constants")] #[spirv(push_constant)] diff --git a/vortx-shaders/src/ml/mod.rs b/vortx-shaders/src/ml/mod.rs index ff269dc..5758022 100644 --- a/vortx-shaders/src/ml/mod.rs +++ b/vortx-shaders/src/ml/mod.rs @@ -1,3 +1,8 @@ +//! Machine-learning kernels for shaders. + +// Shader entry points and their constants look dead on the host, but are used on GPU. +#![allow(dead_code, non_snake_case)] + // TODO: keep the modules private? pub mod activation; pub mod batched_multiquery_attention; @@ -30,24 +35,83 @@ pub mod softmax; pub mod unary; pub mod win_part; -pub use activation::*; -pub use batched_multiquery_attention::*; -pub use concat::*; -pub use conv2d::*; -pub use conv_transpose_2d::*; -pub use fused_attention::*; -pub use gather::*; -pub use get_rel_pos::*; -pub use im2col::*; -pub use layernorm::*; -pub use optim::*; -pub use pool2d::*; -pub use ppo::*; -pub use reduce_axis::*; -pub use rms_norm::*; -pub use rope::*; -pub use select::*; -pub use silu::*; -pub use softmax::*; -pub use unary::*; -pub use win_part::*; +// Parameter and configuration structs, shared by the host and the GPU. +pub use batched_multiquery_attention::AttentionParams; +pub use im2col::Im2ColParams; +pub use optim::AdamParams; +pub use ppo::{PpoActorParams, PpoStageParams, PpoValueParams}; +pub use rms_norm::RmsNormConfig; +pub use rope::RoPEConfig; + +// Generated ShaderArgs structs (host only). The per-module `MAX_NUM_THREADS` +// constants are not re-exported: callers reach those through the module path. +#[cfg(not(target_arch_is_gpu))] +pub use activation::{GpuEluBackward, GpuTanhBackward}; +#[cfg(not(target_arch_is_gpu))] +pub use batched_multiquery_attention::MultMaskAttn; +#[cfg(not(target_arch_is_gpu))] +pub use concat::ConcatCopy; +#[cfg(not(target_arch_is_gpu))] +pub use conv2d::Conv2dNchw; +#[cfg(not(target_arch_is_gpu))] +pub use conv_transpose_2d::{ + ConvTranspose2d, ConvTranspose2dRef, InitDest, InitSrcA, InitSrcB, InitWdata, +}; +#[cfg(not(target_arch_is_gpu))] +pub use fused_attention::{FlashAttention, FusedAttention, FusedAttentionOnline}; +#[cfg(not(target_arch_is_gpu))] +pub use gather::Gather; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q4_0x2::GemvQ40x2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q4_1x2::GemvQ41x2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q4_k::GemvQ4K; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q5_0x2::GemvQ50x2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q5_1x2::GemvQ51x2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q5_k::GemvQ5K; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q6_kx2::GemvQ6Kx2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q8_0x2::GemvQ80x2; +#[cfg(not(target_arch_is_gpu))] +pub use gemv_quant_q8_k::GemvQ8K; +#[cfg(not(target_arch_is_gpu))] +pub use get_rel_pos::{AddRelPosPhaseA, AddRelPosPhaseB, GetRelPos}; +#[cfg(not(target_arch_is_gpu))] +pub use im2col::Im2col; +#[cfg(not(target_arch_is_gpu))] +pub use layernorm::{LayernormCols, LayernormRows}; +#[cfg(not(target_arch_is_gpu))] +pub use optim::GpuAdam; +#[cfg(not(target_arch_is_gpu))] +pub use pool2d::{AvgPool2d, GlobalAvgPool2d, GlobalMaxPool2d, MaxPool2d}; +#[cfg(not(target_arch_is_gpu))] +pub use ppo::{GpuPpoActorGrad, GpuPpoStageBatch, GpuPpoValueGrad}; +#[cfg(not(target_arch_is_gpu))] +pub use reduce_axis::{ReduceMaxAxis, ReduceMeanAxis, ReduceMinAxis, ReduceSumAxis}; +#[cfg(not(target_arch_is_gpu))] +pub use rms_norm::RmsNorm; +#[cfg(not(target_arch_is_gpu))] +pub use rope::{Rope, RopeNeox}; +#[cfg(not(target_arch_is_gpu))] +pub use select::Select; +#[cfg(not(target_arch_is_gpu))] +pub use silu::Silu; +#[cfg(not(target_arch_is_gpu))] +pub use softmax::{LogSoftmax, Softmax}; +#[cfg(not(target_arch_is_gpu))] +pub use unary::{ + AbsInplace, AbsOp, AddScalarInplace, AddScalarOp, ClampInplace, ClampOp, CosInplace, CosOp, + EluInplace, EluOp, ErfInplace, ErfOp, ExpInplace, ExpOp, GeluInplace, GeluOp, GeluQuickInplace, + GeluQuickOp, HardSigmoidInplace, HardSigmoidOp, LeakyReluInplace, LeakyReluOp, LogInplace, + LogOp, NegInplace, NegOp, PowInplace, PowOp, ReciprocalInplace, ReciprocalOp, ReluInplace, + ReluOp, ScaleInplace, ScaleOp, SgnInplace, SgnOp, SigmoidInplace, SigmoidOp, SiluInplace, + SiluOp, SinInplace, SinOp, SqrInplace, SqrOp, SqrtInplace, SqrtOp, StepInplace, StepOp, + TanhInplace, TanhOp, +}; +#[cfg(not(target_arch_is_gpu))] +pub use win_part::{WinPart, WinUnpart}; diff --git a/vortx-shaders/src/ml/optim.rs b/vortx-shaders/src/ml/optim.rs index 263b303..375ff7c 100644 --- a/vortx-shaders/src/ml/optim.rs +++ b/vortx-shaders/src/ml/optim.rs @@ -17,6 +17,13 @@ const WORKGROUP_SIZE: u32 = 256; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + /// Scalar parameters for one Adam step (uniform buffer; padded to 32 bytes). #[repr(C)] #[derive(Clone, Copy)] diff --git a/vortx-shaders/src/ml/ppo.rs b/vortx-shaders/src/ml/ppo.rs index f72a14b..6d70d1e 100644 --- a/vortx-shaders/src/ml/ppo.rs +++ b/vortx-shaders/src/ml/ppo.rs @@ -23,6 +23,13 @@ const WORKGROUP_SIZE: u32 = 256; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + /// Scalar parameters for the actor PPO gradient (uniform buffer; 32 bytes). #[repr(C)] #[derive(Clone, Copy)] diff --git a/vortx-shaders/src/ml/unary.rs b/vortx-shaders/src/ml/unary.rs index 50085c3..0b6f033 100644 --- a/vortx-shaders/src/ml/unary.rs +++ b/vortx-shaders/src/ml/unary.rs @@ -17,6 +17,13 @@ const WORKGROUP_SIZE: u32 = 64; /// overruns the 65535-workgroup limit). pub const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; +// Guards `WORKGROUP_SIZE` against the `threads(...)` attribute it duplicates. +#[cfg(not(target_arch_is_gpu))] +static_assertions::const_assert_eq!( + WORKGROUP_SIZE, + as khal::shader::ShaderArgsType>::WORKGROUP_SIZE[0] +); + // GELU constants const GELU_COEF_A: f32 = 0.044715; const SQRT_2_OVER_PI: f32 = 0.79788456080286535587989211986876;