Added fill bench

candle-core/Cargo.toml

@@ -49,3 +49,7 @@ metal = ["dep:metal", "dep:candle-metal-kernels"]
 name = "matmul"
 harness = false
 
+[[bench]]
+name = "fill"
+harness = false
+

candle-core/benches/fill.rs

@@ -0,0 +1,57 @@
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(shape: (usize, usize, usize), dtype: DType, device: &Device) {
+    Tensor::ones(shape, dtype, device).unwrap();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let b = 1;
+    let rows = 4096;
+    let columns = 4096;
+
+    let flops = b * rows * columns;
+
+    let device1 = Device::new_metal(0).unwrap();
+    let device2 = device1.clone();
+
+    let mut group = c.benchmark_group("fill_metal_u8");
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |bencher| {
+        bencher.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box((b, rows, columns)), black_box(DType::U8), black_box(&device1));
+            }
+            if let Device::Metal(device) = &device1 {
+                device.wait_until_completed().unwrap();
+            } else {
+                panic!("Expected metal device");
+            }
+            start.elapsed()
+        })
+    });
+    group.finish();
+
+    let mut group = c.benchmark_group("fill_metal_f32");
+    group.throughput(Throughput::Bytes((flops * DType::F32.size_in_bytes()) as u64));
+    group.bench_function("iter", move |bencher| {
+        bencher.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box((b, rows, columns)), black_box(DType::F32), black_box(&device2));
+            }
+            if let Device::Metal(device) = &device2 {
+                device.wait_until_completed().unwrap();
+            } else {
+                panic!("Expected metal device");
+            }
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

candle-core/src/metal_backend.rs

@@ -3,7 +3,7 @@ use crate::conv::{ParamsConv1D, ParamsConv2D, ParamsConvTranspose1D, ParamsConvT
 use crate::op::{BinaryOpT, CmpOp, ReduceOp, UnaryOpT};
 use crate::{CpuStorage, DType, Layout, Result, Shape};
 use candle_metal_kernels;
-use candle_metal_kernels::{CallFill, Fill, Kernels};
+use candle_metal_kernels::{FillOp, Unary, Kernels};
 use half::{bf16, f16};
 use metal;
 use metal::{Buffer, CommandBuffer, CommandQueue, MTLResourceOptions, NSUInteger};
@@ -1405,9 +1405,9 @@ impl BackendDevice for MetalDevice {
         let command_buffer = self.command_buffer()?;
         command_buffer.set_label("zeros");
 
-        // This assumes the specific zero type DType is equal to 0x00u8
+        // This assumes the zero value of this DType is equal to 0x00u8
         // (which is true for all current types)
-        Fill::call_fill(
+        Unary::fill(
             &self.device,
             &command_buffer,
             &self.kernels,
@@ -1421,13 +1421,13 @@ impl BackendDevice for MetalDevice {
     }
 
     fn ones_impl(&self, shape: &Shape, dtype: DType) -> Result<Self::Storage> {
-        let buffer = self.new_buffer(shape.elem_count(), dtype, "zeros")?;
+        let buffer = self.new_buffer(shape.elem_count(), dtype, "ones")?;
         let command_buffer = self.command_buffer()?;
         command_buffer.set_label("ones");
 
         macro_rules! fill {
             ($value:expr) => {
-                Fill::call_fill(
+                Unary::fill(
                     &self.device,
                     &command_buffer,
                     &self.kernels,

candle-metal-kernels/Cargo.toml

@@ -19,3 +19,8 @@ num-traits = "0.2.17"
 
 [dev-dependencies]
 rand = "0.8.5"
+criterion = "0.5.1"
+
+[[bench]]
+name = "fill"
+harness = false

candle-metal-kernels/src/lib.rs

@@ -1573,19 +1573,19 @@ pub fn call_upsample_nearest_2d(
     Ok(())
 }
 
-#[inline]
+#[inline(always)]
 fn divide(m: usize, b: usize) -> NSUInteger {
     ((m + b - 1) / b) as NSUInteger
 }
 
-pub struct Fill<T> {
+pub struct Unary<T> {
     _marker: PhantomData<T>,
 }
 
-pub trait CallFill<T> {
-    const KERNEL_NAME: &'static str;
+pub trait FillOp<T> {
+    const FILL_KERNEL: &'static str;
 
-    fn call_fill(
+    fn fill(
         device: &Device,
         command_buffer: &CommandBufferRef,
         kernels: &Kernels,
@@ -1598,66 +1598,12 @@ pub trait CallFill<T> {
 macro_rules ! impl_call_fill {
     ($($t:ty),*) => {
         $(
-            impl CallFill<$t> for Fill<$t> {
-                const KERNEL_NAME: &'static str = concat!("fill_", stringify!($t));
+            impl FillOp<$t> for Unary<$t> {
+                const FILL_KERNEL: &'static str = concat!("fill_", stringify!($t));
 
-                fn call_fill(device: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, elem_count: usize, buffer: &Buffer, value: $t) -> Result<(), MetalKernelError> {
-                    _call_fill(device, command_buffer, kernels, Self::KERNEL_NAME, elem_count, buffer, value)
-                }
-            }
-        )*
-    };
-}
-impl_call_fill!(u32, i64, f16, bf16, f32);
-
-impl CallFill<u8> for Fill<u8> {
-    const KERNEL_NAME: &'static str = "";
-
-    fn call_fill(
-        _: &Device,
-        command_buffer: &CommandBufferRef,
-        kernels: &Kernels,
-        elem_count: usize,
-        buffer: &Buffer,
-        value: u8,
-    ) -> Result<(), MetalKernelError> {
-        _call_blit_fill(command_buffer, kernels, elem_count, buffer, value)
-    }
-}
-
-fn _call_blit_fill(
-    command_buffer: &CommandBufferRef,
-    kernels: &Kernels,
-    elem_count: usize,
-    buffer: &Buffer,
-    value: u8,
-) -> Result<(), MetalKernelError> {
-    let blit = command_buffer.new_blit_command_encoder();
-    blit.wait_for_fence(&kernels.fence);
-    blit.fill_buffer(
-        &buffer,
-        metal::NSRange {
-            location: 0,
-            length: elem_count as NSUInteger,
-        },
-        value,
-    );
-    blit.update_fence(&kernels.fence);
-    blit.end_encoding();
-
-    Ok(())
-}
-
-fn _call_fill<D: EncoderParam>(
-    device: &Device,
-    command_buffer: &CommandBufferRef,
-    kernels: &Kernels,
-    kernel_name: &'static str,
-    elem_count: usize,
-    buffer: &Buffer,
-    value: D,
-) -> Result<(), MetalKernelError> {
-    let pipeline = kernels.load_pipeline(device, Source::Fill, kernel_name)?;
+                #[inline(always)]
+                fn fill(device: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, elem_count: usize, buffer: &Buffer, value: $t) -> Result<(), MetalKernelError> {
+                    let pipeline = kernels.load_pipeline(device, Source::Fill, Self::FILL_KERNEL)?;
                     let encoder = command_buffer.new_compute_command_encoder();
                     encoder.wait_for_fence(&kernels.fence);
                     encoder.set_compute_pipeline_state(&pipeline);
@@ -1672,6 +1618,40 @@ fn _call_fill<D: EncoderParam>(
                     encoder.end_encoding();
 
                     Ok(())
+                }
+            }
+        )*
+    };
+}
+impl_call_fill!(u32, i64, f16, bf16, f32);
+
+impl FillOp<u8> for Unary<u8> {
+    const FILL_KERNEL: &'static str = "";
+
+    #[inline(always)]
+    fn fill(
+        _: &Device,
+        command_buffer: &CommandBufferRef,
+        kernels: &Kernels,
+        elem_count: usize,
+        buffer: &Buffer,
+        value: u8,
+    ) -> Result<(), MetalKernelError> {
+        let blit = command_buffer.new_blit_command_encoder();
+        blit.wait_for_fence(&kernels.fence);
+        blit.fill_buffer(
+            &buffer,
+            metal::NSRange {
+                location: 0,
+                length: elem_count as NSUInteger,
+            },
+            value,
+        );
+        blit.update_fence(&kernels.fence);
+        blit.end_encoding();
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

candle-metal-kernels/src/tests.rs

@@ -808,7 +808,7 @@ fn gemm() {
 
 fn run_fill<T: EncoderParam + Clone>(elem_count: usize, value: T) -> Vec<T>
 where
-    Fill<T>: CallFill<T>,
+    Unary<T>: FillOp<T>,
 {
     let device = device();
     let fence = device.new_fence();
@@ -816,7 +816,7 @@ where
     let command_queue = device.new_command_queue();
     let command_buffer = command_queue.new_command_buffer();
     let buffer = new_buffer(&device, &vec![0.0f32; elem_count]);
-    Fill::<T>::call_fill(
+    Unary::<T>::fill(
         &device,
         command_buffer,
         &kernels,
@@ -835,7 +835,7 @@ where
 fn fill() {
     fn assert_fill<T: EncoderParam + Copy + std::fmt::Debug + PartialEq>(value: T)
     where
-        Fill<T>: CallFill<T>,
+        Unary<T>: FillOp<T>,
     {
         for i in 0..4 {
             assert_eq!(run_fill(8 ^ i, value), vec![value; 8 ^ i]);