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

Implement generic fill. u8 uses speedy blit encoder

Browse Source
This commit is contained in:
Ivar Flakstad
2023-12-29 16:02:29 +01:00
parent 0a29d2e9b8
commit 7fc26764b6
4 changed files with 134 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
candle-core/src/metal_backend.rs
View File

@ -3,7 +3,8 @@ use crate::conv::{ParamsConv1D, ParamsConv2D, ParamsConvTranspose1D, ParamsConvT
use crate::op::{BinaryOpT, CmpOp, ReduceOp, UnaryOpT};
use crate::{CpuStorage, DType, Layout, Result, Shape};
use candle_metal_kernels;
use candle_metal_kernels::Kernels;
use candle_metal_kernels::{CallFill, Fill, Kernels};
use half::{bf16, f16};
use metal;
use metal::{Buffer, CommandBuffer, CommandQueue, MTLResourceOptions, NSUInteger};
use std::collections::HashMap;
@ -1403,25 +1404,52 @@ impl BackendDevice for MetalDevice {
let buffer = self.new_buffer(shape.elem_count(), dtype, "zeros")?;
let command_buffer = self.command_buffer()?;
command_buffer.set_label("zeros");
let blit = command_buffer.new_blit_command_encoder();
blit.wait_for_fence(&self.fence);
blit.fill_buffer(
// This assumes the specific zero type DType is equal to 0x00u8
// (which is true for all current types)
Fill::call_fill(
&self.device,
&command_buffer,
&self.kernels,
shape.elem_count(),
&buffer,
metal::NSRange {
location: 0,
length: buffer.length(),
},
0,
);
blit.update_fence(&self.fence);
blit.end_encoding();
0u8,
)
.map_err(MetalError::from)?;
Ok(MetalStorage::new(buffer, self.clone(), dtype))
}
fn ones_impl(&self, shape: &Shape, dtype: DType) -> Result<Self::Storage> {
// TODO Is there a faster way ?
let cpu_storage = crate::cpu_backend::CpuDevice.ones_impl(shape, dtype)?;
self.storage_from_cpu_storage(&cpu_storage)
let buffer = self.new_buffer(shape.elem_count(), dtype, "zeros")?;
let command_buffer = self.command_buffer()?;
command_buffer.set_label("ones");
macro_rules! fill {
($value:expr) => {
Fill::call_fill(
&self.device,
&command_buffer,
&self.kernels,
shape.elem_count(),
&buffer,
$value,
)
.map_err(MetalError::from)?
};
}
match dtype {
DType::U8 => fill!(1u8),
DType::U32 => fill!(1u32),
DType::I64 => fill!(1i64),
DType::BF16 => fill!(bf16::ONE),
DType::F16 => fill!(f16::ONE),
DType::F32 => fill!(1f32),
DType::F64 => {
return Err(MetalError::Message(format!("metal doesn't support double")).into())
}
}
Ok(MetalStorage::new(buffer, self.clone(), dtype))
}
fn storage_from_cpu_storage(&self, storage: &CpuStorage) -> Result<Self::Storage> {