Tmp commit for the heap experiment (heap is indeed decreasing).

Cargo.toml

@@ -63,8 +63,6 @@ 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"] }
 metal = { path = "../metal-rs", features = ["mps"] }
-dispatch = "0.2.0"
-rustc-hash = "1.1"
 
 [profile.release-with-debug]
 inherits = "release"

candle-core/Cargo.toml

@@ -30,8 +30,6 @@ safetensors = { workspace = true }
 thiserror = { workspace = true }
 yoke = { workspace = true }
 zip = { workspace = true }
-dispatch = { workspace = true, optional = true }
-rustc-hash = { workspace = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
@@ -43,4 +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", "dep:candle-metal-kernels", "dep:dispatch"]
+metal = ["dep:metal", "dep:candle-metal-kernels"]

candle-core/src/metal_backend.rs

@@ -6,11 +6,8 @@ use candle_metal_kernels;
 use candle_metal_kernels::Kernels;
 use half::f16;
 use metal;
-use metal::{Buffer, CommandBuffer, CommandQueue, MTLResourceOptions, NSUInteger};
+use metal::{Buffer, CommandBuffer, CommandQueue, HeapDescriptor, MTLResourceOptions, NSUInteger};
 use std::sync::{Arc, RwLock};
-use std::collections::HashMap;
-use rustc_hash::FxHashMap;
-use dispatch::{Queue, QueueAttribute};
 
 /// Metal related errors
 #[derive(thiserror::Error, Debug)]
@@ -38,9 +35,8 @@ impl From<String> for MetalError {
 pub struct MetalDevice {
     device: metal::Device,
     command_queue: metal::CommandQueue,
+    heap: metal::Heap,
     command_buffer: Arc<RwLock<metal::CommandBuffer>>,
-    buffers: Arc<RwLock<FxHashMap<usize, Vec<Buffer>>>>,
-    queue : Queue,
     kernels: Arc<candle_metal_kernels::Kernels>,
 }
 
@@ -68,15 +64,31 @@ impl MetalDevice {
     }
 
     pub fn command_buffer(&self) -> std::sync::RwLockReadGuard<CommandBuffer> {
-        self.command_buffer.try_read().unwrap()
+        self.command_buffer.read().unwrap()
     }
 
-    pub fn wait_until_completed(&self) {
+    pub fn commit_wait_until_completed(&self) {
         let mut old = self.command_buffer.try_write().unwrap();
+        let status = old.status();
+        use metal::MTLCommandBufferStatus::{
+            Committed, Completed, Enqueued, Error, NotEnqueued, Scheduled,
+        };
+        // match old.status() {}
+        if old.status() == metal::MTLCommandBufferStatus::Completed {
+            return;
+        }
         old.commit();
         old.wait_until_completed();
+        // let count = old.retain_count();
+        // println!("Count {count:?}");
         let command_buffer = self.command_queue.new_command_buffer().to_owned();
+
         *old = command_buffer;
+        // let count = old.retain_count();
+        // // println!("Count after {count:?}");
+        // old.release();
+        // let count = old.retain_count();
+        // println!("Count after release {count:?}");
         // self.command_buffer.replace_with(|_| command_buffer)
     }
 
@@ -89,34 +101,22 @@ impl MetalDevice {
     }
 
     pub fn new_buffer(&self, element_count: usize, dtype: DType) -> Buffer {
-        let size = element_count * dtype.size_in_bytes();
-        let mut buffers = self.buffers.try_write().unwrap();
-        let subbuffers = buffers.entry(size).or_insert(vec![]);
-
-        for sub in &mut *subbuffers{
-            if sub.retain_count() == 1{
-                return sub.clone();
-                // println!("{size } {:?}", sub.retain_count());
-            }
-        }
-        let new_buffer = self.device
-            .new_buffer(size as NSUInteger, MTLResourceOptions::StorageModePrivate);
-        subbuffers.push(new_buffer.clone());
-        new_buffer
-    }
-
-    pub fn new_buffer_managed(&self, size: NSUInteger) -> Buffer {
-        self.device
-            .new_buffer(size, MTLResourceOptions::StorageModeManaged)
+        let size = (element_count * dtype.size_in_bytes()) as NSUInteger;
+        // println!("Creating buffer {size}");
+        let buffer = self
+            .heap
+            .new_buffer(size, MTLResourceOptions::StorageModeShared)
+            .expect("New buffer");
+        // println!("{:?}", self.heap.used_size());
+        buffer
     }
 
     pub fn new_buffer_with_data<T>(&self, data: &[T]) -> Buffer {
-        let option = metal::MTLResourceOptions::StorageModeManaged;
-        self.device.new_buffer_with_data(
-            data.as_ptr() as *const core::ffi::c_void,
-            core::mem::size_of_val(data) as NSUInteger,
-            option,
-        )
+        let size = core::mem::size_of_val(data) as NSUInteger;
+        let option = metal::MTLResourceOptions::StorageModeShared;
+        // println!("Creating data buffer {size}");
+        self.device
+            .new_buffer_with_data(data.as_ptr() as *const core::ffi::c_void, size, option)
     }
 }
 
@@ -143,39 +143,29 @@ impl BackendStorage for MetalStorage {
     }
 
     fn to_cpu_storage(&self) -> Result<CpuStorage> {
-        let buffer = self.device.new_buffer_managed(self.buffer.length());
-        {
-        let command = self.device.command_buffer();
-        let blit = command.new_blit_command_encoder();
-        blit.copy_from_buffer(&self.buffer, 0, &buffer, 0, self.buffer.length());
-        blit.end_encoding();
-
-        }
-
-
-        self.device.wait_until_completed();
+        self.device.commit_wait_until_completed();
 
         match self.dtype {
             DType::U8 => Ok(CpuStorage::U8(
-                buffer.read_to_vec(buffer.length() as usize),
+                self.buffer.read_to_vec(self.buffer.length() as usize),
             )),
             DType::U32 => Ok(CpuStorage::U32(
-                buffer.read_to_vec(buffer.length() as usize / 4),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 4),
             )),
             DType::I64 => Ok(CpuStorage::I64(
-                buffer.read_to_vec(buffer.length() as usize / 8),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 8),
             )),
             DType::F16 => Ok(CpuStorage::F16(
-                buffer.read_to_vec(buffer.length() as usize / 2),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 2),
             )),
             DType::BF16 => Ok(CpuStorage::BF16(
-                buffer.read_to_vec(buffer.length() as usize / 2),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 2),
             )),
             DType::F32 => Ok(CpuStorage::F32(
-                buffer.read_to_vec(buffer.length() as usize / 4),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 4),
             )),
             DType::F64 => Ok(CpuStorage::F64(
-                buffer.read_to_vec(buffer.length() as usize / 8),
+                self.buffer.read_to_vec(self.buffer.length() as usize / 8),
             )),
         }
     }
@@ -359,23 +349,14 @@ impl BackendStorage for MetalStorage {
     }
 
     fn unary_impl<B: UnaryOpT>(&self, layout: &Layout) -> Result<Self> {
-
-
         let device = self.device();
         let dtype = self.dtype;
         let shape = layout.shape();
         let el_count = shape.elem_count();
-        let buffer = device.new_buffer(el_count, dtype);
-        let metal = self.device.clone();
-        let mut cloned = buffer.clone();
-        let inbuffer = self.buffer.clone();
-        let ldims = layout.dims().to_owned();
-        let lstride = layout.stride().to_owned();
-        let loffset = layout.start_offset() * dtype.size_in_bytes();
-        if layout.is_contiguous() && layout.start_offset() == 0 {
-        // self.device.queue.exec_async(move || {
-            let device = metal;
+        let mut buffer = device.new_buffer(el_count, dtype);
+        {
             let command_buffer = device.command_buffer();
+            if layout.is_contiguous() && layout.start_offset() == 0 {
                 use candle_metal_kernels::unary::contiguous;
 
                 let kernel_name = match (B::KERNEL, dtype) {
@@ -413,16 +394,11 @@ impl BackendStorage for MetalStorage {
                     &device.kernels,
                     kernel_name,
                     el_count,
-                &inbuffer,
-                &mut cloned,
+                    &self.buffer,
+                    &mut buffer,
                 )
-            .unwrap();
-        // });
-
+                .map_err(MetalError::from)?;
             } else {
-        // self.device.queue.exec_async(move || {
-            let device = metal;
-        let command_buffer = device.command_buffer();
                 use candle_metal_kernels::unary::strided;
                 let kernel_name = match (B::KERNEL, dtype) {
                     ("ucos", DType::F32) => strided::cos::FLOAT,
@@ -458,17 +434,16 @@ impl BackendStorage for MetalStorage {
                     &command_buffer,
                     &device.kernels,
                     kernel_name,
-                &ldims,
-                &inbuffer,
-                &lstride,
-                loffset,
-                &mut cloned,
+                    layout.dims(),
+                    &self.buffer,
+                    layout.stride(),
+                    layout.start_offset() * self.dtype.size_in_bytes(),
+                    &mut buffer,
                     0,
                 )
-            .unwrap();
-        // });
+                .map_err(MetalError::from)?;
+            }
         }
-
         Ok(Self {
             buffer,
             device: device.clone(),
@@ -815,6 +790,7 @@ impl BackendStorage for MetalStorage {
 
         let out_buffer = self.device.new_buffer(elem_count, self.dtype);
 
+        {
             let command_buffer = self.device.command_buffer();
             for bi in 0..b {
                 // Create matrix objects
@@ -869,6 +845,7 @@ impl BackendStorage for MetalStorage {
                     &result_matrix,
                 );
             }
+        }
 
         Ok(Self {
             buffer: out_buffer,
@@ -928,30 +905,22 @@ impl BackendDevice for MetalDevice {
     fn new(ordinal: usize) -> Result<Self> {
         let device = metal::Device::all().swap_remove(ordinal);
 
-        
-        // let capture = metal::CaptureManager::shared();
-        // let descriptor = metal::CaptureDescriptor::new();
-        // descriptor.set_destination(metal::MTLCaptureDestination::GpuTraceDocument);
-        // descriptor.set_capture_device(&device);
-        // let mut dir = std::env::current_dir()?;
-        // dir.push("out.gputrace");
-        // descriptor.set_output_url(dir);
-
-        // capture
-        //     .start_capture(&descriptor)
-        //     .map_err(MetalError::from)?;
-
         let command_queue = device.new_command_queue();
+
+        let descriptor = HeapDescriptor::new();
+        let mut size =
+            device.heap_buffer_size_and_align(100_000_000, MTLResourceOptions::StorageModeShared);
+        size.size += (size.size & (size.align - 1)) + size.align;
+        descriptor.set_size(size.size);
+        descriptor.set_storage_mode(metal::MTLStorageMode::Shared);
+        let heap = device.new_heap(&descriptor);
         let command_buffer = Arc::new(RwLock::new(command_queue.new_command_buffer().to_owned()));
         let kernels = Arc::new(Kernels::new());
-        let queue = Queue::create("co.huggingface.candle", QueueAttribute::Serial);
-        let buffers = Arc::new(RwLock::new(FxHashMap::default()));
         Ok(Self {
             device,
+            heap,
             command_queue,
             command_buffer,
-            buffers,
-            queue,
             kernels,
         })
     }

candle-metal-kernels/src/lib.rs

@@ -300,9 +300,6 @@ pub fn call_unary_contiguous(
     input: &Buffer,
     output: &mut Buffer,
 ) -> Result<(), MetalKernelError> {
-    // println!("Kernel {:?}", kernel_name.0);
-    // assert_eq!(input.length(), output.length());
-
     let pipeline = kernels.load_pipeline(device, Source::Unary, kernel_name.0)?;
     let encoder = command_buffer.new_compute_command_encoder();
     encoder.set_compute_pipeline_state(&pipeline);
@@ -1136,6 +1133,7 @@ mod tests {
         let device = Device::system_default().expect("no device found");
 
         let options = CompileOptions::new();
+        options.set_fast_math_enabled(true);
         let library = device.new_library_with_source(INDEXING, &options).unwrap();
 
         let left = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];