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

Cargo.toml

@@ -61,7 +61,8 @@ 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"] }
+# metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
+metal = { path = "../metal-rs", features = ["mps"] }
 
 [profile.release-with-debug]
 inherits = "release"

candle-core/src/metal_backend.rs

@@ -67,12 +67,28 @@ impl MetalDevice {
         self.command_buffer.read().unwrap()
     }
 
-    pub fn wait_until_completed(&self) {
-        let mut old = self.command_buffer.write().unwrap();
+    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 command_buffer = self.command_queue.new_owned_command_buffer();
+        // 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)
     }
 
@@ -86,18 +102,21 @@ impl MetalDevice {
 
     pub fn new_buffer(&self, element_count: usize, dtype: DType) -> Buffer {
         let size = (element_count * dtype.size_in_bytes()) as NSUInteger;
-        self.heap
+        // println!("Creating buffer {size}");
+        let buffer = self
+            .heap
             .new_buffer(size, MTLResourceOptions::StorageModeShared)
-            .expect(" New buffer")
+            .expect("New buffer");
+        // println!("{:?}", self.heap.used_size());
+        buffer
     }
 
     pub fn new_buffer_with_data<T>(&self, data: &[T]) -> Buffer {
+        let size = core::mem::size_of_val(data) as NSUInteger;
         let option = metal::MTLResourceOptions::StorageModeShared;
-        self.device.new_buffer_with_data(
-            data.as_ptr() as *const core::ffi::c_void,
-            core::mem::size_of_val(data) as NSUInteger,
-            option,
-        )
+        // println!("Creating data buffer {size}");
+        self.device
+            .new_buffer_with_data(data.as_ptr() as *const core::ffi::c_void, size, option)
     }
 }
 
@@ -124,7 +143,7 @@ impl BackendStorage for MetalStorage {
     }
 
     fn to_cpu_storage(&self) -> Result<CpuStorage> {
-        self.device.wait_until_completed();
+        self.device.commit_wait_until_completed();
 
         match self.dtype {
             DType::U8 => Ok(CpuStorage::U8(
@@ -335,6 +354,7 @@ impl BackendStorage for MetalStorage {
         let shape = layout.shape();
         let el_count = shape.elem_count();
         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;
@@ -423,6 +443,7 @@ impl BackendStorage for MetalStorage {
                 )
                 .map_err(MetalError::from)?;
             }
+        }
         Ok(Self {
             buffer,
             device: device.clone(),
@@ -769,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
@@ -823,6 +845,7 @@ impl BackendStorage for MetalStorage {
                     &result_matrix,
                 );
             }
+        }
 
         Ok(Self {
             buffer: out_buffer,
@@ -891,7 +914,7 @@ impl BackendDevice for MetalDevice {
         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_owned_command_buffer()));
+        let command_buffer = Arc::new(RwLock::new(command_queue.new_command_buffer().to_owned()));
         let kernels = Arc::new(Kernels::new());
         Ok(Self {
             device,

candle-metal-kernels/Cargo.toml

@@ -10,7 +10,8 @@ categories = ["science"]
 license = "MIT OR Apache-2.0"
 
 [dependencies]
-metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
+# metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
+metal = { path = "../../metal-rs", features = ["mps"] }
 once_cell = "1.18.0"
 thiserror = "1"
 tracing = "0.1.37"

candle-metal-kernels/src/lib.rs

@@ -1133,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];