feat: support metal

.github/workflows/binaries.yml

@@ -102,3 +102,37 @@ jobs:
           draft: true
           files: ${{ steps.precompile.outputs.file-path }}
         if: startsWith(github.ref, 'refs/tags/')
+
+  build_metal:
+    name: metal / ${{ matrix.target }} / ${{ matrix.os }}
+    runs-on: macos-13
+    permissions:
+      contents: write
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - target: aarch64-apple-darwin
+            os: macos-13
+
+    steps:
+      - uses: actions/checkout@v4
+      - run: rustup target add ${{ matrix.target }}
+
+      - uses: philss/rustler-precompiled-action@main
+        id: precompile
+        with:
+          project-dir: ${{ env.PROJECT_DIR }}
+          project-name: ${{ env.PROJECT_NAME }}
+          project-version: ${{ env.PROJECT_VERSION }}
+          target: ${{ matrix.target }}
+          use-cross: null
+          nif-version: ${{ env.NIF_VERSION }}
+          variant: metal
+          cargo-args: "--features metal"
+
+      - uses: softprops/action-gh-release@v1
+        with:
+          draft: true
+          files: ${{ steps.precompile.outputs.file-path }}
+        if: startsWith(github.ref, 'refs/tags/')

README.md

@@ -57,6 +57,7 @@ only the host CPU.
 | --- | --- |
 | cpu | |
 | cuda | CUDA 12.x |
+| metal | Metal ? |
 
 To use Candlex with NVidia GPU you need [CUDA](https://developer.nvidia.com/cuda-downloads) compatible with your
 GPU drivers.

config/config.exs

@@ -1,3 +1,5 @@
 import Config
 
-config :candlex, use_cuda: System.get_env("CANDLEX_NIF_TARGET") == "cuda"
+config :candlex,
+  use_cuda: System.get_env("CANDLEX_NIF_TARGET") == "cuda",
+  use_metal: System.get_env("CANDLEX_NIF_TARGET") == "metal"

lib/candlex/backend.ex

@@ -12,6 +12,7 @@ defmodule Candlex.Backend do
 
   @device_cuda :cuda
   @device_cpu :cpu
+  @device_metal :metal
 
   @impl true
   def init(opts) do
@@ -1198,10 +1199,15 @@ defmodule Candlex.Backend do
   end
 
   defp default_device do
-    if cuda_available?() do
-      @device_cuda
-    else
-      @device_cpu
+    cond do
+      cuda_available?() ->
+        @device_cuda
+
+      metal_available?() ->
+        @device_metal
+
+      true ->
+        @device_cpu
     end
   end
 
@@ -1217,6 +1223,10 @@ defmodule Candlex.Backend do
     Native.is_cuda_available()
   end
 
+  def metal_available? do
+    Native.is_metal_available()
+  end
+
   defp unsupported_dtype(t) do
     raise("Unsupported candle dtype for #{inspect(t)}")
   end

lib/candlex/native.ex

@@ -9,9 +9,16 @@ defmodule Candlex.Native do
   # mode = if Mix.env() in [:dev, :test], do: :debug, else: :release
   mode = :release
 
+  features =
+    cond do
+      Application.compile_env(:candlex, :use_cuda) -> [:cuda]
+      Application.compile_env(:candlex, :use_metal) -> [:metal]
+      true -> []
+    end
+
   use RustlerPrecompiled,
     otp_app: :candlex,
-    features: if(Application.compile_env(:candlex, :use_cuda), do: [:cuda], else: []),
+    features: features,
     base_url: "#{source_url}/releases/download/v#{version}",
     force_build: System.get_env("CANDLEX_NIF_BUILD") in ["1", "true"],
     mode: mode,
@@ -25,6 +32,7 @@ defmodule Candlex.Native do
       "x86_64-unknown-linux-gnu"
     ],
     variants: %{
+      "aarch64-apple-darwin" => [metal: fn -> Application.compile_env(:candlex, :use_metal) end],
       "x86_64-unknown-linux-gnu" => [cuda: fn -> Application.compile_env(:candlex, :use_cuda) end]
     }
 
@@ -138,6 +146,7 @@ defmodule Candlex.Native do
   end
 
   def is_cuda_available(), do: error()
+  def is_metal_available(), do: error()
   def to_device(_tensor, _device), do: error()
 
   defp error(), do: :erlang.nif_error(:nif_not_loaded)

native/candlex/Cargo.toml

@@ -22,3 +22,4 @@ bindgen_cuda = { version = "0.1.1", optional = true }
 
 [features]
 cuda = ["candle-core/cuda", "dep:bindgen_cuda"]
+metal = ["candle-core/metal"]

native/candlex/src/devices.rs

@@ -2,3 +2,8 @@
 pub fn is_cuda_available() -> bool {
     candle_core::utils::cuda_is_available()
 }
+
+#[rustler::nif(schedule = "DirtyCpu")]
+pub fn is_metal_available() -> bool {
+    candle_core::utils::metal_is_available()
+}

native/candlex/src/lib.rs

@@ -121,7 +121,8 @@ rustler::init! {
         tensors::sum_pool2d,
         tensors::max_pool2d,
         tensors::contiguous,
-        devices::is_cuda_available
+        devices::is_cuda_available,
+        devices::is_metal_available
     ],
     load = load
 }