Metal updates

Author: EricLBuehler

candle-core/src/lib.rs

@@ -102,7 +102,7 @@ pub use shape::{Shape, D};
 pub use storage::Storage;
 pub use streaming::{StreamTensor, StreamingBinOp, StreamingModule};
 pub use strided_index::{StridedBlocks, StridedIndex};
-pub use tensor::{Tensor, TensorId};
+pub use tensor::{from_storage_no_op, Tensor, TensorId};
 pub use variable::Var;
 
 #[cfg(feature = "cuda")]

candle-core/src/metal_backend/device.rs

@@ -58,14 +58,21 @@ impl Commands {
         })
     }
 
+    pub fn flush_command_buffer(&mut self) -> Result<()> {
+        self.command_buffer.commit();
+        let command_buffer = self.command_queue.new_command_buffer().to_owned();
+        self.command_buffer = command_buffer.clone();
+        self.command_buffer_index = 0;
+
+        Ok(())
+    }
+
     pub fn command_buffer(&mut self) -> Result<(bool, CommandBuffer)> {
         let mut command_buffer = self.command_buffer.to_owned();
         let mut flushed = false;
         if self.command_buffer_index > self.compute_per_buffer {
-            self.command_buffer.commit();
-            command_buffer = self.command_queue.new_command_buffer().to_owned();
-            self.command_buffer = command_buffer.clone();
-            self.command_buffer_index = 0;
+            self.flush_command_buffer()?;
+            command_buffer = self.command_buffer.to_owned();
             flushed = true;
         }
         self.command_buffer_index += 1;
@@ -181,6 +188,13 @@ impl MetalDevice {
         Ok(())
     }
 
+    pub fn flush_command_buffer(&self) -> Result<()> {
+        let mut commands = self.commands.write().map_err(MetalError::from)?;
+        commands.flush_command_buffer()?;
+
+        Ok(())
+    }
+
     pub fn command_buffer(&self) -> Result<CommandBuffer> {
         let mut commands = self.commands.write().map_err(MetalError::from)?;
         let (flushed, command_buffer) = commands.command_buffer()?;
@@ -190,6 +204,11 @@ impl MetalDevice {
         Ok(command_buffer)
     }
 
+    pub fn new_command_buffer(&self) -> Result<CommandBuffer> {
+        let commands = self.commands.write().map_err(MetalError::from)?;
+        Ok(commands.command_queue.new_command_buffer().to_owned())
+    }
+
     pub fn wait_until_completed(&self) -> Result<()> {
         let mut commands = self.commands.write().map_err(MetalError::from)?;
         commands.wait_until_completed()

candle-core/src/tensor.rs

@@ -176,6 +176,22 @@ pub(crate) fn from_storage<S: Into<Shape>>(
     Tensor(Arc::new(tensor_))
 }
 
+/// Creates a fresh tensor structure based on a storage and a shape, this uses contiguous strides. This has a BackpropOp:none().
+pub fn from_storage_no_op<S: Into<Shape>>(storage: Storage, shape: S, is_variable: bool) -> Tensor {
+    let dtype = storage.dtype();
+    let device = storage.device();
+    let tensor_ = Tensor_ {
+        id: TensorId::new(),
+        storage: Arc::new(RwLock::new(storage)),
+        layout: Layout::contiguous(shape),
+        op: BackpropOp::none(),
+        is_variable,
+        dtype,
+        device,
+    };
+    Tensor(Arc::new(tensor_))
+}
+
 impl Tensor {
     pub(crate) fn ones_impl<S: Into<Shape>>(
         shape: S,