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<()> {
+    fn set_seed(&self, seed: u64) -> Result<()> {
-        crate::bail!("set_seed")
+        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,
+        min: f64,
-        stddev: f64,
+        max: f64,
     ) -> Result<Self::Storage> {
-        // TODO is there a better way ?
+        let name = match dtype {
-        let cpu_storage = crate::cpu_backend::CpuDevice.rand_uniform(shape, dtype, mean, stddev)?;
+            DType::F32 => "rand_uniform_f32",
-        self.storage_from_cpu_storage(&cpu_storage)
+            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(

candle-metal-kernels/src/lib.rs

@@ -12,8 +12,9 @@ const UNARY: &str = include_str!("unary.metal");
 const BINARY: &str = include_str!("binary.metal");
 const TERNARY: &str = include_str!("ternary.metal");
 const CAST: &str = include_str!("cast.metal");
-const REDUCE: &str = include_str!("reduce.metal");
 const CONV: &str = include_str!("conv.metal");
+const REDUCE: &str = include_str!("reduce.metal");
+const RANDOM: &str = include_str!("random.metal");
 const MFA: &[u8] = include_bytes!("libMetalFlashAttention.metallib");
 
 /// Most kernels apply similarly across the tensors
@@ -45,7 +46,7 @@ fn set_param<P: EncoderParam>(encoder: &ComputeCommandEncoderRef, position: u64,
 /// Helper functions to create the various objects on the compute command encoder
 /// on a single line.
 /// Prevents getting wrong some arguments number and mixing length and size in bytes.
-trait EncoderParam {
+pub trait EncoderParam {
     fn set_param(encoder: &ComputeCommandEncoderRef, position: u64, data: Self);
 }
 macro_rules! primitive {
@@ -61,8 +62,10 @@ macro_rules! primitive {
         }
     };
 }
+primitive!(bool);
 primitive!(usize);
 primitive!(u32);
+primitive!(u64);
 primitive!(f32);
 
 impl<T> EncoderParam for &[T] {
@@ -117,6 +120,7 @@ pub enum Source {
     Reduce,
     Mfa,
     Conv,
+    Random,
 }
 
 macro_rules! ops{
@@ -228,6 +232,7 @@ impl Kernels {
             Source::Cast => CAST,
             Source::Reduce => REDUCE,
             Source::Conv => CONV,
+            Source::Random => RANDOM,
             Source::Mfa => panic!("Invalid lib"),
         }
     }
@@ -1566,5 +1571,69 @@ fn divide(m: usize, b: usize) -> NSUInteger {
     ((m + b - 1) / b) as NSUInteger
 }
 
+#[allow(clippy::too_many_arguments)]
+pub fn call_random_uniform(
+    device: &Device,
+    command_buffer: &CommandBufferRef,
+    kernels: &Kernels,
+    name: &'static str,
+    seed: u64,
+    min: f32,
+    max: f32,
+    length: usize,
+    buffer: &Buffer,
+) -> Result<(), MetalKernelError> {
+    if min >= max {
+        return Err(MetalKernelError::LoadLibraryError(
+            "min must be less than max".to_string(),
+        ));
+    }
+
+    let size: usize = match name {
+        "rand_uniform_f32" => 4,
+        "rand_uniform_f16" | "rand_uniform_bf16" => 2,
+        _ => Err(MetalKernelError::LoadLibraryError(format!(
+            "{name} is not a valid kernel for random"
+        )))?,
+    };
+
+    let elems_per_key = length;
+    let bytes_per_key = size * elems_per_key;
+
+    let out_per_key = (bytes_per_key + 4 - 1) / 4;
+    let half_size = out_per_key / 2;
+    let odd = length % 2 != 0;
+
+    let pipeline = kernels.load_pipeline(device, Source::Random, name)?;
+    let encoder = command_buffer.new_compute_command_encoder();
+
+    let thread_group_count = MTLSize {
+        width: length as u64,
+        height: half_size as u64 + odd as u64,
+        depth: 1,
+    };
+    let threads = std::cmp::min(
+        (half_size + odd as usize) as NSUInteger,
+        pipeline.max_total_threads_per_threadgroup(),
+    );
+    let thread_group_size = MTLSize {
+        width: threads,
+        height: 1,
+        depth: 1,
+    };
+
+    encoder.wait_for_fence(&kernels.fence);
+    encoder.set_compute_pipeline_state(&pipeline);
+
+    set_params!(encoder, (length, seed, min, max, buffer));
+
+    encoder.use_resource(buffer, metal::MTLResourceUsage::Write);
+    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);
+    encoder.update_fence(&kernels.fence);
+    encoder.end_encoding();
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests;

candle-metal-kernels/src/random.metal

@@ -0,0 +1,139 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// Constants
+// 2^32 and 1/2^32. Useful for converting between float and uint.
+static constexpr constant ulong UNIF01_NORM32 = 4294967296;
+static constexpr constant float UNIF01_INV32 = 2.328306436538696289e-10;
+// 2 * pi
+static constexpr constant float TWO_PI = 2.0 * M_PI_F;
+static constexpr constant int3 S1 = {13, 19, 12};
+static constexpr constant int3 S2 = {2, 25, 4};
+static constexpr constant int3 S3 = {3, 11, 17};
+
+static constexpr constant uint64_t PHI[16] = {
+    0x9E3779B97F4A7C15,
+    0xF39CC0605CEDC834,
+    0x1082276BF3A27251,
+    0xF86C6A11D0C18E95,
+    0x2767F0B153D27B7F,
+    0x0347045B5BF1827F,
+    0x01886F0928403002,
+    0xC1D64BA40F335E36,
+    0xF06AD7AE9717877E,
+    0x85839D6EFFBD7DC6,
+    0x64D325D1C5371682,
+    0xCADD0CCCFDFFBBE1,
+    0x626E33B8D04B4331,
+    0xBBF73C790D94F79D,
+    0x471C4AB3ED3D82A5,
+    0xFEC507705E4AE6E5,
+};
+
+// Combined Tausworthe and LCG Random Number Generator.
+// https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-37-efficient-random-number-generation-and-application
+// https://indico.cern.ch/event/93877/contributions/2118070/attachments/1104200/1575343/acat3_revised_final.pdf
+class HybridTaus {
+private:
+    thread float seed;
+
+    // Generate seeds for each thread.
+    thread uint4 seed_per_thread(const ulong4 seeds) {
+        return uint4(ulong4(seeds) * ulong4(PHI[0], PHI[1], PHI[2], PHI[3]) * ulong4(1099087573UL));
+    }
+
+    // Tausworthe generator.
+    thread uint taus(const uint z, const int3 s, const uint M) {
+        uint b = (((z << s.x) ^ z) >> s.y);
+        return (((z & M) << s.z) ^ b);
+    }
+
+    // LCG generator.
+    thread uint lcg(const uint z) {
+        return (1664525 * z + 1013904223UL);
+    }
+
+public:
+    thread HybridTaus(const ulong4 seeds) {
+        uint4 seed = this->seed_per_thread(seeds);
+
+        // Seed #1
+        uint z1 = taus(seed.x, S1, 4294967294UL);
+        uint z2 = taus(seed.y, S2, 4294967288UL);
+        uint z3 = taus(seed.z, S3, 4294967280UL);
+        uint z4 = lcg(seed.x);
+
+        // Seed #2
+        uint r1 = (z1^z2^z3^z4^seed.y);
+        z1 = taus(r1, S1, 429496729UL);
+        z2 = taus(r1, S2, 4294967288UL);
+        z3 = taus(r1, S3, 429496280UL);
+        z4 = lcg(r1);
+
+        // Seed #3
+        r1 = (z1^z2^z3^z4^seed.z);
+        z1 = taus(r1, S1, 429496729UL);
+        z2 = taus(r1, S2, 4294967288UL);
+        z3 = taus(r1, S3, 429496280UL);
+        z4 = lcg(r1);
+
+        // Seed #4
+        r1 = (z1^z2^z3^z4^seed.w);
+        z1 = taus(r1, S1, 429496729UL);
+        z2 = taus(r1, S2, 4294967288UL);
+        z3 = taus(r1, S3, 429496280UL);
+        z4 = lcg(r1);
+
+        this->seed = (z1^z2^z3^z4) * UNIF01_INV32;
+    }
+
+    thread float rand() {
+        uint seed = this->seed * UNIF01_NORM32;
+        uint z1 = taus(seed, S1, 429496729UL);
+        uint z2 = taus(seed, S2, 4294967288UL);
+        uint z3 = taus(seed, S3, 429496280UL);
+        uint z4 = lcg(seed);
+
+        thread float old_seed = this->seed;
+        this->seed = (z1^z2^z3^z4) * UNIF01_INV32;
+        return old_seed;
+    }
+};
+
+template<typename T> METAL_FUNC void rand_uniform(
+    constant size_t &elem_count,
+    constant ulong &seed,
+    constant float &min,
+    constant float &max,
+    device T *out,
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= elem_count) {
+        return;
+    }
+    float diff = max - min;
+    HybridTaus rng = HybridTaus({seed, tid, 1, 1});
+    out[tid] = static_cast<T>(rng.rand() * diff + min);
+}
+
+#define UNIFORM_OP(NAME, T)                                 \
+kernel void rand_uniform_##NAME(                            \
+    constant size_t &elem_count,                            \
+    constant ulong &seed,                                   \
+    constant float &min,                                    \
+    constant float &max,                                    \
+    device T *out,                                          \
+    uint tid [[thread_position_in_grid]]                    \
+) {                                                         \
+    rand_uniform<T>(elem_count, seed, min, max, out, tid);  \
+}                                                           \
+
+#define RANDOM_OPS(NAME, T) \
+UNIFORM_OP(NAME, T)         \
+
+RANDOM_OPS(f32, float)
+RANDOM_OPS(f16, half)
+
+#if __METAL_VERSION__ >= 310
+RANDOM_OPS(bf16, bfloat)
+#endif

candle-metal-kernels/src/tests.rs

@@ -11,7 +11,7 @@ fn read_to_vec<T: Clone>(buffer: &Buffer, n: usize) -> Vec<T> {
 
 fn new_buffer<T>(device: &Device, data: &[T]) -> Buffer {
     let options = MTLResourceOptions::StorageModeManaged;
-    let ptr = data.as_ptr() as *const core::ffi::c_void;
+    let ptr = data.as_ptr() as *const c_void;
     let size = (data.len() * std::mem::size_of::<T>()) as u64;
     device.new_buffer_with_data(ptr, size, options)
 }
@@ -590,7 +590,6 @@ fn softmax() {
     }
     let results = run_softmax(&v, last_dim, "softmax_f32");
     let results = approx(results, 4);
-    println!("{results:?}");
     assert_eq!(
         results.iter().map(|&s| s.round() as usize).sum::<usize>(),
         n
@@ -806,3 +805,56 @@ fn gemm() {
         vec![56.0, 59.0, 62.0, 65.0, 200.0, 212.0, 224.0, 236.0]
     );
 }
+
+fn run_random<T: Clone>(seed: u64, shape: &[usize], name: &'static str, min: f32, max: f32) -> Vec<T> {
+    let device = device();
+    let fence = device.new_fence();
+    let kernels = Kernels::new(fence);
+    let command_queue = device.new_command_queue();
+    let command_buffer = command_queue.new_command_buffer();
+    let options = MTLResourceOptions::StorageModeManaged;
+    let length = shape.iter().product::<usize>();
+    let output = device.new_buffer((length * core::mem::size_of::<T>()) as u64, options);
+
+    call_random_uniform(
+        &device,
+        command_buffer,
+        &kernels,
+        name,
+        seed,
+        min,
+        max,
+        length,
+        &output,
+    )
+    .unwrap();
+
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    read_to_vec(&output, length)
+}
+
+#[test]
+fn random() {
+    use std::fs::File;
+    use std::io::prelude::*;
+
+    let shape = vec![1024, 4];
+    let seed = 299792458;
+    let min = -30.0;
+    let max = 30.0;
+    let results = run_random::<f32>(seed, &shape, "rand_uniform_f32", min, max);
+    for &v in &results {
+        assert!(v >= min && v <= max);
+    }
+
+    // Writing bytes to file for testing with ENT
+    // https://www.fourmilab.ch/random/
+    // TODO: Remove before merge
+    let (head, body, tail) = unsafe { results.align_to::<u8>() };
+    assert!(head.is_empty());
+    assert!(tail.is_empty());
+    let mut file = File::create("test").unwrap();
+    file.write_all(body).unwrap();
+}