More cleanup.

Fix NaN errors for Gelu in Metal

Cargo.toml

@@ -61,7 +61,7 @@ tracing-subscriber = "0.3.7"
 wav = "1.0.0"
 yoke = { version = "0.7.2", features = ["derive"] }
 zip = { version = "0.6.6", default-features = false }
-metal = { version = "0.27.1", features = ["mps"], package="candle-metal" }
+metal = { version = "0.27.0", features = ["mps"]}
 
 [profile.release-with-debug]
 inherits = "release"

candle-core/examples/basics.rs

@@ -5,43 +5,13 @@ extern crate intel_mkl_src;
 extern crate accelerate_src;
 
 use anyhow::Result;
-use candle::{CpuStorage, Device, Layout, Shape, Tensor};
-use candle_core as candle;
-
-struct ArgSort;
-impl candle::CustomOp1 for ArgSort {
-    fn name(&self) -> &'static str {
-        "arg-sort"
-    }
-
-    fn cpu_fwd(
-        &self,
-        storage: &CpuStorage,
-        layout: &Layout,
-    ) -> candle::Result<(CpuStorage, Shape)> {
-        if layout.shape().rank() != 1 {
-            candle::bail!(
-                "input should have a single dimension, got {:?}",
-                layout.shape()
-            )
-        }
-        let slice = storage.as_slice::<f32>()?;
-        let src = match layout.contiguous_offsets() {
-            None => candle::bail!("input has to be contiguous"),
-            Some((o1, o2)) => &slice[o1..o2],
-        };
-        let mut dst = (0..src.len() as u32).collect::<Vec<u32>>();
-        dst.sort_by(|&i, &j| src[i as usize].total_cmp(&src[j as usize]));
-        let storage = candle::WithDType::to_cpu_storage_owned(dst);
-        Ok((storage, layout.shape().clone()))
-    }
-}
+use candle_core::{Device, Tensor};
 
 fn main() -> Result<()> {
-    let a = Tensor::new(&[0.0f32, 1.0, 3.0, 2.0, -12.0, 4.0, 3.5], &Device::Cpu)?;
-    let indices = a.apply_op1(ArgSort)?;
-    let a_sorted = a.gather(&indices, 0)?;
-    println!("{indices}");
-    println!("{a_sorted}");
+    let a = Tensor::new(&[[0.0f32, 1.0, 2.0], [3.0, 4.0, 5.0]], &Device::Cpu)?;
+    let b = Tensor::new(&[[88.0f32, 99.0]], &Device::Cpu)?;
+    let new_a = a.slice_scatter(&b, 1, 2)?;
+    assert_eq!(a.to_vec2::<f32>()?, [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
+    assert_eq!(new_a.to_vec2::<f32>()?, [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
     Ok(())
 }

candle-core/src/tensor.rs

@@ -1864,7 +1864,7 @@ impl Tensor {
                 }
                 (Storage::Cuda(storage), Device::Cpu) => Storage::Cpu(storage.to_cpu_storage()?),
                 (Storage::Metal(storage), Device::Cpu) => {
-                    println!("{storage:?} - {:?}", storage.to_cpu_storage()?);
+                    // println!("{storage:?} - {:?}", storage.to_cpu_storage()?);
                     Storage::Cpu(storage.to_cpu_storage()?)
                 }
                 (Storage::Cuda(storage), Device::Cuda(cuda)) => {

candle-core/tests/tensor_tests.rs

@@ -900,7 +900,9 @@ fn matmul(device: &Device) -> Result<()> {
     let b = Tensor::from_slice(&data, (2, 2), device)?;
 
     let c = a.matmul(&b)?;
+    let d = a.matmul(&c)?;
     assert_eq!(c.to_vec2::<f32>()?, &[[7.0f32, 10.0], [15.0, 22.0]]);
+    assert_eq!(d.to_vec2::<f32>()?, &[[37.0, 54.0], [81.0, 118.0]]);
 
     let data = vec![1.0f32, 2.0];
     let a = Tensor::from_slice(&data, (2, 1), device)?;

candle-examples/Cargo.toml

@@ -57,6 +57,7 @@ flash-attn = ["cuda", "candle-transformers/flash-attn", "dep:candle-flash-attn"]
 mkl = ["dep:intel-mkl-src", "candle/mkl", "candle-nn/mkl", "candle-transformers/mkl"]
 nccl = ["cuda", "cudarc/nccl", "dep:half"]
 onnx = ["candle-onnx"]
+metal = ["candle/metal", "candle-nn/metal"]
 
 [[example]]
 name = "llama_multiprocess"

candle-metal-kernels/Cargo.toml

@@ -10,7 +10,7 @@ categories = ["science"]
 license = "MIT OR Apache-2.0"
 
 [dependencies]
-metal = { version = "0.27.1", features = ["mps"], package="candle-metal" }
+metal = { version = "0.27.0", features = ["mps"]}
 once_cell = "1.18.0"
 thiserror = "1"
 tracing = "0.1.37"

candle-metal-kernels/src/affine.metal

@@ -29,15 +29,96 @@ kernel void FN_NAME( \
     if (id >= dim) { \
         return; \
     } \
-    const TYPENAME m = TYPENAME(mul); \
-    const TYPENAME a = TYPENAME(add); \
-    output[id] = input[id] * m + a; \
+    output[id] = TYPENAME(float(input[id]) * mul + add); \
 } \
+kernel void FN_NAME##_strided( \
+    constant size_t &dim, \
+    constant size_t &num_dims, \
+    constant size_t *dims, \
+    constant size_t *strides, \
+    constant float &mul, \
+    constant float &add, \
+    device const TYPENAME *input,  \
+    device TYPENAME *output, \
+    uint id [[ thread_position_in_grid ]] \
+) { \
+    if (id >= dim) { \
+        return; \
+    } \
+    output[id] = TYPENAME(float(input[get_strided_index(id, num_dims, dims, strides)]) * mul + add); \
+}
+
+#define POWF(FN_NAME, TYPENAME) \
+kernel void FN_NAME( \
+    constant size_t &dim, \
+    constant float &mul, \
+    device const TYPENAME *input,  \
+    device TYPENAME *output, \
+    uint id [[ thread_position_in_grid ]] \
+) { \
+    if (id >= dim) { \
+        return; \
+    } \
+    output[id] = TYPENAME(pow(input[id], TYPENAME(mul))); \
+} \
+kernel void FN_NAME##_strided( \
+    constant size_t &dim, \
+    constant size_t &num_dims, \
+    constant size_t *dims, \
+    constant size_t *strides, \
+    constant float &mul, \
+    device const TYPENAME *input,  \
+    device TYPENAME *output, \
+    uint id [[ thread_position_in_grid ]] \
+) { \
+    if (id >= dim) { \
+        return; \
+    } \
+    output[id] = TYPENAME(pow(input[get_strided_index(id, num_dims, dims, strides)], TYPENAME(mul))); \
+}
+
+#define ELU(FN_NAME, TYPENAME) \
+kernel void FN_NAME( \
+    constant size_t &dim, \
+    constant float &mul, \
+    device const TYPENAME *input,  \
+    device TYPENAME *output, \
+    uint id [[ thread_position_in_grid ]] \
+) { \
+    if (id >= dim) { \
+        return; \
+    } \
+    const TYPENAME x = input[id]; \
+    output[id] = TYPENAME((x > 0)?x: mul * exp(x - 1)); \
+} \
+kernel void FN_NAME##_strided( \
+    constant size_t &dim, \
+    constant size_t &num_dims, \
+    constant size_t *dims, \
+    constant size_t *strides, \
+    constant float &mul, \
+    device const TYPENAME *input,  \
+    device TYPENAME *output, \
+    uint id [[ thread_position_in_grid ]] \
+) { \
+    if (id >= dim) { \
+        return; \
+    } \
+    const TYPENAME x = input[get_strided_index(id, num_dims, dims, strides)]; \
+    output[id] = TYPENAME((x > 0)?x: mul * exp(x - 1)); \
+} \
+
 
 AFFINE(affine_float, float)
 AFFINE(affine_half, half)
+POWF(powf_float, float)
+POWF(powf_half, half)
+ELU(elu_float, float)
+ELU(elu_half, half)
 
 
 #if __METAL_VERSION__ >= 310
 AFFINE(affine_bfloat, bfloat);
+POWF(powf_bfloat, bfloat);
+ELU(elu_bfloat, bfloat);
 #endif

candle-metal-kernels/src/cast.metal

@@ -23,12 +23,12 @@ kernel void FN_NAME( \
     constant size_t &dim, \
     device const LEFT_TYPENAME *input,  \
     device RIGHT_TYPENAME *output, \
-    uint thread_position_in_grid [[ thread_position_in_grid ]] \
+    uint tid [[ thread_position_in_grid ]] \
 ) { \
-    if (thread_position_in_grid >= dim) { \
+    if (tid >= dim) { \
         return; \
     } \
-    output[thread_position_in_grid] = RIGHT_TYPENAME(input[thread_position_in_grid]); \
+    output[tid] = RIGHT_TYPENAME(input[tid]); \
 } \
 kernel void FN_NAME_STRIDED( \
     constant size_t &dim, \
@@ -37,15 +37,19 @@ kernel void FN_NAME_STRIDED( \
     constant size_t *strides, \
     device const LEFT_TYPENAME *input,  \
     device RIGHT_TYPENAME *output, \
-    uint i [[ thread_position_in_grid ]] \
+    uint tid [[ thread_position_in_grid ]] \
 ) { \
-    if (i >= dim) { \
+    if (tid >= dim) { \
         return; \
     } \
-    output[i] = RIGHT_TYPENAME(input[get_strided_index(i, num_dims, dims, strides)]); \
+    output[tid] = RIGHT_TYPENAME(input[get_strided_index(tid, num_dims, dims, strides)]); \
 } \
 
-CAST(cast_u32_f32, cast_u32_f32_strided, int32_t, float)
+CAST(cast_u32_f32, cast_u32_f32_strided, uint32_t, float)
+CAST(cast_u32_u8, cast_u32_u8_strided, uint32_t, uint8_t)
+CAST(cast_u8_u32, cast_u8_u32_strided, uint8_t, uint32_t)
+CAST(cast_f16_f32, cast_f16_f32_strided, half, float)
+CAST(cast_f32_f16, cast_f32_f16_strided, float, half)
 
 #if __METAL_VERSION__ >= 310
 #endif

candle-metal-kernels/src/indexing.metal

@@ -16,16 +16,16 @@ kernel void NAME( \
     if (gid >= dst_size) { \
         return; \
     } \
-    const size_t id_i = gid / right_size / left_size; \
+    const size_t id_i = (gid / right_size) % ids_size; \
+    const INDEX_TYPENAME input_i = min(input_ids[id_i], (INDEX_TYPENAME)(src_dim_size - 1)); \
     const size_t right_rank_i = gid % right_size; \
-    const size_t left_rank_i = gid % left_size; \
+    const size_t left_rank_i = gid / right_size / ids_size; \
     /* \
     // Force prevent out of bounds indexing \
     // since there doesn't seem to be a good way to force crash \
     // No need to check for zero we're only allowing unsized. \
     */ \
-    const INDEX_TYPENAME input_i = min(input_ids[id_i], (INDEX_TYPENAME)(src_dim_size - 1)); \
-    const size_t src_i = ((input_i * right_size) + right_rank_i) * left_size + left_rank_i; \
+    const size_t src_i = left_rank_i * src_dim_size * right_size + input_i * right_size + right_rank_i; \
     output[gid] = input[src_i]; \
 }
 
@@ -75,6 +75,7 @@ kernel void FN_NAME( \
 
 
 INDEX_OP(is_u32_f32, uint, float)
+INDEX_OP(is_u32_f16, uint, half)
 
 
 #if __METAL_VERSION__ >= 310

candle-metal-kernels/src/reduce.metal

@@ -1,6 +1,8 @@
 #include <metal_stdlib>
 using namespace metal;
 
+#define MAX(x, y) ((x) > (y) ? (x) : (y))
+
 METAL_FUNC uint get_strided_index(
     uint idx,
     constant size_t &num_dims,
@@ -16,21 +18,21 @@ METAL_FUNC uint get_strided_index(
     return strided_i;
 }
 
-constant int THREADGROUP_SIZE = 256;
+constant int THREADGROUP_SIZE = 2048;
 
-# define REDUCE(FN, NAME, TYPENAME) \
+# define REDUCE(FN, NAME, T) \
 kernel void NAME( \
     constant size_t &src_numel, \
     constant size_t &el_to_sum_per_block, \
-    device const TYPENAME *src,  \
-    device TYPENAME *dst, \
+    device const T *src,  \
+    device T *dst, \
     uint id [[ thread_position_in_grid ]], \
     uint tid [[ thread_index_in_threadgroup ]], \
     uint dst_id [[ threadgroup_position_in_grid ]], \
-    uint blockDim [[ threads_per_threadgroup ]] \
+    uint block_dim [[ threads_per_threadgroup ]] \
 ) { \
      \
-   threadgroup float shared_memory[THREADGROUP_SIZE]; \
+   threadgroup T shared_memory[THREADGROUP_SIZE]; \
       \
    shared_memory[tid] = 0; \
    /* \
@@ -45,10 +47,10 @@ kernel void NAME( \
      // TODO: Fast version for the contiguous case. \
      // size_t strided_i = get_strided_index(idx, num_dims, dims, strides); \
      */ \
-     TYPENAME x = shared_memory[tid]; \
-     TYPENAME y = src[idx]; \
+     T x = shared_memory[tid]; \
+     T y = src[idx]; \
      shared_memory[tid] = FN; \
-     idx += blockDim; \
+     idx += block_dim; \
    } \
       \
    threadgroup_barrier(mem_flags::mem_none); \
@@ -56,10 +58,10 @@ kernel void NAME( \
    /* \
    // reduction in shared memory \
    */ \
-   for (uint s = blockDim / 2; s > 0; s >>= 1) { \
+   for (uint s = block_dim / 2; s > 0; s >>= 1) { \
        if (tid < s) { \
-           TYPENAME x = shared_memory[tid]; \
-           TYPENAME y = shared_memory[tid + s]; \
+           T x = shared_memory[tid]; \
+           T y = shared_memory[tid + s]; \
            shared_memory[tid] = FN; \
        } \
        threadgroup_barrier(mem_flags::mem_none); \
@@ -68,72 +70,80 @@ kernel void NAME( \
    dst[dst_id] = shared_memory[0]; \
 } \
 
-kernel void softmax_float(
-    constant size_t &src_numel,
-    constant size_t &el_to_sum_per_block,
-    device const float *src, 
-    device float *dst,
-    uint id [[ thread_position_in_grid ]],
-    uint tid [[ thread_index_in_threadgroup ]],
-    uint dst_id [[ threadgroup_position_in_grid ]],
-    uint blockDim [[ threads_per_threadgroup ]]
-) {
-    
-   threadgroup float shared_memory[THREADGROUP_SIZE];
-     
-   shared_memory[tid] = -INFINITY;
-   // Elements summed in this block range from dst_id * el_to_sum_per_block
-   // to (dst_id + 1) * el_to_sum_per_block.
-   size_t start_idx = dst_id * el_to_sum_per_block;
-   size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel);
-   size_t idx = start_idx + tid;
-
-   while (idx < stop_idx) {
-     // TODO: Fast version for the contiguous case.
-     shared_memory[tid] = max(shared_memory[tid], src[idx]);
-     idx += blockDim;
-   }
-     
-   threadgroup_barrier(mem_flags::mem_none);
-   
-   // reduction in shared memory
-   for (uint s = blockDim / 2; s > 0; s >>= 1) {
-       if (tid < s) {
-           shared_memory[tid] = max(shared_memory[tid], shared_memory[tid + s]);
-       }
-       threadgroup_barrier(mem_flags::mem_none);
-   }
-   
-   float max = shared_memory[0];
-
-   shared_memory[tid] = 0;
-
-   // Restart
-   idx = start_idx + tid;
-   while (idx < stop_idx) {
-     // TODO: Fast version for the contiguous case.
-     const float val = exp(src[idx] - max);
-     dst[idx] = val; 
-     shared_memory[tid] += val;
-     idx += blockDim;
-   }
-   // reduction in shared memory
-   for (uint s = blockDim / 2; s > 0; s >>= 1) {
-       if (tid < s) {
-           shared_memory[tid] += shared_memory[tid + s];
-       }
-       threadgroup_barrier(mem_flags::mem_none);
-   }
-
-   const float inv_acc = 1/shared_memory[0];
-   idx = start_idx + tid;
-   while (idx < stop_idx) {
-     dst[idx] *= inv_acc; 
-     idx += blockDim;
-   }
-}
-
 
 REDUCE(x + y, fast_sum_float, float)
 REDUCE(x * y, fast_mul_float, float)
 REDUCE(max(x, y), fast_max_float, float)
+
+#define SOFTMAX(NAME, T)                                                          \
+kernel void NAME(                                                                 \
+    constant size_t &src_numel,                                                   \
+    constant size_t &el_to_sum_per_block,                                         \
+    device const T *src,                                                          \
+    device T *dst,                                                                \
+                                                                                  \
+    uint id [[ thread_position_in_grid ]],                                        \
+    uint tid [[ thread_index_in_threadgroup ]],                                   \
+    uint dst_id [[ threadgroup_position_in_grid ]],                               \
+    uint block_dim [[ threads_per_threadgroup ]]                                  \
+) {                                                                               \
+    threadgroup float shared_memory[THREADGROUP_SIZE];                                \
+    shared_memory[tid] = -INFINITY;                                            \
+    size_t start_idx = dst_id * el_to_sum_per_block;                              \
+    size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel);            \
+    size_t idx = start_idx + tid;                                                 \
+                                                                                  \
+                                                                                  \
+    float tmp = -INFINITY; \
+    while (idx < stop_idx) {                                                      \
+        tmp = MAX(tmp, float(src[idx]));                   \
+        idx += block_dim;                                                         \
+    }                                                                             \
+    shared_memory[tid] = tmp; \
+                                                                                  \
+    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+                                                                                  \
+    for (uint s = block_dim / 2; s > 0; s >>= 1) {                                \
+        if (tid < s) {                                                            \
+            shared_memory[tid] = MAX(shared_memory[tid], shared_memory[tid + s]); \
+        }                                                                         \
+        threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+    }                                                                             \
+                                                                                  \
+    /* wait for shared_memory[0] to be filled */ \
+    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+                                                                                  \
+    float _max = shared_memory[0];                                                    \
+                                                                                  \
+    /* prevent tid=0 from overwriting _max before other threads have written it */ \
+    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+    shared_memory[tid] = 0;                                                       \
+                                                                                  \
+    idx = start_idx + tid;                                                        \
+    while (idx < stop_idx) {                                                      \
+        const float val = exp(float(src[idx]) - _max);                                    \
+        dst[idx] = T(val);                                                           \
+        shared_memory[tid] += val;                                                \
+        idx += block_dim;                                                         \
+    }                                                                             \
+    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+    for (uint s = block_dim / 2; s > 0; s >>= 1) {                                \
+        if (tid < s) {                                                            \
+            shared_memory[tid] += shared_memory[tid + s];                         \
+        }                                                                         \
+        threadgroup_barrier(mem_flags::mem_threadgroup);                              \
+    }                                                                             \
+                                                                                  \
+    const T inv_acc = T(1.0/shared_memory[0]);                                         \
+    idx = start_idx + tid;                                                        \
+    while (idx < stop_idx) {                                                      \
+        dst[idx] *= inv_acc;                                                      \
+        idx += block_dim;                                                         \
+    }                                                                             \
+}                                                                                 \
+
+SOFTMAX(softmax_float, float)
+SOFTMAX(softmax_half, half)
+#if __METAL_VERSION__ >= 310
+SOFTMAX(softmax_bfloat, bfloat)
+#endif

candle-metal-kernels/src/ternary.metal

@@ -32,6 +32,9 @@ kernel void FN_NAME(  \
     device TYPENAME *out ,\
     uint i [[ thread_position_in_grid ]] \
 ) {  \
+   if (i >= numel){ \
+       return; \
+   } \
    uint strided_i = get_strided_index(i, num_dims, dims, strides); \
    uint strided_i_t = get_strided_index(i, num_dims, dims, strides_t); \
    uint strided_i_f = get_strided_index(i, num_dims, dims, strides_f); \

candle-metal-kernels/src/test.swift

@@ -0,0 +1,209 @@
+
+import Metal
+import MetalPerformanceShadersGraph
+
+
+
+let type = MTLDataType.float;
+let dataType = type;
+var B = 2;
+var M = 2;
+var N = 2;
+var K = 2;
+var A_trans = false;
+var B_trans = false;
+var D_trans = false;
+var alpha = Float(1.0);
+var beta = Float(0.0);
+var batched = B > 1;
+var fused_activation = false;
+var fused_bias = false;
+let constants = MTLFunctionConstantValues()
+constants.setConstantValue(&M, type: .uint, index: 0)
+constants.setConstantValue(&N, type: .uint, index: 1)
+constants.setConstantValue(&K, type: .uint, index: 2)
+constants.setConstantValue(&A_trans, type: .bool, index: 10)
+constants.setConstantValue(&B_trans, type: .bool, index: 11)
+constants.setConstantValue(&D_trans, type: .bool, index: 13)
+constants.setConstantValue(&alpha, type: .float, index: 20)
+constants.setConstantValue(&beta, type: .float, index: 21)
+constants.setConstantValue(&batched, type: .bool, index: 100)
+constants.setConstantValue(&fused_activation, type: .bool, index: 101)
+constants.setConstantValue(&fused_bias, type: .bool, index: 50001)
+
+
+var M_simd = UInt16(16)
+var N_simd = UInt16(16)
+var K_simd = UInt16(32)
+var M_splits = UInt16(2)
+var N_splits = UInt16(2)
+constants.setConstantValue(&M_simd, type: .ushort, index: 200)
+constants.setConstantValue(&N_simd, type: .ushort, index: 201)
+constants.setConstantValue(&K_simd, type: .ushort, index: 202)
+constants.setConstantValue(&M_splits, type: .ushort, index: 210)
+constants.setConstantValue(&N_splits, type: .ushort, index: 211)
+
+let M_group = M_simd * M_splits
+let N_group = N_simd * N_splits
+
+// Satisfy Metal API validation.
+#if DEBUG
+do {
+  var garbage: SIMD4<UInt64> = .zero
+  constants.setConstantValue(&garbage, type: .bool, index: 102)
+  constants.setConstantValue(&garbage, type: .bool, index: 103)
+  constants.setConstantValue(&garbage, type: .bool, index: 113)
+  constants.setConstantValue(&garbage, type: .bool, index: 50000)
+}
+#endif
+
+let device = MTLCopyAllDevices().first!
+device.shouldMaximizeConcurrentCompilation = true
+
+var libraryURL = URL.init(string: "/Users/nicolas/src/candle/candle-metal-kernels/")!;
+libraryURL.append(component: "src")
+libraryURL.append(component: "libMetalFlashAttention.metallib")
+let library = try! device.makeLibrary(URL: libraryURL)
+
+var name: String
+    switch dataType {
+    case .half: name = "hgemm"
+    case .float: name = "sgemm"
+    default: fatalError()
+    }
+let function = try! library.makeFunction(
+  name: name, constantValues: constants)
+
+let A_block_length = M_group * K_simd
+let B_block_length = K_simd * N_group
+
+var blockElements = A_block_length + B_block_length;
+if (M % 8 != 0) && (N % 8 != 0) {
+  let C_block_length = M_group * N_group;
+  blockElements = max(C_block_length, blockElements)
+}
+if fused_bias {
+  if D_trans {
+    blockElements = max(blockElements, M_group)
+  } else {
+    blockElements = max(blockElements, N_group)
+  }
+}
+// let blockBytes = blockElements * UInt16(dataType.size)
+let elementSize = 4
+let blockBytes = blockElements * UInt16(elementSize)
+
+func ceilDivide(target: Int, granularity: UInt16) -> Int {
+  (target + Int(granularity) - 1) / Int(granularity)
+}
+var gridSize = MTLSize(
+  width: ceilDivide(target: N, granularity: N_group),
+  height: ceilDivide(target: M, granularity: M_group),
+  depth: 1)
+let groupSize = MTLSize(
+  width: Int(32 * M_splits * N_splits),
+  height: 1,
+  depth: 1)
+
+let commandQueue = device.makeCommandQueue()!
+
+let threadgroupMemoryLength = blockBytes;
+
+let rowsA = M;
+let columnsA = K;
+let rowsB = K;
+let columnsB = N;
+let rowsC = M;
+let columnsC = N;
+var arrayA = [Float](repeating: 0, count: B * rowsA * columnsA)
+
+var arrayB = [Float](repeating: 0, count: B * rowsB * columnsB)
+
+var arrayC = [Float](repeating: 0, count: B * rowsC * columnsC)
+var arrayD = [Float](repeating: 0, count: B * rowsC * columnsC)
+for i in 0..<arrayA.count {
+  arrayA[i] = Float(i)
+}
+
+for i in 0..<arrayB.count {
+  arrayB[i] = Float(i)
+}
+
+let bufferA = device.makeBuffer(bytes: arrayA, length: B * rowsA * columnsA * MemoryLayout<Float>.stride, options: [])!
+
+let bufferB = device.makeBuffer(bytes: arrayB, length: B * rowsB * columnsB * MemoryLayout<Float>.stride, options: [])!
+
+let bufferC = device.makeBuffer(length: B * rowsC * columnsC * MemoryLayout<Float>.stride, options: [])!
+let bufferD = device.makeBuffer(length: B * rowsC * columnsC * MemoryLayout<Float>.stride, options: [])!
+
+
+let pipeline = try device.makeComputePipelineState(function: function)
+
+func call(bufferA: MTLBuffer, bufferB: MTLBuffer, bufferC: MTLBuffer){ 
+    let encoder = commandBuffer.makeComputeCommandEncoder(dispatchType: MTLDispatchType.serial)!
+    encoder.setComputePipelineState(pipeline)
+    encoder.setThreadgroupMemoryLength(Int(threadgroupMemoryLength), index: 0)
+
+    encoder.setBuffer(bufferA, offset: 0, index: 0)
+    encoder.setBuffer(bufferB, offset: 0, index: 1)
+    encoder.setBuffer(bufferC, offset: 0, index: 2)
+    let gridZ: Int = B
+    if batched{
+      func byteStride(shape: [Int]) -> Int {
+        let rank = shape.count
+        var output = elementSize * shape[rank - 2] * shape[rank - 1]
+        if shape.dropLast(2).reduce(1, *) == 1 {
+          output = 0
+        }
+        return output
+      }
+      let byteStrideA = M*K*elementSize
+      let byteStrideB = N*K*elementSize
+      let byteStrideC = M*N*elementSize
+      
+      let byteStrideD = 0
+      withUnsafeTemporaryAllocation(
+        of: SIMD4<UInt64>.self, capacity: gridZ
+      ) { buffer in
+        for i in 0..<buffer.count {
+          buffer[i] = SIMD4(
+            UInt64(truncatingIfNeeded: i * byteStrideA),
+            UInt64(truncatingIfNeeded: i * byteStrideB),
+            UInt64(truncatingIfNeeded: i * byteStrideC),
+            UInt64(truncatingIfNeeded: i * byteStrideD))
+        }
+        
+        let bufferLength = buffer.count * MemoryLayout<SIMD4<UInt64>>.stride
+        assert(MemoryLayout<SIMD4<UInt64>>.stride == 8 * 4)
+        encoder.setBytes(buffer.baseAddress!, length: bufferLength, index: 10)
+      }
+    }
+    gridSize.depth = gridZ
+
+
+    encoder.dispatchThreadgroups(
+      gridSize, threadsPerThreadgroup: groupSize
+    )
+    encoder.endEncoding()
+}
+
+var commandBuffer = commandQueue.makeCommandBuffer()!
+call(bufferA:bufferA, bufferB:bufferB, bufferC:bufferC)
+commandBuffer.commit()
+commandBuffer = commandQueue.makeCommandBuffer()!
+commandBuffer.encodeWaitForEvent(event, value: 2)
+call(bufferA:bufferA, bufferB:bufferC, bufferC:bufferD)
+commandBuffer.commit()
+
+commandBuffer.waitUntilCompleted()
+var contents = bufferC.contents();
+var count = B * rowsA * columnsB;
+var typedPointer = contents.bindMemory(to: Float.self, capacity: count)
+var bufferedPointer = UnsafeBufferPointer(start: typedPointer, count: count)
+print("First matmul is OK", Array(bufferedPointer))
+
+contents = bufferD.contents();
+count = B * rowsA * columnsB;
+typedPointer = contents.bindMemory(to: Float.self, capacity: count)
+bufferedPointer = UnsafeBufferPointer(start: typedPointer, count: count)
+print("This should be filled", Array(bufferedPointer))

candle-metal-kernels/src/tests.rs

@@ -1,7 +1,14 @@
 use super::*;
-use half::f16;
+use half::{bf16, f16};
 use metal::{CompileOptions, Device, MTLResourceOptions, MTLSize, NSUInteger};
 
+fn read_to_vec<T: Clone>(buffer: &Buffer, n: usize) -> Vec<T> {
+    let ptr = buffer.contents() as *const T;
+    assert!(!ptr.is_null());
+    let slice = unsafe { std::slice::from_raw_parts(ptr, n) };
+    slice.to_vec()
+}
+
 fn new_buffer<T>(device: &Device, data: &[T]) -> Buffer {
     let options = MTLResourceOptions::StorageModeManaged;
     let ptr = data.as_ptr() as *const core::ffi::c_void;
@@ -23,13 +30,19 @@ fn approx_f16(v: Vec<f16>, digits: i32) -> Vec<f32> {
     v.iter().map(|t| f32::round(t.to_f32() * b) / b).collect()
 }
 
+fn approx_bf16(v: Vec<bf16>, digits: i32) -> Vec<f32> {
+    let b = 10f32.powi(digits);
+    v.iter().map(|t| f32::round(t.to_f32() * b) / b).collect()
+}
+
 fn run<T: Clone>(v: &[T], name: unary::contiguous::Kernel) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let input = new_buffer(&device, v);
-    let mut output = new_buffer(&device, v);
+    let output = new_buffer(&device, v);
     call_unary_contiguous(
         &device,
         command_buffer,
@@ -37,23 +50,24 @@ fn run<T: Clone>(v: &[T], name: unary::contiguous::Kernel) -> Vec<T> {
         name,
         v.len(),
         &input,
-        &mut output,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
-    output.read_to_vec::<T>(v.len())
+    read_to_vec(&output, v.len())
 }
 
 fn run_binary<T: Clone>(x: &[T], y: &[T], name: binary::contiguous::Kernel) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let options = MTLResourceOptions::StorageModeManaged;
     let left = new_buffer(&device, x);
     let right = new_buffer(&device, y);
-    let mut output = device.new_buffer(std::mem::size_of_val(x) as u64, options);
+    let output = device.new_buffer(std::mem::size_of_val(x) as u64, options);
     call_binary_contiguous(
         &device,
         command_buffer,
@@ -62,12 +76,12 @@ fn run_binary<T: Clone>(x: &[T], y: &[T], name: binary::contiguous::Kernel) -> V
         x.len(),
         &left,
         &right,
-        &mut output,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
-    output.read_to_vec::<T>(x.len())
+    read_to_vec(&output, x.len())
 }
 
 fn run_strided<T: Clone>(
@@ -81,8 +95,9 @@ fn run_strided<T: Clone>(
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let input = new_buffer(&device, v);
-    let mut output = new_buffer(&device, v);
-    let kernels = Kernels::new();
+    let output = new_buffer(&device, v);
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     call_unary_strided(
         &device,
         command_buffer,
@@ -92,13 +107,13 @@ fn run_strided<T: Clone>(
         &input,
         strides,
         offset,
-        &mut output,
+        &output,
         0,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
-    output.read_to_vec::<T>(v.len())
+    read_to_vec(&output, v.len())
 }
 
 #[test]
@@ -200,6 +215,25 @@ fn cos_strided_random() {
     );
 }
 
+#[test]
+fn gelu_f16() {
+    let v: Vec<f16> = [-10f32, -1.0, 0., 1., 2., 3., 10.0, 20.0]
+        .iter()
+        .map(|v| f16::from_f32(*v))
+        .collect();
+    let expected: Vec<f32> = vec![-0.0, -0.16, 0.0, 0.84, 1.96, 3.0, 10.0, 20.0];
+    let results = run(&v, unary::contiguous::gelu::HALF);
+    assert_eq!(approx_f16(results, 2), expected);
+}
+
+#[test]
+fn gelu_f32() {
+    let v: Vec<f32> = vec![-10f32, -1.0, 0., 1., 2., 3., 10.0, 20.0];
+    let expected: Vec<f32> = vec![-0.0, -0.159, 0.0, 0.841, 1.955, 2.996, 10.0, 20.0];
+    let results = run(&v, unary::contiguous::gelu::FLOAT);
+    assert_eq!(approx(results, 3), expected);
+}
+
 #[test]
 fn binary_add_f32() {
     let left = vec![1.0f32, 2.0, 3.0];
@@ -216,11 +250,14 @@ fn binary_add_f32() {
 
 fn cast<T: Clone, U: Clone>(v: &[T], name: &'static str) -> Vec<U> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let input = new_buffer(&device, v);
-    let mut output = new_buffer(&device, v);
+    let options = MTLResourceOptions::StorageModeManaged;
+    let size = (v.len() * std::mem::size_of::<U>()) as u64;
+    let output = device.new_buffer(size, options);
 
     call_cast_contiguous(
         &device,
@@ -229,12 +266,13 @@ fn cast<T: Clone, U: Clone>(v: &[T], name: &'static str) -> Vec<U> {
         name,
         v.len(),
         &input,
-        &mut output,
+        0,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
-    output.read_to_vec::<U>(v.len())
+    read_to_vec(&output, v.len())
 }
 
 #[test]
@@ -245,21 +283,28 @@ fn cast_u32_f32() {
     assert_eq!(approx(results, 4), vec![1.0f32, 2.0, 3.0]);
     assert_eq!(approx(expected, 4), vec![1.0f32, 2.0, 3.0]);
 
+    let v = vec![1.0f32, 2.0, 3.0];
+    let input: Vec<f16> = v.iter().map(|v| f16::from_f32(*v)).collect();
+    let results: Vec<f32> = cast(&input, "cast_f16_f32");
+    assert_eq!(results, vec![1.0f32, 2.0, 3.0]);
+
     let v = vec![1.0f32; 10_000];
-    let results = run(&v, unary::contiguous::cos::FLOAT);
-    let expected: Vec<_> = v.iter().map(|v| v.cos()).collect();
-    assert_eq!(approx(results, 4), vec![0.5403; 10_000]);
-    assert_eq!(approx(expected, 4), vec![0.5403; 10_000]);
+    let input: Vec<f16> = v.iter().map(|v| f16::from_f32(*v)).collect();
+    let results: Vec<f32> = cast(&input, "cast_f16_f32");
+    assert_eq!(results.len(), 10_000);
+    assert_eq!(&results[..10], vec![1.0f32; 10]);
+    assert_eq!(results, vec![1.0f32; 10_000]);
 }
 
 fn run_affine<T: Clone>(v: &[T], mul: f64, add: f64) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
 
     let input = new_buffer(&device, v);
-    let mut output = new_buffer(&device, v);
+    let output = new_buffer(&device, v);
 
     let size = v.len();
 
@@ -267,9 +312,10 @@ fn run_affine<T: Clone>(v: &[T], mul: f64, add: f64) -> Vec<T> {
         &device,
         command_buffer,
         &kernels,
+        "affine_float",
         size,
         &input,
-        &mut output,
+        &output,
         mul as f32,
         add as f32,
     )
@@ -277,7 +323,44 @@ fn run_affine<T: Clone>(v: &[T], mul: f64, add: f64) -> Vec<T> {
     command_buffer.commit();
     command_buffer.wait_until_completed();
 
-    output.read_to_vec::<T>(v.len())
+    read_to_vec(&output, v.len())
+}
+
+fn run_affine_strided<T: Clone>(
+    v: &[T],
+    shape: &[usize],
+    strides: &[usize],
+    mul: f64,
+    add: f64,
+) -> Vec<T> {
+    let device = device();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
+    let command_queue = device.new_command_queue();
+    let command_buffer = command_queue.new_command_buffer();
+
+    let input = new_buffer(&device, v);
+    let output = new_buffer(&device, v);
+
+    call_affine_strided(
+        &device,
+        command_buffer,
+        &kernels,
+        "affine_float_strided",
+        shape,
+        &input,
+        strides,
+        0,
+        &output,
+        mul as f32,
+        add as f32,
+    )
+    .unwrap();
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    let len: usize = shape.iter().product();
+    read_to_vec(&output, len)
 }
 
 #[test]
@@ -295,6 +378,18 @@ fn affine() {
     assert_eq!(result, vec![2.6; 40_000]);
 }
 
+#[test]
+fn affine_strided() {
+    let input = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
+    let mul = 1.5;
+    let add = 1.1;
+    let shape = [4];
+    let strides = [2];
+    let result = run_affine_strided(&input, &shape, &strides, mul, add);
+    // 1 on 2
+    assert_eq!(result, vec![2.6, 5.6, 8.6, 11.6]);
+}
+
 #[test]
 fn index_select() {
     let embedding = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
@@ -313,7 +408,26 @@ fn index_select() {
         result,
         vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 1.0f32, 2.0, 3.0, 4.0, 5.0]
     );
+}
 
+#[test]
+fn index_select_f16() {
+    let embedding: Vec<_> = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+        .into_iter()
+        .map(|x| f16::from_f32(x))
+        .collect();
+    let shape = [5, 2];
+    let ids = [0u32, 4, 2];
+    let dim = 0;
+    let result = run_index_select(&embedding, &shape, &ids, dim);
+    assert_eq!(
+        approx_f16(result, 4),
+        vec![1.0f32, 2.0, 9.0, 10.0, 5.0, 6.0]
+    );
+}
+
+#[test]
+fn index_select_dim1() {
     let embedding = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
     let shape = [5, 2];
     let ids = [0u32, 1, 0];
@@ -321,7 +435,7 @@ fn index_select() {
     let result = run_index_select(&embedding, &shape, &ids, dim);
     assert_eq!(
         result,
-        vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 1.0f32, 2.0, 3.0, 4.0, 5.0]
+        vec![1.0f32, 2.0, 1.0, 3.0, 4.0, 3.0, 5.0, 6.0, 5.0, 7.0, 8.0f32, 7.0, 9.0, 10.0, 9.0]
     );
 }
 
@@ -341,27 +455,34 @@ fn run_index_select<T: Clone, I: Clone + std::fmt::Debug>(
     let left_size: usize = shape[..dim].iter().product();
     let right_size: usize = shape[dim + 1..].iter().product();
     let dst_el = ids.len() * left_size * right_size;
-    let mut dst_buffer = new_buffer(&device, &vec![0.0f32; dst_el]);
+    let dst_buffer = new_buffer(&device, &vec![0.0f32; dst_el]);
 
-    let kernels = Kernels::new();
+    let name = match core::mem::size_of::<T>() {
+        4 => "is_u32_f32",
+        2 => "is_u32_f16",
+        _ => unimplemented!(),
+    };
+
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     call_index_select(
         &device,
         &command_buffer,
         &kernels,
-        "is_u32_f32",
+        name,
         shape,
         ids.len(),
         dim,
         &embeddings_buffer,
         &ids_buffer,
-        &mut dst_buffer,
+        &dst_buffer,
     )
     .unwrap();
 
     command_buffer.commit();
     command_buffer.wait_until_completed();
 
-    dst_buffer.read_to_vec::<T>(dst_el)
+    read_to_vec(&dst_buffer, dst_el)
 }
 
 #[test]
@@ -427,7 +548,7 @@ fn index_add() {
     let expected = vec![
         2.0, 3.0, 4.0, 1.0, 1.0, 1.0, 8.0, 9.0, 10.0, 1.0, 1.0, 1.0, 5.0, 6.0, 7.0,
     ];
-    let result = outputs_buffer.read_to_vec::<f32>(right.len());
+    let result: Vec<f32> = read_to_vec(&outputs_buffer, right.len());
     assert_eq!(result, expected);
 }
 
@@ -439,19 +560,20 @@ fn cos_f16() {
         .collect();
     let results = run(&v, unary::contiguous::cos::HALF);
     let expected: Vec<f16> = v.iter().map(|v| f16::from_f32(v.to_f32().cos())).collect();
-    assert_eq!(approx_f16(results, 4), vec![0.5405, -0.4163, -0.9902]);
-    assert_eq!(approx_f16(expected, 4), vec![0.5405, -0.4163, -0.9902]);
+    assert_eq!(approx_f16(results, 2), vec![0.54, -0.42, -0.99]);
+    assert_eq!(approx_f16(expected, 2), vec![0.54, -0.42, -0.99]);
 }
 
 fn run_reduce<T: Clone>(v: &[T], out_length: usize, name: &'static str) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let input = new_buffer(&device, v);
 
     let options = MTLResourceOptions::StorageModeManaged;
-    let mut output = device.new_buffer((out_length * core::mem::size_of::<T>()) as u64, options);
+    let output = device.new_buffer((out_length * core::mem::size_of::<T>()) as u64, options);
     call_reduce_contiguous(
         &device,
         command_buffer,
@@ -460,22 +582,24 @@ fn run_reduce<T: Clone>(v: &[T], out_length: usize, name: &'static str) -> Vec<T
         v.len(),
         out_length,
         &input,
-        &mut output,
+        0,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
 
-    output.read_to_vec::<T>(out_length)
+    read_to_vec(&output, out_length)
 }
 
 fn run_softmax<T: Clone + std::fmt::Debug>(v: &[T], last_dim: usize, name: &'static str) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let input = new_buffer(&device, v);
-    let mut output = new_buffer(&device, v);
+    let output = new_buffer(&device, v);
     call_last_softmax(
         &device,
         command_buffer,
@@ -484,13 +608,13 @@ fn run_softmax<T: Clone + std::fmt::Debug>(v: &[T], last_dim: usize, name: &'sta
         v.len(),
         last_dim,
         &input,
-        &mut output,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
 
-    output.read_to_vec::<T>(v.len())
+    read_to_vec(&output, v.len())
 }
 
 #[test]
@@ -521,6 +645,24 @@ fn softmax() {
         vec![0.0043, 0.0116, 0.0315, 0.0858, 0.2331, 0.6337]
     );
 
+    let last_dim = 4096;
+    let n = 200;
+    let mut v = vec![0.0; n * last_dim];
+    for i in 0..n {
+        v[i * last_dim] = 20.0;
+    }
+    let results = run_softmax(&v, last_dim, "softmax_float");
+    let results = approx(results, 4);
+    println!("{results:?}");
+    assert_eq!(
+        results.iter().map(|&s| s.round() as usize).sum::<usize>(),
+        n
+    );
+    assert_eq!(results[0], 1.0);
+    assert_eq!(results[1], 0.0);
+    assert_eq!(results[last_dim], 1.0);
+    assert_eq!(results[2 * last_dim], 1.0);
+
     let v = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0];
     let last_dim = 6;
     let results = run_softmax(&v, last_dim, "softmax_float");
@@ -536,6 +678,28 @@ fn softmax() {
         approx(results, 4),
         vec![0.0900, 0.2447, 0.6652, 0.0900, 0.2447, 0.6652]
     );
+
+    let v = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]
+        .iter()
+        .map(|v| f16::from_f32(*v))
+        .collect::<Vec<_>>();
+    let last_dim = 6;
+    let results = run_softmax(&v, last_dim, "softmax_half");
+    assert_eq!(
+        approx_f16(results, 4),
+        vec![0.0043, 0.0116, 0.0316, 0.0858, 0.2332, 0.6338]
+    );
+
+    let v = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]
+        .iter()
+        .map(|v| bf16::from_f32(*v))
+        .collect::<Vec<_>>();
+    let last_dim = 6;
+    let results = run_softmax(&v, last_dim, "softmax_bfloat");
+    assert_eq!(
+        approx_bf16(results, 4),
+        vec![0.0043, 0.0116, 0.0315, 0.0859, 0.2324, 0.6328]
+    );
 }
 
 fn run_where_cond<I: Clone, T: Clone>(
@@ -549,7 +713,8 @@ fn run_where_cond<I: Clone, T: Clone>(
     name: &'static str,
 ) -> Vec<T> {
     let device = device();
-    let kernels = Kernels::new();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let options = MTLResourceOptions::StorageModeManaged;
@@ -571,7 +736,7 @@ fn run_where_cond<I: Clone, T: Clone>(
         options,
     );
 
-    let mut output = device.new_buffer((length * core::mem::size_of::<T>()) as u64, options);
+    let output = device.new_buffer((length * core::mem::size_of::<T>()) as u64, options);
     call_where_cond_strided(
         &device,
         command_buffer,
@@ -584,13 +749,13 @@ fn run_where_cond<I: Clone, T: Clone>(
         (&left_stride, left_offset),
         &right,
         (&cond_stride, cond_offset),
-        &mut output,
+        &output,
     )
     .unwrap();
     command_buffer.commit();
     command_buffer.wait_until_completed();
 
-    output.read_to_vec::<T>(length)
+    read_to_vec(&output, length)
 }
 
 #[test]
@@ -614,3 +779,93 @@ fn where_cond() {
     );
     assert_eq!(approx(results, 4), vec![-1.0f32, 2.0, -3.0, -4.0, 5.0, 6.0]);
 }
+
+fn run_gemm<T: Clone>(
+    (b, m, n, k): (usize, usize, usize, usize),
+    lhs: &[T],
+    lhs_stride: Vec<usize>,
+    lhs_offset: usize,
+    rhs: &[T],
+    rhs_stride: Vec<usize>,
+    rhs_offset: usize,
+) -> Vec<T> {
+    let device = device();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
+    let command_queue = device.new_command_queue();
+    let command_buffer = command_queue.new_command_buffer();
+    let options = MTLResourceOptions::StorageModeManaged;
+
+    let lhs = device.new_buffer_with_data(
+        lhs.as_ptr() as *const core::ffi::c_void,
+        std::mem::size_of_val(lhs) as u64,
+        options,
+    );
+    let rhs = device.new_buffer_with_data(
+        rhs.as_ptr() as *const core::ffi::c_void,
+        std::mem::size_of_val(rhs) as u64,
+        options,
+    );
+    let length = b * m * n;
+    let output = device.new_buffer((length * core::mem::size_of::<T>()) as u64, options);
+    call_gemm(
+        &device,
+        command_buffer,
+        &kernels,
+        "sgemm",
+        (b, m, n, k),
+        &lhs_stride,
+        lhs_offset,
+        &lhs,
+        &rhs_stride,
+        rhs_offset,
+        &rhs,
+        &output,
+    )
+    .unwrap();
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    read_to_vec(&output, length)
+}
+
+#[test]
+fn gemm() {
+    let (b, m, n, k) = (1, 2, 4, 3);
+    let lhs_stride = vec![m * k, k, 1];
+    let lhs: Vec<f32> = (0..b * m * k).map(|f| f as f32).collect();
+    let rhs_stride = vec![n * k, n, 1];
+    let rhs: Vec<f32> = (0..b * n * k).map(|f| f as f32).collect();
+    let results = run_gemm((b, m, n, k), &lhs, lhs_stride, 0, &rhs, rhs_stride, 0);
+    assert_eq!(
+        approx(results, 4),
+        vec![20.0, 23.0, 26.0, 29.0, 56.0, 68.0, 80.0, 92.0]
+    );
+
+    let (b, m, n, k) = (2, 2, 4, 3);
+    let lhs_stride = vec![m * k, k, 1];
+    let lhs: Vec<f32> = (0..b * m * k).map(|f| f as f32).collect();
+    let rhs_stride = vec![n * k, n, 1];
+    let rhs: Vec<f32> = (0..b * n * k).map(|f| f as f32).collect();
+    let results = run_gemm((b, m, n, k), &lhs, lhs_stride, 0, &rhs, rhs_stride, 0);
+    assert_eq!(
+        approx(results, 4),
+        vec![
+            20.0, 23.0, 26.0, 29.0, 56.0, 68.0, 80.0, 92.0, 344.0, 365.0, 386.0, 407.0, 488.0,
+            518.0, 548.0, 578.0
+        ]
+    );
+
+    // OFFSET
+    let (b, m, n, k) = (2, 2, 4, 3);
+    let lhs_stride = vec![m * k, k, 1];
+    let lhs: Vec<f32> = (0..b * m * k).map(|f| f as f32).collect();
+    let rhs_stride = vec![n * k, n, 1];
+    let rhs: Vec<f32> = (0..b * n * k).map(|f| f as f32).collect();
+    // Manually set batch_size=1 and offset 12 elements * 4 the number of bytes for f32
+    let results = run_gemm((1, m, n, k), &lhs, lhs_stride, 0, &rhs, rhs_stride, 12 * 4);
+    assert_eq!(
+        approx(results, 4),
+        vec![56.0, 59.0, 62.0, 65.0, 200.0, 212.0, 224.0, 236.0]
+    );
+}

candle-metal-kernels/src/unary.metal

@@ -1,4 +1,7 @@
 #include <metal_stdlib>
+#include <metal_math>
+#
+using namespace metal;
 
 METAL_FUNC uint get_strided_index(
     uint idx,
@@ -17,10 +20,44 @@ METAL_FUNC uint get_strided_index(
 
 template <typename T> METAL_FUNC T sqr(T in){ return in * in; }
 template <typename T> METAL_FUNC T neg(T in){ return -in; }
-template <typename T> METAL_FUNC T id(T in){ return in; }
+template <typename T> METAL_FUNC T erf(T in){
+    float x = (float) in;
+    // constants
+    float a1 =  0.254829592;
+    float a2 = -0.284496736;
+    float a3 =  1.421413741;
+    float a4 = -1.453152027;
+    float a5 =  1.061405429;
+    float p  =  0.3275911;
+
+    // Save the sign of x
+    int sign = 1;
+    if (x < 0)
+        sign = -1;
+    x = fabs(x);
+
+    // A&S formula 7.1.26
+    float t = 1.0/(1.0 + p*x);
+    float y = 1.0 - (((((a5*t + a4)*t) + a3)*t + a2)*t + a1)*t*exp(-x*x);
+
+    return T(sign*y);
+}
+template <typename T> METAL_FUNC T id(T in) { return in; }
+template <typename T> METAL_FUNC T gelu_erf(T x) {
+    return T(x * (1 + erf(x * M_SQRT1_2_F)) / 2);
+}
+template <typename T> METAL_FUNC T gelu(T x) {
+    if (x > 5) {
+        return x;
+    }
+    T x_sq = x * x;
+    T x_cube = x_sq * x;
+    T alpha = x + static_cast<T>(0.044715) * x_cube;
+    T beta =  (static_cast<T>(M_2_SQRTPI_F * M_SQRT1_2_F) * alpha);
+    return static_cast<T>(0.5) * x * (static_cast<T>(1.0) + T(tanh(beta)));
+}
 
 
-using namespace metal;
 
 #define UNARY(FN, TYPENAME, FN_NAME, FN_NAME_STRIDED) \
 kernel void FN_NAME( \
@@ -32,7 +69,7 @@ kernel void FN_NAME( \
     if (thread_position_in_grid >= dim) { \
         return; \
     } \
-    output[thread_position_in_grid] = TYPENAME(FN(input[thread_position_in_grid])); \
+    output[thread_position_in_grid] = TYPENAME(FN(float(input[thread_position_in_grid]))); \
 }\
 kernel void FN_NAME_STRIDED( \
     constant size_t &dim, \
@@ -46,7 +83,7 @@ kernel void FN_NAME_STRIDED( \
     if (thread_position_in_grid >= dim) { \
         return; \
     } \
-    output[thread_position_in_grid] = TYPENAME(FN(input[get_strided_index(thread_position_in_grid, num_dims, dims, strides)])); \
+    output[thread_position_in_grid] = TYPENAME(FN(float(input[get_strided_index(thread_position_in_grid, num_dims, dims, strides)]))); \
 }
 
 #define UNARY_OP(NAME) \
@@ -64,8 +101,17 @@ UNARY_OP(sqrt)
 UNARY_OP(neg)
 UNARY_OP(exp)
 UNARY_OP(log)
+UNARY_OP(gelu)
+UNARY_OP(ceil)
+UNARY_OP(floor)
+UNARY_OP(round)
+UNARY_OP(gelu_erf)
+UNARY_OP(erf)
+UNARY_OP(tanh)
 UNARY(id, float, copy_float, copy_float_strided)
 UNARY(id, half, copy_half, copy_half_strided)
+UNARY(id, uint8_t, copy_u8, copy_u8_strided)
+UNARY(id, uint32_t, copy_u32, copy_u32_strided)
 
 #if __METAL_VERSION__ >= 310
 BFLOAT_UNARY_OP(cos)
@@ -75,6 +121,13 @@ BFLOAT_UNARY_OP(sqrt)
 BFLOAT_UNARY_OP(neg)
 BFLOAT_UNARY_OP(exp)
 BFLOAT_UNARY_OP(log)
+BFLOAT_UNARY_OP(gelu)
+BFLOAT_UNARY_OP(ceil)
+BFLOAT_UNARY_OP(floor)
+BFLOAT_UNARY_OP(round)
+BFLOAT_UNARY_OP(gelu_erf)
+BFLOAT_UNARY_OP(erf)
+BFLOAT_UNARY_OP(tanh)
 
 UNARY(id, bfloat, copy_bfloat, copy_bfloat_strided)
 #endif

candle-metal-kernels/tmp/affine.rs

@@ -50,6 +50,7 @@ fn run_affine_bench<T: Clone>(device: &Device, kernels: &Kernels, v: &[T]) {
                 &device,
                 command_buffer,
                 &kernels,
+                "affine_float",
                 v.len(),
                 &input,
                 &mut output,

candle-metal-kernels/tmp/unary.rs

@@ -147,7 +147,7 @@ fn run_unary_bench<T: Clone>(
         println!(
             "{0: <5} | {1: <19} | {2: <6} | {3: <5} | {4: <11?} | {5: <11?}",
             type_name::<T>().split("::").last().unwrap(),
-            kernel_name.to_string(),
+            kernel_name.0,
             v.len(),
             iterations,
             total_time,
@@ -159,7 +159,7 @@ fn run_unary_bench<T: Clone>(
     let shape = vec![2, 5_000];
     let strides = vec![2, 1];
     let offset = 0;
-    for kernel_name in strided {
+    for kernel_name in &strided {
         let total_time = autoreleasepool(|| {
             let command_buffer = command_queue.new_command_buffer();
             let start = Instant::now();
@@ -187,7 +187,7 @@ fn run_unary_bench<T: Clone>(
         println!(
             "{0: <5} | {1: <19} | {2: <6} | {3: <5} | {4: <11?} | {5: <11?}",
             type_name::<T>().split("::").last().unwrap(),
-            kernel_name.to_string(),
+            kernel_name.0,
             v.len(),
             iterations,
             total_time,

candle-nn/Cargo.toml

@@ -19,6 +19,8 @@ num-traits = { workspace = true }
 rayon = { workspace = true }
 safetensors = { workspace = true }
 serde = { workspace = true }
+metal = { workspace = true, optional = true }
+candle-metal-kernels = { path = "../candle-metal-kernels", version = "0.3.0", optional = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
@@ -29,3 +31,4 @@ default = []
 accelerate = ["dep:accelerate-src", "candle/accelerate"]
 cuda = ["candle/cuda"]
 mkl = ["dep:intel-mkl-src", "candle/mkl"]
+metal = ["candle/metal", "dep:candle-metal-kernels", "dep:metal"]

candle-nn/src/ops.rs

@@ -201,6 +201,46 @@ impl candle::CustomOp1 for SoftmaxLastDim {
         };
         Ok((dst, layout.shape().clone()))
     }
+
+    #[cfg(feature = "metal")]
+    fn metal_fwd(
+        &self,
+        storage: &candle::MetalStorage,
+        layout: &Layout,
+    ) -> Result<(candle::MetalStorage, Shape)> {
+        use candle::{backend::BackendStorage, DType};
+        let device = storage.device();
+        let command_buffer = device.command_buffer();
+        let kernels = device.kernels();
+        let name = match storage.dtype() {
+            DType::F32 => "softmax_float",
+            DType::F16 => "softmax_half",
+            DType::BF16 => "softmax_bfloat",
+            dtype => candle::bail!("softmax-last-dim is not implemented for {dtype:?}"),
+        };
+
+        let n = layout.stride().len();
+        if !(layout.is_contiguous() && layout.stride()[n - 1] == 1 && layout.start_offset() == 0) {
+            candle::bail!("Non contiguous softmax-last-dim is not implemented");
+        }
+
+        let last_dim = layout.dims()[layout.shape().rank() - 1];
+        let elem_count = layout.shape().elem_count();
+        let mut output = device.new_buffer(elem_count, storage.dtype(), "softmax");
+        candle_metal_kernels::call_last_softmax(
+            device.metal_device(),
+            &command_buffer,
+            &kernels,
+            name,
+            elem_count,
+            last_dim,
+            storage.buffer(),
+            &mut output,
+        )
+        .unwrap();
+        let newstorage = candle::MetalStorage::new(output, device.clone(), storage.dtype());
+        Ok((newstorage, layout.shape().clone()))
+    }
 }
 
 pub fn softmax_last_dim(xs: &Tensor) -> Result<Tensor> {

candle-transformers/Cargo.toml

@@ -31,3 +31,4 @@ accelerate = ["dep:accelerate-src", "candle/accelerate", "candle-nn/accelerate"]
 cuda = ["candle/cuda", "candle-nn/cuda"]
 flash-attn = ["cuda", "dep:candle-flash-attn"]
 mkl = ["dep:intel-mkl-src", "candle/mkl", "candle-nn/mkl"]
+metal = ["candle/metal", "candle-nn/metal"]

candle-transformers/src/models/mixformer.rs

@@ -142,10 +142,9 @@ impl RotaryEmbedding {
             .to_dtype(DType::F32)?
             .reshape((max_seq_len, 1))?;
         let freqs = t.matmul(&inv_freq)?;
-        Ok(Self {
-            sin: freqs.sin()?,
-            cos: freqs.cos()?,
-        })
+        let sin = freqs.sin()?;
+        let cos = freqs.cos()?;
+        Ok(Self { sin, cos })
     }
 
     fn apply_rotary_emb_qkv(
@@ -408,3 +407,38 @@ impl MixFormerSequentialForCausalLM {
         self.blocks.iter_mut().for_each(|b| b.clear_kv_cache())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_rotary() {
+        let dev = Device::new_metal(0).unwrap();
+        for i in 0..10000 {
+            let dim = 8;
+            let max_seq_len = 12;
+            let inv_freq: Vec<_> = (0..dim)
+                .step_by(2)
+                .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32))
+                .collect();
+            let inv_freq_len = inv_freq.len();
+            let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), &dev).unwrap();
+            let t = Tensor::arange(0u32, max_seq_len as u32, &dev)
+                .unwrap()
+                .to_dtype(DType::F32)
+                .unwrap()
+                .reshape((max_seq_len, 1))
+                .unwrap();
+            let x: f32 = t.i((1, 0)).unwrap().to_scalar().unwrap();
+            assert_eq!(x, 1.0);
+            let x: f32 = inv_freq.i((0, 1)).unwrap().to_scalar().unwrap();
+            assert_eq!(x, 0.1);
+            let freqs = t.matmul(&inv_freq).unwrap();
+            let x: f32 = freqs.i((1, 1)).unwrap().to_scalar().unwrap();
+            assert_eq!(x, 0.1);
+            let sin = freqs.sin().unwrap().contiguous().unwrap();
+            let x: f32 = sin.i((1, 1)).unwrap().to_scalar().unwrap();
+            assert_eq!(x, 0.099833414);
+        }
+    }
+}