mlxrs 0.1.0

Safe Rust bindings for Apple's MLX array framework, with LM, VLM, audio, and embeddings support
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

//! Arithmetic ops benches — M2 representative set.
//!
//! Each bench pre-constructs the input arrays outside the timing loop and
//! `eval()`s the result inside, so we measure the op call + Metal eval rather
//! than FFI graph construction alone. Shapes are fixed at 1024x1024 f32.

use std::hint::black_box;

use criterion::{Criterion, criterion_group, criterion_main};

fn bench_add(c: &mut Criterion) {
  let a = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let b = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let mut group = c.benchmark_group("arithmetic");
  group.bench_function("add 1024x1024 f32", |bencher| {
    bencher.iter(|| {
      let mut r = mlxrs::ops::arithmetic::add(black_box(&a), black_box(&b)).unwrap();
      r.eval().unwrap();
      black_box(r);
    });
  });
  group.finish();
}

fn bench_multiply(c: &mut Criterion) {
  let a = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let b = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let mut group = c.benchmark_group("arithmetic");
  group.bench_function("multiply 1024x1024 f32", |bencher| {
    bencher.iter(|| {
      let mut r = mlxrs::ops::arithmetic::multiply(black_box(&a), black_box(&b)).unwrap();
      r.eval().unwrap();
      black_box(r);
    });
  });
  group.finish();
}

fn bench_negative(c: &mut Criterion) {
  let a = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let mut group = c.benchmark_group("arithmetic");
  group.bench_function("negative 1024x1024 f32", |bencher| {
    bencher.iter(|| {
      let mut r = mlxrs::ops::arithmetic::negative(black_box(&a)).unwrap();
      r.eval().unwrap();
      black_box(r);
    });
  });
  group.finish();
}

fn bench_sqrt(c: &mut Criterion) {
  let a = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let mut group = c.benchmark_group("arithmetic");
  group.bench_function("sqrt 1024x1024 f32", |bencher| {
    bencher.iter(|| {
      let mut r = mlxrs::ops::arithmetic::sqrt(black_box(&a)).unwrap();
      r.eval().unwrap();
      black_box(r);
    });
  });
  group.finish();
}

fn bench_exp(c: &mut Criterion) {
  let a = mlxrs::Array::ones::<f32>(&(1024usize, 1024)).unwrap();
  let mut group = c.benchmark_group("arithmetic");
  group.bench_function("exp 1024x1024 f32", |bencher| {
    bencher.iter(|| {
      let mut r = mlxrs::ops::arithmetic::exp(black_box(&a)).unwrap();
      r.eval().unwrap();
      black_box(r);
    });
  });
  group.finish();
}

criterion_group!(
  benches,
  bench_add,
  bench_multiply,
  bench_negative,
  bench_sqrt,
  bench_exp,
);
criterion_main!(benches);