mirror of
https://github.com/huggingface/candle.git
synced 2025-06-19 03:54:56 +00:00
Metal commands refactoring (#2489)
* Split out the commands part of the metal device. * Make most fields private. * Move the allocator back. * Rework the encoder provider type.
This commit is contained in:
@ -168,17 +168,22 @@ pub trait EncoderProvider {
|
||||
fn encoder(&self) -> Self::Encoder<'_>;
|
||||
}
|
||||
|
||||
pub struct WrappedEncoder<'a>(&'a ComputeCommandEncoderRef);
|
||||
pub struct WrappedEncoder<'a> {
|
||||
inner: &'a ComputeCommandEncoderRef,
|
||||
end_encoding_on_drop: bool,
|
||||
}
|
||||
|
||||
impl<'a> Drop for WrappedEncoder<'a> {
|
||||
fn drop(&mut self) {
|
||||
self.0.end_encoding()
|
||||
if self.end_encoding_on_drop {
|
||||
self.inner.end_encoding()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> AsRef<metal::ComputeCommandEncoderRef> for WrappedEncoder<'a> {
|
||||
fn as_ref(&self) -> &metal::ComputeCommandEncoderRef {
|
||||
self.0
|
||||
self.inner
|
||||
}
|
||||
}
|
||||
|
||||
@ -187,7 +192,10 @@ impl EncoderProvider for &metal::CommandBuffer {
|
||||
where
|
||||
Self: 'a;
|
||||
fn encoder(&self) -> Self::Encoder<'_> {
|
||||
WrappedEncoder(self.new_compute_command_encoder())
|
||||
WrappedEncoder {
|
||||
inner: self.new_compute_command_encoder(),
|
||||
end_encoding_on_drop: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -196,6 +204,21 @@ impl EncoderProvider for &metal::CommandBufferRef {
|
||||
where
|
||||
Self: 'a;
|
||||
fn encoder(&self) -> Self::Encoder<'_> {
|
||||
WrappedEncoder(self.new_compute_command_encoder())
|
||||
WrappedEncoder {
|
||||
inner: self.new_compute_command_encoder(),
|
||||
end_encoding_on_drop: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl EncoderProvider for &ComputeCommandEncoderRef {
|
||||
type Encoder<'a> = WrappedEncoder<'a>
|
||||
where
|
||||
Self: 'a;
|
||||
fn encoder(&self) -> Self::Encoder<'_> {
|
||||
WrappedEncoder {
|
||||
inner: self,
|
||||
end_encoding_on_drop: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user