Metal part 1 - Scaffolding for metal.

Cargo.toml

@@ -59,6 +59,7 @@ tracing-subscriber = "0.3.7"
 wav = "1.0.0"
 yoke = { version = "0.7.2", features = ["derive"] }
 zip = { version = "0.6.6", default-features = false }
+metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
 
 [profile.release-with-debug]
 inherits = "release"

candle-core/Cargo.toml

@@ -12,7 +12,9 @@ readme = "README.md"
 [dependencies]
 accelerate-src = { workspace = true, optional = true }
 byteorder = { workspace = true }
+tracing = { workspace = true }
 candle-kernels = { path = "../candle-kernels", version = "0.3.0", optional = true }
+metal = { workspace = true, optional = true}
 cudarc = { workspace = true, optional = true }
 gemm = { workspace = true }
 half = { workspace = true }
@@ -39,3 +41,4 @@ cuda = ["cudarc", "dep:candle-kernels"]
 cudnn = ["cuda", "cudarc/cudnn"]
 mkl = ["dep:libc", "dep:intel-mkl-src"]
 accelerate = ["dep:libc", "dep:accelerate-src"]
+metal = ["dep:metal"]

candle-core/src/device.rs

@@ -1,19 +1,21 @@
 use crate::backend::BackendDevice;
 use crate::cpu_backend::CpuDevice;
-use crate::{CpuStorage, DType, Result, Shape, Storage, WithDType};
+use crate::{bail, CpuStorage, DType, Result, Shape, Storage, WithDType};
 
 /// A `DeviceLocation` represents a physical device whereas multiple `Device`
 /// can live on the same location (typically for cuda devices).
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum DeviceLocation {
  Cpu,
  Cuda { gpu_id: usize },
+ Metal,
 }
 
 #[derive(Debug, Clone)]
 pub enum Device {
  Cpu,
  Cuda(crate::CudaDevice),
+ Metal(crate::MetalDevice),
 }
 
 pub trait NdArray {
@@ -103,14 +105,14 @@ impl<S: WithDType, const N1: usize, const N2: usize, const N3: usize, const N4:
 impl<S: NdArray> NdArray for Vec<S> {
  fn shape(&self) -> Result<Shape> {
  if self.is_empty() {
- crate::bail!("empty array")
+ bail!("empty array")
  }
  let shape0 = self[0].shape()?;
  let n = self.len();
  for v in self.iter() {
  let shape = v.shape()?;
  if shape != shape0 {
- crate::bail!("two elements have different shapes {shape:?} {shape0:?}")
+ bail!("two elements have different shapes {shape:?} {shape0:?}")
  }
  }
  Ok(Shape::from([[n].as_slice(), shape0.dims()].concat()))
@@ -128,10 +130,15 @@ impl Device {
  Ok(Self::Cuda(crate::CudaDevice::new(ordinal)?))
  }
 
+ pub fn new_metal(ordinal: usize) -> Result<Self> {
+ Ok(Self::Metal(crate::MetalDevice::new(ordinal)?))
+ }
+
  pub fn set_seed(&self, seed: u64) -> Result<()> {
  match self {
- Self::Cpu => crate::cpu_backend::CpuDevice.set_seed(seed),
+ Self::Cpu => CpuDevice.set_seed(seed),
  Self::Cuda(c) => c.set_seed(seed),
+ Self::Metal(m) => m.set_seed(seed),
  }
  }
 
@@ -147,21 +154,16 @@ impl Device {
  match self {
  Self::Cpu => DeviceLocation::Cpu,
  Self::Cuda(device) => device.location(),
+ Device::Metal(device) => device.location(),
  }
  }
 
  pub fn is_cpu(&self) -> bool {
- match self {
- Self::Cpu => true,
- Self::Cuda(_) => false,
- }
+ matches!(self, Self::Cpu)
  }
 
  pub fn is_cuda(&self) -> bool {
- match self {
- Self::Cpu => false,
- Self::Cuda(_) => true,
- }
+ matches!(self, Self::Cuda(_))
  }
 
  pub fn cuda_if_available(ordinal: usize) -> Result<Self> {
@@ -194,6 +196,11 @@ impl Device {
  Ok(Storage::Cuda(storage))
  }
  }
+ Device::Metal(_device) => {
+ // let storage = device.rand_uniform(shape, dtype, lo, up)?;
+ // Ok(Storage::Metal(storage))
+ bail!("Metal rand_uniform not implemented")
+ }
  }
  }
 
@@ -228,6 +235,10 @@ impl Device {
  Ok(Storage::Cuda(storage))
  }
  }
+ Device::Metal(device) => {
+ let storage = device.rand_normal(shape, dtype, mean, std)?;
+ Ok(Storage::Metal(storage))
+ }
  }
  }
 
@@ -250,6 +261,10 @@ impl Device {
  let storage = device.ones_impl(shape, dtype)?;
  Ok(Storage::Cuda(storage))
  }
+ Device::Metal(device) => {
+ let storage = device.ones_impl(shape, dtype)?;
+ Ok(Storage::Metal(storage))
+ }
  }
  }
 
@@ -263,6 +278,10 @@ impl Device {
  let storage = device.zeros_impl(shape, dtype)?;
  Ok(Storage::Cuda(storage))
  }
+ Device::Metal(device) => {
+ let storage = device.zeros_impl(shape, dtype)?;
+ Ok(Storage::Metal(storage))
+ }
  }
  }
 
@@ -274,6 +293,11 @@ impl Device {
  let storage = device.storage_from_cpu_storage(&storage)?;
  Ok(Storage::Cuda(storage))
  }
+ Device::Metal(device) => {
+ let storage = array.to_cpu_storage();
+ let storage = device.storage_from_cpu_storage(&storage)?;
+ Ok(Storage::Metal(storage))
+ }
  }
  }
 
@@ -285,6 +309,11 @@ impl Device {
  let storage = device.storage_from_cpu_storage(&storage)?;
  Ok(Storage::Cuda(storage))
  }
+ Device::Metal(device) => {
+ let storage = S::to_cpu_storage_owned(data);
+ let storage = device.storage_from_cpu_storage(&storage)?;
+ Ok(Storage::Metal(storage))
+ }
  }
  }
 }

candle-core/src/display.rs

@@ -14,6 +14,7 @@ impl Tensor {
  crate::DeviceLocation::Cuda { gpu_id } => {
  format!(", cuda:{}", gpu_id)
  }
+ _ => todo!(),
  };
 
  write!(f, "Tensor[")?;
@@ -476,6 +477,7 @@ impl std::fmt::Display for Tensor {
  crate::DeviceLocation::Cuda { gpu_id } => {
  format!(", cuda:{}", gpu_id)
  }
+ crate::DeviceLocation::Metal => todo!(),
  };
 
  write!(