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,