mlx-sys 0.1.2-release

Low-level interface and binding generation for the mlx library
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

/* Copyright © 2023-2024 Apple Inc. */

#include <stdio.h>
#include "mlx/c/mlx.h"

void print_array(const char* msg, mlx_array arr) {
  mlx_string str = mlx_string_new();
  mlx_array_tostring(&str, arr);
  printf("%s\n%s\n", msg, mlx_string_data(str));
  mlx_string_free(str);
}

void exp_elemwise(
    mlx_array* output_,
    const mlx_array input,
    mlx_stream stream) {
  const char* source =
      "uint elem = thread_position_in_grid.x;"
      "T tmp = inp[elem];"
      "out[elem] = metal::exp(tmp);";
  mlx_fast_metal_kernel kernel = mlx_fast_metal_kernel_new("myexp", source, "");
  mlx_fast_metal_kernel_set_input_names(kernel, 1, "inp");
  mlx_fast_metal_kernel_set_output_names(kernel, 1, "out");

  mlx_vector_array inputs = mlx_vector_array_new_value(input);
  mlx_fast_metal_kernel_add_template_arg_dtype(kernel, "T", MLX_FLOAT32);
  mlx_fast_metal_kernel_set_grid(kernel, mlx_array_size(input), 1, 1);
  mlx_fast_metal_kernel_set_thread_group(kernel, 256, 1, 1);
  mlx_fast_metal_kernel_add_output_arg(
      kernel,
      mlx_array_shape(input),
      mlx_array_ndim(input),
      mlx_array_dtype(input));

  mlx_vector_array outputs = mlx_vector_array_new();
  mlx_fast_metal_kernel_apply(&outputs, kernel, inputs, stream);
  mlx_vector_array_get(output_, outputs, 0);

  mlx_fast_metal_kernel_free(kernel);
  mlx_vector_array_free(inputs);
  mlx_vector_array_free(outputs);
}
int main() {
  mlx_stream stream = mlx_default_gpu_stream_new();
  mlx_array input = mlx_array_new();
  mlx_array output = mlx_array_new();

  int dims[2] = {4, 16};
  mlx_random_normal(
      &input, dims, 2, MLX_FLOAT32, 0, 1, mlx_array_empty, stream);

  exp_elemwise(&output, input, stream);

  print_array("input", input);
  print_array("output", output);

  mlx_array_free(input);
  mlx_array_free(output);
  mlx_stream_free(stream);

  return 0;
}