gazagnaire.org / ocaml-metal
Metal GPU matrix multiplication with a CPU fallback
0
Fork 0

Configure Feed

Select the types of activity you want to include in your feed.

OCaml 68.2%
C 20.8%
Dune 2.9%
Other 8.1%
19 1 0

Clone this repository

https://git.vm.fail/gazagnaire.org/ocaml-metal https://git.vm.fail/did:plc:zbma44t57f5fbacc6xwzshex
ssh://git@git.recoil.org:2222/gazagnaire.org/ocaml-metal ssh://git@git.recoil.org:2222/did:plc:zbma44t57f5fbacc6xwzshex

For self-hosted knots, clone URLs may differ based on your setup.


Download tar.gz Download .zip
README.md

nox-metal#

Metal GPU matrix multiplication in OCaml, with a CPU fallback everywhere else.

A thin binding to MPSMatrixMultiplication -- the GEMM at the core of neural-network inference -- over float32 row-major matrices carried as Bigarray. On a Mac with a Metal device the product runs on the GPU; on any other platform, or a Mac without one, the same call computes on the CPU, so code built on it runs unchanged.

# let a = Metal.of_array [| 1.; 2.; 3.; 4. |]   (* 2x2 *)
val a : Metal.t = <abstr>
# let b = Metal.of_array [| 5.; 6.; 7.; 8. |]   (* 2x2 *)
val b : Metal.t = <abstr>
# Metal.to_array (Metal.gemm ~m:2 ~k:2 ~n:2 a b)
- : float array = [|19.; 22.; 43.; 50.|]

Metal.available () reports whether the GPU path is active; Metal.backend () is Metal or Cpu.

Installation#

Install with opam:

$ opam install nox-metal

If opam cannot find the package, it may not yet be released in the public opam-repository. Add the overlay repository, then install it:

$ opam repo add samoht https://tangled.org/gazagnaire.org/opam-overlay.git
$ opam update
$ opam install nox-metal

Building#

The single C stub (lib/metal_stubs.c) is compiled as Objective-C under ARC on macOS and linked against Metal and MetalPerformanceShaders; on other platforms it is a raising fallback the OCaml side never reaches. No configuration is needed -- the dune rules pick the frameworks on Darwin and skip them elsewhere.

  • mirage/metal and other bindings expose more of the Metal API surface; this package is deliberately just the GEMM primitive an inference backend needs, with a portable fallback.
  • Apple's Metal Performance Shaders is the underlying framework.

License#

ISC. See LICENSE.md.