Add missing bfloat unary strided kernels

candle-core/benches/benchmarks/matmul.rs

@@ -13,11 +13,11 @@ fn run_bench(c: &mut Criterion, device: &Device) {
     let n = 2048;
     let k = 2048;
 
-    let dtype = DType::F32;
+    let dtype = DType::BF16;
     let lhs = Tensor::zeros((b, m, k), dtype, device).unwrap();
     let rhs = Tensor::zeros((b, n, k), dtype, device).unwrap();
 
-    let flops = b * m * n * k;
+    let flops = b * m * n * k * dtype.size_in_bytes();
 
     let mut group = c.benchmark_group(device.bench_name("matmul"));
     group.throughput(Throughput::Bytes(flops as u64));

candle-core/src/metal_backend/mod.rs

@@ -540,6 +540,7 @@ impl BackendStorage for MetalStorage {
                 ("urelu", DType::F32) => strided::relu::FLOAT,
                 ("uround", DType::F32) => strided::round::FLOAT,
                 ("utanh", DType::F32) => strided::tanh::FLOAT,
+
                 ("ucos", DType::F16) => strided::cos::HALF,
                 ("usin", DType::F16) => strided::sin::HALF,
                 ("usqr", DType::F16) => strided::sqr::HALF,
@@ -557,6 +558,25 @@ impl BackendStorage for MetalStorage {
                 ("urelu", DType::F16) => strided::relu::HALF,
                 ("uround", DType::F16) => strided::round::HALF,
                 ("utanh", DType::F16) => strided::tanh::HALF,
+
+                ("ucos", DType::BF16) => strided::cos::BFLOAT,
+                ("usin", DType::BF16) => strided::sin::BFLOAT,
+                ("usqr", DType::BF16) => strided::sqr::BFLOAT,
+                ("usqrt", DType::BF16) => strided::sqrt::BFLOAT,
+                ("uneg", DType::BF16) => strided::neg::BFLOAT,
+                ("uexp", DType::BF16) => strided::exp::BFLOAT,
+                ("ulog", DType::BF16) => strided::log::BFLOAT,
+                ("ugelu", DType::BF16) => strided::gelu::BFLOAT,
+                ("ugelu_erf", DType::BF16) => strided::gelu_erf::BFLOAT,
+                ("uerf", DType::BF16) => strided::erf::BFLOAT,
+                ("usilu", DType::BF16) => strided::silu::BFLOAT,
+                ("uabs", DType::BF16) => strided::abs::BFLOAT,
+                ("uceil", DType::BF16) => strided::ceil::BFLOAT,
+                ("ufloor", DType::BF16) => strided::floor::BFLOAT,
+                ("urelu", DType::BF16) => strided::relu::BFLOAT,
+                ("uround", DType::BF16) => strided::round::BFLOAT,
+                ("utanh", DType::BF16) => strided::tanh::BFLOAT,
+                
                 (name, dtype) => {
                     crate::bail!("Metal strided unary {name} {dtype:?} not implemented")
                 }
@@ -1235,6 +1255,7 @@ impl BackendStorage for MetalStorage {
         let name = match self.dtype {
             DType::F32 => "sgemm",
             DType::F16 => "hgemm",
+            DType::BF16 => "bgemm",
             dtype => {
                 return Err(MetalError::Message(format!("matmul doesn't support {dtype:?}")).into())
             }

candle-metal-kernels/src/lib.rs

@@ -1451,6 +1451,7 @@ pub fn call_gemm(
     let bytes = match name {
         "sgemm" => 4,
         "hgemm" => 2,
+        "bgemm" => 2,
         other => {
             return Err(MetalKernelError::LoadLibraryError(format!(
                 "{other} is not a valid kernel for gemm"

candle-metal-kernels/src/tests.rs

@@ -1024,7 +1024,20 @@ 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>(
+trait Gemmable: Clone {
+    const gemm_name: &'static str;
+}
+impl Gemmable for f32 {
+    const gemm_name: &'static str = "sgemm";
+}
+impl Gemmable for f16 {
+    const gemm_name: &'static str = "hgemm";
+}
+impl Gemmable for bf16 {
+    const gemm_name: &'static str = "bgemm";
+}
+
+fn run_gemm<T: Gemmable>(
     (b, m, n, k): (usize, usize, usize, usize),
     lhs: &[T],
     lhs_stride: Vec<usize>,
@@ -1033,6 +1046,7 @@ fn run_gemm<T: Clone>(
     rhs_stride: Vec<usize>,
     rhs_offset: usize,
 ) -> Vec<T> {
+
     let device = device();
     let kernels = Kernels::new();
     let command_queue = device.new_command_queue();
@@ -1055,7 +1069,7 @@ fn run_gemm<T: Clone>(
         &device,
         command_buffer,
         &kernels,
-        "sgemm",
+        T::gemm_name,
         (b, m, n, k),
         &lhs_stride,
         lhs_offset,
@@ -1076,23 +1090,23 @@ fn run_gemm<T: Clone>(
 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 lhs: Vec<bf16> = (0..b * m * k).map(|f| bf16::from_f32(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 rhs: Vec<bf16> = (0..b * n * k).map(|f| bf16::from_f32(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),
+        approx_bf16(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 lhs: Vec<bf16> = (0..b * m * k).map(|f| bf16::from_f32(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 rhs: Vec<bf16> = (0..b * n * k).map(|f| bf16::from_f32(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),
+        approx_bf16(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
@@ -1102,13 +1116,13 @@ fn gemm() {
     // 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 lhs: Vec<bf16> = (0..b * m * k).map(|f| bf16::from_f32(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 rhs: Vec<bf16> = (0..b * n * k).map(|f| bf16::from_f32(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),
+        approx_bf16(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

@@ -175,5 +175,5 @@ BFLOAT_UNARY_OP(sign)
 
 UNARY(id, bfloat, copy_bf16, copy_bf16_strided)
 
-COPY2D(copy2d_bf64, bfloat)
+COPY2D(copy2d_bf16, bfloat)
 #endif