Merge pull request #1533 from huggingface/ivarflakstad/metal-prng

3 files changed, 402 insertions, 4 deletions

diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs
index 201af97e..fe969372 100644
--- a/candle-metal-kernels/src/lib.rs
+++ b/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");
 const QUANTIZED: &str = include_str!("quantized.metal");
 
@@ -62,10 +63,12 @@ macro_rules! primitive {
         }
     };
 }
+primitive!(bool);
 primitive!(usize);
-primitive!(i64);
 primitive!(i32);
+primitive!(i64);
 primitive!(u32);
+primitive!(u64);
 primitive!(f32);
 
 impl<T> EncoderParam for &[T] {
@@ -120,6 +123,7 @@ pub enum Source {
     Reduce,
     Mfa,
     Conv,
+    Random,
     Quantized,
 }
 
@@ -241,6 +245,7 @@ impl Kernels {
             Source::Cast => CAST,
             Source::Reduce => REDUCE,
             Source::Conv => CONV,
+            Source::Random => RANDOM,
             Source::Quantized => QUANTIZED,
             Source::Mfa => panic!("Invalid lib"),
         }
@@ -1527,6 +1532,73 @@ pub fn call_upsample_nearest_2d(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
+pub fn call_random_uniform(
+    device: &Device,
+    command_buffer: &CommandBufferRef,
+    kernels: &Kernels,
+    name: &'static str,
+    min: f32,
+    max: f32,
+    length: usize,
+    seed: &Buffer,
+    buffer: &Buffer,
+) -> Result<(), MetalKernelError> {
+    if min >= max {
+        return Err(MetalKernelError::LoadLibraryError(
+            "min must be less than max".to_string(),
+        ));
+    }
+    let pipeline = kernels.load_pipeline(device, Source::Random, name)?;
+    let encoder = command_buffer.new_compute_command_encoder();
+
+    let odd = (length % 2 != 0) as usize;
+    let (thread_group_count, thread_group_size) = linear_split(&pipeline, length / 2 + odd);
+
+    encoder.set_compute_pipeline_state(&pipeline);
+
+    set_params!(encoder, (length, min, max, seed, buffer));
+
+    encoder.use_resource(seed, metal::MTLResourceUsage::Read);
+    encoder.use_resource(seed, metal::MTLResourceUsage::Write);
+    encoder.use_resource(buffer, metal::MTLResourceUsage::Write);
+    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);
+    encoder.end_encoding();
+
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn call_random_normal(
+    device: &Device,
+    command_buffer: &CommandBufferRef,
+    kernels: &Kernels,
+    name: &'static str,
+    mean: f32,
+    stddev: f32,
+    length: usize,
+    seed: &Buffer,
+    buffer: &Buffer,
+) -> Result<(), MetalKernelError> {
+    let pipeline = kernels.load_pipeline(device, Source::Random, name)?;
+    let encoder = command_buffer.new_compute_command_encoder();
+
+    let odd = (length % 2 != 0) as usize;
+    let (thread_group_count, thread_group_size) = linear_split(&pipeline, length / 2 + odd);
+
+    encoder.set_compute_pipeline_state(&pipeline);
+
+    set_params!(encoder, (length, mean, stddev, seed, buffer));
+
+    encoder.use_resource(seed, metal::MTLResourceUsage::Read);
+    encoder.use_resource(seed, metal::MTLResourceUsage::Write);
+    encoder.use_resource(buffer, metal::MTLResourceUsage::Write);
+    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);
+    encoder.end_encoding();
+
+    Ok(())
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum GgmlDType {
     Q4_0,
diff --git a/candle-metal-kernels/src/random.metal b/candle-metal-kernels/src/random.metal
new file mode 100644
index 00000000..a7e48393
--- /dev/null
+++ b/candle-metal-kernels/src/random.metal
@@ -0,0 +1,206 @@
+#include <metal_stdlib>
+#include <metal_integer>
+#include <metal_atomic>
+
+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};
+
+// Used to prevent bad seeds.
+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
+struct HybridTaus {
+
+    float state;
+
+    HybridTaus() thread = default;
+    HybridTaus() threadgroup = default;
+    HybridTaus() device = default;
+    HybridTaus() constant = default;
+
+    // Generate seeds for each thread.
+    METAL_FUNC static uint4 seed_per_thread(const ulong4 seeds) {
+        return uint4(ulong4(seeds) * ulong4(PHI[0], PHI[1], PHI[2], PHI[3]) * ulong4(1099087573UL));
+    }
+
+    // Tausworthe generator.
+    METAL_FUNC static 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.
+    METAL_FUNC static uint lcg(const uint z) {
+        return (1664525 * z + 1013904223UL);
+    }
+
+    // Initialize the RNG state.
+    METAL_FUNC static HybridTaus init(const ulong4 seeds) {
+        uint4 seed = 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);
+
+        HybridTaus rng;
+        rng.state = (z1^z2^z3^z4) * UNIF01_INV32;
+        return rng;
+    }
+
+    METAL_FUNC float rand() {
+        uint seed = this->state * 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 result = this->state;
+        this->state = (z1^z2^z3^z4) * UNIF01_INV32;
+        return result;
+    }
+};
+
+template<typename T> METAL_FUNC void rand_uniform(
+    constant size_t &size,
+    constant float &min,
+    constant float &max,
+    device atomic_uint *seed,
+    device T *out,
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= size) {
+        return;
+    }
+
+    float diff = abs(min - max);
+    HybridTaus rng = HybridTaus::init({ulong(seed), tid, 1, 1});
+    out[tid] = static_cast<T>(rng.rand() * diff + min);
+    if (tid == 0) {
+        atomic_store_explicit(seed, uint(rng.rand() * UNIF01_NORM32), memory_order_relaxed);
+        // Return early if tid == 0, otherwise we will write to out[size].
+        return;
+    }
+    // Use symmetry to fill the other half of the array.
+    out[size - tid] = static_cast<T>(rng.rand() * diff + min);
+}
+
+// Create Gaussian normal distribution using Box-Muller transform:
+// https://en.wikipedia.org/wiki/Box–Muller_transform
+template<typename T> METAL_FUNC void normal(
+    constant size_t &size,
+    constant float &mean,
+    constant float &stddev,
+    device atomic_uint *seed,
+    device T *out,
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= size) {
+        return;
+    }
+    HybridTaus rng = HybridTaus::init({ulong(seed), tid, 1, 1});
+    float u1 = rng.rand();
+    float u2 = rng.rand();
+
+    float cosval;
+    float sinval = sincos(TWO_PI * u2, cosval);
+    float mag = stddev * sqrt(-2.0 * log(u1));
+    float z0  = mag * cosval + mean;
+    float z1  = mag * sinval + mean;
+
+    out[tid] = static_cast<T>(z0);
+
+    if (tid == 0) {
+        atomic_store_explicit(seed, uint(rng.rand() * UNIF01_NORM32), memory_order_relaxed);
+        // Return early if tid == 0, otherwise we will write to out[size].
+        return;
+    }
+    // Use symmetry to fill the other half of the array.
+    out[size - tid] = static_cast<T>(z1);
+}
+
+#define UNIFORM_OP(NAME, T)                             \
+kernel void rand_uniform_##NAME(                        \
+    constant size_t &size,                              \
+    constant float &min,                                \
+    constant float &max,                                \
+    device atomic_uint *seed,                           \
+    device T *out,                                      \
+    uint tid [[thread_position_in_grid]]                \
+) {                                                     \
+    rand_uniform<T>(size, min, max, seed, out, tid);    \
+}                                                       \
+
+#define NORMAL_OP(NAME, T)                              \
+kernel void rand_normal_##NAME(                         \
+    constant size_t &size,                              \
+    constant float &mean,                               \
+    constant float &stddev,                             \
+    device atomic_uint *seed,                           \
+    device T *out,                                      \
+    uint tid [[thread_position_in_grid]]                \
+) {                                                     \
+    normal<T>(size, mean, stddev, seed, out, tid);      \
+}                                                       \
+
+
+#define RANDOM_OPS(NAME, T) \
+UNIFORM_OP(NAME, T)         \
+NORMAL_OP(NAME, T)          \
+
+RANDOM_OPS(f32, float)
+RANDOM_OPS(f16, half)
+
+#if __METAL_VERSION__ >= 310
+RANDOM_OPS(bf16, bfloat)
+#endif
diff --git a/candle-metal-kernels/src/tests.rs b/candle-metal-kernels/src/tests.rs
index 787a7d45..655161e5 100644
--- a/candle-metal-kernels/src/tests.rs
+++ b/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)
 }
@@ -713,7 +713,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
@@ -927,3 +926,124 @@ fn gemm() {
         vec![56.0, 59.0, 62.0, 65.0, 200.0, 212.0, 224.0, 236.0]
     );
 }
+
+fn run_random<T: Clone>(name: &'static str, seed: u32, length: usize, a: f32, b: f32) -> Vec<T> {
+    let device = device();
+    let kernels = Kernels::new();
+    let command_queue = device.new_command_queue();
+    let command_buffer = command_queue.new_command_buffer();
+
+    let options = MTLResourceOptions::StorageModeManaged;
+    let output = device.new_buffer((length * core::mem::size_of::<T>()) as NSUInteger, options);
+
+    let seed = device.new_buffer_with_data(
+        &seed as *const u32 as *const core::ffi::c_void,
+        std::mem::size_of::<u32>() as NSUInteger,
+        options,
+    );
+
+    if name.starts_with("rand_uniform") {
+        call_random_uniform(
+            &device,
+            command_buffer,
+            &kernels,
+            name,
+            a,
+            b,
+            length,
+            &seed,
+            &output,
+        )
+        .unwrap();
+    } else {
+        call_random_normal(
+            &device,
+            command_buffer,
+            &kernels,
+            name,
+            a,
+            b,
+            length,
+            &seed,
+            &output,
+        )
+        .unwrap();
+    }
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    read_to_vec(&output, length)
+}
+
+#[test]
+fn random() {
+    fn calc_mean(data: &[f32]) -> f32 {
+        let sum = data.iter().sum::<f32>() as f32;
+        let count = data.len();
+        assert!(count > 0);
+        sum / count as f32
+    }
+
+    fn calc_stddev(data: &[f32]) -> f32 {
+        let mean = calc_mean(data);
+        let count = data.len();
+        assert!(count > 0);
+
+        let variance = data
+            .iter()
+            .map(|value| {
+                let diff = mean - (*value as f32);
+                diff * diff
+            })
+            .sum::<f32>()
+            / count as f32;
+
+        variance.sqrt()
+    }
+
+    let shape = vec![1024, 10];
+
+    let length = shape.iter().product::<usize>();
+    let seed = 299792458;
+
+    let min = -30.0;
+    let max = 30.0;
+    let mean = 100.0;
+    let stddev = 50.0;
+
+    macro_rules! validate_random {
+        ($type:ty) => {
+            let results: Vec<f32> = run_random::<$type>(
+                concat!("rand_uniform_", stringify!($type)),
+                seed,
+                length,
+                min,
+                max,
+            )
+            .into_iter()
+            .map(f32::from)
+            .collect();
+            results.iter().for_each(|v| {
+                assert!(*v >= min && *v <= max);
+            });
+            assert!(calc_mean(&results) > -1.0 && calc_mean(&results) < 1.0);
+
+            let results: Vec<f32> = run_random::<$type>(
+                concat!("rand_normal_", stringify!($type)),
+                seed,
+                length,
+                mean,
+                stddev,
+            )
+            .into_iter()
+            .map(f32::from)
+            .collect();
+            assert!((calc_mean(&results) - mean).abs() < mean / 10.0);
+            assert!((calc_stddev(&results) - stddev).abs() < stddev / 10.0);
+        };
+    }
+
+    validate_random!(f32);
+    validate_random!(f16);
+    validate_random!(bf16);
+}