Implement hybrid Tausworthe + LCG psuedo random number generator in metal

candle-core/Cargo.toml

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

candle-core/benches/random.rs

@@ -0,0 +1,41 @@
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(a: &Tensor) {
+    a.rand_like(0.0, 1.0).unwrap();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let b = 1;
+
+    let rows = 2048;
+    let cols = 2048;
+
+    let device = Device::new_metal(0).unwrap();
+    let dtype = DType::F32;
+    let tensor = Tensor::zeros((b, rows, cols), dtype, &device).unwrap();
+
+    let flops = b * rows * cols;
+
+    let mut group = c.benchmark_group("random_metal");
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |benches| {
+        benches.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&tensor));
+            }
+            if let Device::Metal(device) = &device {
+                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

@@ -8,7 +8,7 @@ use metal;
 use metal::{Buffer, CommandBuffer, CommandQueue, MTLResourceOptions, NSUInteger};
 use std::collections::HashMap;
 use std::path::Path;
-use std::sync::{Arc, RwLock, TryLockError};
+use std::sync::{Arc, Mutex, RwLock, TryLockError};
 
 /// Simple way to catch lock error without
 /// depending on T
@@ -106,6 +106,8 @@ pub struct MetalDevice {
     /// Whenever we actually allocate a new buffer, we make a full sweep to cleanup unused buffers
     /// (strong_count = 1).
     buffers: AllocatedBuffers,
+
+    seed: Arc<Mutex<u64>>,
 }
 
 impl std::fmt::Debug for MetalDevice {
@@ -1373,6 +1375,7 @@ impl BackendDevice for MetalDevice {
             Ok(val) => val.parse()?,
             _ => 20,
         };
+        let seed = Arc::new(Mutex::new(299792458));
         Ok(Self {
             device,
             fence,
@@ -1382,11 +1385,14 @@ impl BackendDevice for MetalDevice {
             compute_per_buffer,
             buffers,
             kernels,
+            seed
         })
     }
 
-    fn set_seed(&self, _seed: u64) -> Result<()> {
-        crate::bail!("set_seed")
+    fn set_seed(&self, seed: u64) -> Result<()> {
+        let mut s = self.seed.try_lock().map_err(MetalError::from)?;
+        *s = seed;
+        Ok(())
     }
 
     fn location(&self) -> crate::DeviceLocation {
@@ -1441,12 +1447,30 @@ impl BackendDevice for MetalDevice {
         &self,
         shape: &Shape,
         dtype: DType,
-        mean: f64,
-        stddev: f64,
+        min: f64,
+        max: f64,
     ) -> Result<Self::Storage> {
-        // TODO is there a better way ?
-        let cpu_storage = crate::cpu_backend::CpuDevice.rand_uniform(shape, dtype, mean, stddev)?;
-        self.storage_from_cpu_storage(&cpu_storage)
+        let name = match dtype {
+            DType::F32 => "rand_uniform_f32",
+            DType::F16 => "rand_uniform_f16",
+            DType::BF16 => "rand_uniform_bf16",
+            dtype => crate::bail!("rand_uniform not implemented for {dtype:?}"),
+        };
+        let buffer = self.new_buffer(shape.elem_count(), dtype, "rand_uniform")?;
+        let command_buffer = self.command_buffer()?;
+        candle_metal_kernels::call_random_uniform(
+            &self.device,
+            &command_buffer,
+            &self.kernels,
+            name,
+            *self.seed.lock().unwrap(),
+            min as f32,
+            max as f32,
+            shape.elem_count(),
+            &buffer
+        ).map_err(MetalError::from)?;
+
+        Ok(Self::Storage::new(buffer, self.clone(), dtype))
     }
 
     fn rand_normal(