[metal] Add CubeCL metal compiler support (#2993)

Author: syl20bnr

crates/burn-wgpu/Cargo.toml

@@ -23,10 +23,12 @@ template = ["burn-cubecl/template", "cubecl/template"]
 # Backends
 vulkan = ["cubecl-spirv"]
 webgpu = ["cubecl-wgsl"]
+metal = ["cubecl-msl"]
 
 # Compilers
 cubecl-spirv = ["cubecl/wgpu-spirv"]
 cubecl-wgsl = []
+cubecl-msl = ["cubecl/wgpu-cpp-msl"]
 
 [dependencies]
 cubecl = { workspace = true, features = ["wgpu"] }

crates/burn-wgpu/src/lib.rs

@@ -105,18 +105,26 @@ pub type Vulkan<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
 /// Tensor backend that uses the wgpu crate to execute GPU compute shaders written in WGSL.
 pub type WebGpu<F = f32, I = i32, B = u32> = Wgpu<F, I, B>;
 
+#[cfg(feature = "metal")]
+/// Tensor backend that leverages the Metal graphics API to execute GPU compute shaders compiled to MSL.
+pub type Metal<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
+
 #[cfg(test)]
 mod tests {
     use burn_cubecl::CubeBackend;
     #[cfg(feature = "vulkan")]
     pub use half::f16;
+    #[cfg(feature = "metal")]
+    pub use half::f16;
 
     pub type TestRuntime = cubecl::wgpu::WgpuRuntime;
 
     // Don't test `flex32` for now, burn sees it as `f32` but is actually `f16` precision, so it
     // breaks a lot of tests from precision issues
     #[cfg(feature = "vulkan")]
     burn_cubecl::testgen_all!([f16, f32], [i8, i16, i32, i64], [u8, u32]);
-    #[cfg(not(feature = "vulkan"))]
+    #[cfg(feature = "metal")]
+    burn_cubecl::testgen_all!([f16, f32], [i16, i32], [u32]);
+    #[cfg(all(not(feature = "vulkan"), not(feature = "metal")))]
     burn_cubecl::testgen_all!([f32], [i32], [u32]);
 }

crates/burn/Cargo.toml

@@ -89,7 +89,6 @@ fusion = ["burn-wgpu?/fusion", "burn-cuda?/fusion"]
 accelerate = ["burn-candle?/accelerate", "burn-ndarray?/blas-accelerate"]
 autotune = ["burn-wgpu?/autotune", "burn-cuda?/autotune", "burn-hip?/autotune"]
 blas-netlib = ["burn-ndarray?/blas-netlib"]
-metal = ["burn-candle?/metal"]
 openblas = ["burn-ndarray?/blas-openblas"]
 openblas-system = ["burn-ndarray?/blas-openblas-system"]
 remote = ["burn-remote/client"]
@@ -100,12 +99,14 @@ template = ["burn-wgpu?/template"]
 
 candle = ["burn-candle"]
 candle-cuda = ["candle", "burn-candle/cuda"]
+candle-metal = ["burn-candle?/metal"]
 cuda = ["burn-cuda"]
 hip = ["burn-hip"]
 ndarray = ["burn-ndarray"]
 tch = ["burn-tch"]
 vulkan = ["wgpu", "burn-wgpu/vulkan"]
 webgpu = ["wgpu", "burn-wgpu/webgpu"]
+metal = ["wgpu", "burn-wgpu/metal"]
 wgpu = ["burn-wgpu"]
 
 [dependencies]

crates/burn/src/backend.rs

@@ -27,6 +27,9 @@ pub use burn_wgpu::WebGpu;
 #[cfg(feature = "vulkan")]
 pub use burn_wgpu::Vulkan;
 
+#[cfg(feature = "metal")]
+pub use burn_wgpu::Metal;
+
 #[cfg(feature = "cuda")]
 pub use burn_cuda as cuda;
 

examples/mnist/Cargo.toml

@@ -15,6 +15,7 @@ ndarray-blas-openblas = ["burn/ndarray", "burn/openblas"]
 tch-cpu = ["burn/tch"]
 tch-gpu = ["burn/tch"]
 wgpu = ["burn/wgpu"]
+metal = ["burn/metal"]
 
 [dependencies]
 burn = { path = "../../crates/burn", features = ["train"] }

examples/mnist/examples/mnist.rs

@@ -35,7 +35,7 @@ mod tch_gpu {
     }
 }
 
-#[cfg(feature = "wgpu")]
+#[cfg(any(feature = "wgpu", feature = "metal",))]
 mod wgpu {
     use burn::backend::{
         Autodiff,
@@ -75,6 +75,6 @@ fn main() {
     tch_gpu::run();
     #[cfg(feature = "tch-cpu")]
     tch_cpu::run();
-    #[cfg(feature = "wgpu")]
+    #[cfg(any(feature = "wgpu", feature = "metal"))]
     wgpu::run();
 }

examples/text-classification/Cargo.toml

@@ -20,6 +20,7 @@ vulkan = ["burn/vulkan", "burn/default"]
 remote = ["burn/remote"]
 cuda = ["burn/cuda"]
 hip = ["burn/hip"]
+metal = ["burn/metal"]
 
 [dependencies]
 # Burn

examples/text-classification/examples/ag-news-train.rs

@@ -103,6 +103,16 @@ mod vulkan {
     }
 }
 
+#[cfg(feature = "metal")]
+mod metal {
+    use crate::{ElemType, launch};
+    use burn::backend::{Autodiff, Metal};
+
+    pub fn run() {
+        launch::<Autodiff<Metal<ElemType, i32>>>(vec![Default::default()]);
+    }
+}
+
 #[cfg(feature = "remote")]
 mod remote {
     use crate::{ElemType, launch};
@@ -155,4 +165,6 @@ fn main() {
     remote::run();
     #[cfg(feature = "vulkan")]
     vulkan::run();
+    #[cfg(feature = "metal")]
+    metal::run();
 }