Use a trait for the encoder provider (so that encoder can ultimately be reused). (#2352)

1 files changed, 23 insertions, 0 deletions

diff --git a/candle-metal-kernels/src/utils.rs b/candle-metal-kernels/src/utils.rs
index 194cddf4..4ef2162c 100644
--- a/candle-metal-kernels/src/utils.rs
+++ b/candle-metal-kernels/src/utils.rs
@@ -160,3 +160,26 @@ macro_rules! set_params {
         )*
     );
 }
+
+pub trait EncoderProvider {
+    fn encoder(&self) -> &ComputeCommandEncoderRef;
+    fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef);
+}
+
+impl EncoderProvider for &metal::CommandBuffer {
+    fn encoder(&self) -> &ComputeCommandEncoderRef {
+        self.new_compute_command_encoder()
+    }
+    fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef) {
+        enc.end_encoding()
+    }
+}
+
+impl EncoderProvider for &metal::CommandBufferRef {
+    fn encoder(&self) -> &ComputeCommandEncoderRef {
+        self.new_compute_command_encoder()
+    }
+    fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef) {
+        enc.end_encoding()
+    }
+}