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

Add avg_pool2d metal implementation for the metal backend (#1869)

Browse Source 
* implement metal avg pool 2d

* fixX

* add suggested precision workaround for the accumulator
This commit is contained in:
Thomas Santerre
2024-03-18 13:50:14 -04:00
committed by GitHub
parent 58605252e8
commit 04a61a9c72
5 changed files with 236 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1044,8 +1044,46 @@ impl BackendStorage for MetalStorage {
crate::bail!("Metal conv_tranpose2d not implemented")
}
fn avg_pool2d(&self, _: &Layout, _: (usize, usize), _: (usize, usize)) -> Result<Self> {
crate::bail!("Metal avg_pool2d not implemented")
fn avg_pool2d(
&self,
inp_l: &Layout,
(w_k, h_k): (usize, usize),
(w_stride, h_stride): (usize, usize),
) -> Result<Self> {
let shape = inp_l.shape();
let (b_size, channels, width, height) = shape.dims4()?;
let strides = inp_l.stride();
let name = match self.dtype {
DType::F32 => "avg_pool2d_f32",
DType::F16 => "avg_pool2d_f16",
DType::BF16 => "avg_pool2d_bf16",
DType::U8 => "avg_pool2d_u8",
DType::U32 => "avg_pool2d_u32",
dtype => crate::bail!("Metal avg_pool2d {dtype:?} not implemented"),
};
let out_w = (width - w_k) / w_stride + 1;
let out_h = (height - h_k) / h_stride + 1;
let dst_el = out_w * out_h * b_size * channels;
let buffer = self.device.new_buffer(dst_el, self.dtype, "avg_pool2d")?;
let command_buffers = self.device.command_buffer()?;
candle_metal_kernels::call_pool2d(
&self.device.device,
&command_buffers,
&self.device.kernels,
name,
inp_l.dims(),
strides,
out_w,
out_h,
w_k,
h_k,
w_stride,
h_stride,
&self.buffer,
&buffer,
)
.map_err(MetalError::from)?;
Ok(Self::new(buffer, self.device.clone(), dst_el, self.dtype))
}
fn max_pool2d(
@ -1063,14 +1101,14 @@ impl BackendStorage for MetalStorage {
DType::BF16 => "max_pool2d_bf16",
DType::U8 => "max_pool2d_u8",
DType::U32 => "max_pool2d_u32",
dtype => crate::bail!("Metal upsample_nearest2d {dtype:?} not implemented"),
dtype => crate::bail!("Metal max_pool2d {dtype:?} not implemented"),
};
let out_w = (width - w_k) / w_stride + 1;
let out_h = (height - h_k) / h_stride + 1;
let dst_el = out_w * out_h * b_size * channels;
let buffer = self.device.new_buffer(dst_el, self.dtype, "max_pool2d")?;
let command_buffers = self.device.command_buffer()?;
candle_metal_kernels::call_max_pool2d(
candle_metal_kernels::call_pool2d(
&self.device.device,
&command_buffers,
&self.device.kernels,

3
candle-core/tests/pool_tests.rs
View File

@ -2,9 +2,6 @@ use candle_core::{test_device, test_utils, Device, IndexOp, Result, Tensor};
// https://github.com/huggingface/candle/issues/364
fn avg_pool2d(dev: &Device) -> Result<()> {
if dev.is_metal() {
return Ok(());
}
let data: Vec<f32> = vec![
1., 1., 1., 1., 0., 0., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
];