Efficient implementation of Tensor::ones() for metal (#2512)

* WIP: hopefully better const impl

* with GPU

* More tests on

* Reverting primitive for

* Incorporating review changes - added check elem count check in kerner, using  for call strategy

* rustfmt ran

candle-core/src/metal_backend/mod.rs

@@ -1917,10 +1917,38 @@ impl BackendDevice for MetalDevice {
         ))
     }
 
-    fn ones_impl(&self, shape: &Shape, dtype: DType) -> Result<Self::Storage> {
-        // TODO Is there a faster way ?
-        let cpu_storage = crate::cpu_backend::CpuDevice.ones_impl(shape, dtype)?;
-        self.storage_from_cpu_storage(&cpu_storage)
+    fn ones_impl(&self, shape: &Shape, dtype: DType) -> Result<MetalStorage> {
+        let name = match dtype {
+            DType::U8 => "fill_u8",
+            DType::U32 => "fill_u32",
+            DType::I64 => "fill_i64",
+            DType::F16 => "fill_f16",
+            DType::BF16 => "fill_bf16",
+            DType::F32 => "fill_f32",
+            DType::F64 => {
+                let cpu_storage = crate::cpu_backend::CpuDevice.ones_impl(shape, dtype)?;
+                return self.storage_from_cpu_storage(&cpu_storage);
+            }
+        };
+        let buffer = self.new_buffer(shape.elem_count(), dtype, "alloc-ones")?;
+        let command_buffer = self.command_buffer()?;
+        candle_metal_kernels::call_const_fill(
+            &self.device,
+            &command_buffer,
+            &self.kernels,
+            name,
+            shape.elem_count(),
+            &buffer,
+            1.,
+        )
+        .map_err(MetalError::from)?;
+
+        Ok(MetalStorage::new(
+            buffer,
+            self.clone(),
+            shape.elem_count(),
+            dtype,
+        ))
     }
 
     fn storage_from_slice<T: crate::WithDType>(&self, s: &[T]) -> Result<Self::Storage> {

candle-core/tests/tensor_tests.rs

@@ -29,6 +29,36 @@ fn ones(device: &Device) -> Result<()> {
         Tensor::ones((2, 3), DType::F64, device)?.to_vec2::<f64>()?,
         [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
     );
+    assert_eq!(
+        Tensor::ones((2, 3), DType::F16, device)?.to_vec2::<half::f16>()?,
+        [
+            [
+                half::f16::from_f32(1.0),
+                half::f16::from_f32(1.0),
+                half::f16::from_f32(1.0)
+            ],
+            [
+                half::f16::from_f32(1.0),
+                half::f16::from_f32(1.0),
+                half::f16::from_f32(1.0)
+            ]
+        ],
+    );
+    assert_eq!(
+        Tensor::ones((2, 3), DType::BF16, device)?.to_vec2::<half::bf16>()?,
+        [
+            [
+                half::bf16::from_f32(1.0),
+                half::bf16::from_f32(1.0),
+                half::bf16::from_f32(1.0)
+            ],
+            [
+                half::bf16::from_f32(1.0),
+                half::bf16::from_f32(1.0),
+                half::bf16::from_f32(1.0)
+            ]
+        ],
+    );
     Ok(())
 }
 

candle-metal-kernels/src/fill.metal

@@ -0,0 +1,39 @@
+#include <metal_stdlib>
+
+using namespace metal;
+
+template<typename T> METAL_FUNC void fill_with(
+    device T *out,
+    constant float &value,
+    constant size_t &numel,
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= numel) {
+        return;
+    }
+    out[tid] = static_cast<T>(value);
+}
+
+#define FILL_OP(NAME, T)                                \
+kernel void fill_##NAME(                                \
+    device T *out,                                      \
+    constant float &value,                              \
+    constant size_t &numel,                              \
+    uint tid [[thread_position_in_grid]]                \
+) {                                                     \
+    fill_with<T>(out, value, numel, tid);              \
+}                                                       \
+
+
+#define FILL_OPS(NAME, T) \
+FILL_OP(NAME, T)          \
+
+FILL_OPS(u8, uchar)
+FILL_OPS(u32, uint)
+FILL_OPS(i64, long)
+FILL_OPS(f16, half)
+FILL_OPS(f32, float)
+
+#if __METAL_VERSION__ >= 310
+FILL_OPS(bf16, bfloat)
+#endif

candle-metal-kernels/src/lib.rs

@@ -14,6 +14,7 @@ const AFFINE: &str = include_str!("affine.metal");
 const BINARY: &str = include_str!("binary.metal");
 const CAST: &str = include_str!("cast.metal");
 const CONV: &str = include_str!("conv.metal");
+const FILL: &str = include_str!("fill.metal");
 const INDEXING: &str = include_str!("indexing.metal");
 // Current source: https://github.com/ivarflakstad/metal-flash-attention/tree/candle
 const MFA: &[u8] = include_bytes!("libMetalFlashAttention.metallib");
@@ -31,6 +32,7 @@ pub enum Source {
     Binary,
     Cast,
     Conv,
+    Fill,
     Gemm,
     Indexing,
     Mfa,
@@ -196,6 +198,7 @@ impl Kernels {
             Source::Binary => BINARY,
             Source::Cast => CAST,
             Source::Conv => CONV,
+            Source::Fill => FILL,
             Source::Gemm => MLX_GEMM,
             Source::Indexing => INDEXING,
             Source::Quantized => QUANTIZED,
@@ -2357,5 +2360,30 @@ pub fn call_mlx_gemm(
     Ok(())
 }
 
+pub fn call_const_fill(
+    device: &Device,
+    ep: impl EncoderProvider,
+    kernels: &Kernels,
+    name: &'static str,
+    length: usize,
+    output: &Buffer,
+    v: f32,
+) -> Result<(), MetalKernelError> {
+    let pipeline = kernels.load_pipeline(device, Source::Fill, name)?;
+    let encoder = ep.encoder();
+    let encoder: &ComputeCommandEncoderRef = encoder.as_ref();
+
+    encoder.set_compute_pipeline_state(&pipeline);
+
+    set_params!(encoder, (output, v, length));
+
+    let (thread_group_count, thread_group_size) = linear_split(&pipeline, length);
+
+    encoder.use_resource(output, metal::MTLResourceUsage::Write);
+    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests;

candle-metal-kernels/src/tests.rs

@@ -1,6 +1,7 @@
 use super::*;
 use half::{bf16, f16};
 use metal::MTLResourceOptions;
+use rand::Rng;
 
 fn read_to_vec<T: Clone>(buffer: &Buffer, n: usize) -> Vec<T> {
     let ptr = buffer.contents() as *const T;
@@ -2307,3 +2308,67 @@ fn conv_transpose1d_u32() {
     let expected = vec![1, 4, 10, 20, 25, 24, 16];
     assert_eq!(results, expected);
 }
+
+fn constant_fill<T: Clone>(name: &'static str, len: usize, value: f32) -> Vec<T> {
+    let dev = device();
+    let kernels = Kernels::new();
+    let command_queue = dev.new_command_queue();
+    let command_buffer = command_queue.new_command_buffer();
+
+    let buffer = dev.new_buffer(
+        (len * std::mem::size_of::<T>()) as u64,
+        MTLResourceOptions::StorageModePrivate,
+    );
+
+    call_const_fill(&dev, command_buffer, &kernels, name, len, &buffer, value).unwrap();
+
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    read_to_vec::<T>(&buffer, len)
+}
+
+#[test]
+fn const_fill() {
+    let fills = [
+        "fill_u8",
+        "fill_u32",
+        "fill_i64",
+        "fill_f16",
+        "fill_bf16",
+        "fill_f32",
+    ];
+
+    for name in fills {
+        let len = rand::thread_rng().gen_range(2..16) * rand::thread_rng().gen_range(4..16);
+        let value = rand::thread_rng().gen_range(1. ..19.);
+
+        match name {
+            "fill_u8" => {
+                let v = constant_fill::<u8>(name, len, value);
+                assert_eq!(v, vec![value as u8; len])
+            }
+            "fill_u32" => {
+                let v = constant_fill::<u32>(name, len, value);
+                assert_eq!(v, vec![value as u32; len])
+            }
+            "fill_i64" => {
+                let v = constant_fill::<i64>(name, len, value);
+                assert_eq!(v, vec![value as i64; len])
+            }
+            "fill_f16" => {
+                let v = constant_fill::<f16>(name, len, value);
+                assert_eq!(v, vec![f16::from_f32(value); len])
+            }
+            "fill_bf16" => {
+                let v = constant_fill::<bf16>(name, len, value);
+                assert_eq!(v, vec![bf16::from_f32(value); len])
+            }
+            "fill_f32" => {
+                let v = constant_fill::<f32>(name, len, value);
+                assert_eq!(v, vec![value; len])
+            }
+            _ => unimplemented!(),
+        };
+    }
+}