huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 03:54:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use RAII for terminating the encoding. (#2353)

Browse Source
This commit is contained in:
Laurent Mazare
2024-07-24 15:29:56 +01:00
committed by GitHub
parent a925ae6bc6
commit ddafc61055
2 changed files with 69 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
candle-metal-kernels/src/utils.rs
View File

@ -162,24 +162,40 @@ macro_rules! set_params {
}
pub trait EncoderProvider {
fn encoder(&self) -> &ComputeCommandEncoderRef;
fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef);
type Encoder<'a>: AsRef<metal::ComputeCommandEncoderRef>
where
Self: 'a;
fn encoder<'a>(&'a self) -> Self::Encoder<'a>;
}
pub struct WrappedEncoder<'a>(&'a ComputeCommandEncoderRef);
impl<'a> Drop for WrappedEncoder<'a> {
fn drop(&mut self) {
self.0.end_encoding()
}
}
impl<'a> AsRef<metal::ComputeCommandEncoderRef> for WrappedEncoder<'a> {
fn as_ref(&self) -> &metal::ComputeCommandEncoderRef {
&self.0
}
}
impl EncoderProvider for &metal::CommandBuffer {
fn encoder(&self) -> &ComputeCommandEncoderRef {
self.new_compute_command_encoder()
}
fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef) {
enc.end_encoding()
type Encoder<'a> = WrappedEncoder<'a>
where
Self: 'a;
fn encoder<'a>(&'a self) -> Self::Encoder<'a> {
WrappedEncoder(self.new_compute_command_encoder())
}
}
impl EncoderProvider for &metal::CommandBufferRef {
fn encoder(&self) -> &ComputeCommandEncoderRef {
self.new_compute_command_encoder()
}
fn maybe_end_encoding(&self, enc: &ComputeCommandEncoderRef) {
enc.end_encoding()
type Encoder<'a> = WrappedEncoder<'a>
where
Self: 'a;
fn encoder<'a>(&'a self) -> Self::Encoder<'a> {
WrappedEncoder(self.new_compute_command_encoder())
}
}