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

Extend supported dtypes for metal (im2col & upsample_2d) (#1938)

Browse Source 
* update im2col dtype implementations

* update dtypes for upsample
This commit is contained in:
Thomas Santerre
2024-03-26 01:48:56 -04:00
committed by GitHub
parent 196765e995
commit f5dfe883d7
2 changed files with 16 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-metal-kernels/src/conv.metal
View File

@ -486,16 +486,24 @@ kernel void FN_NAME( \
} \
IM2COL_OP(float, im2col_f32)
IM2COL_OP(half, im2col_f16)
IM2COL_OP(uint8_t, im2col_u8)
IM2COL_OP(uint32_t, im2col_u32)
#if defined(__HAVE_BFLOAT__)
IM2COL_OP(bfloat, im2col_bf16)
#endif
IM2COL1D_OP(float, im2col1d_f32)
IM2COL1D_OP(uint8_t, im2col1d_u8)
IM2COL1D_OP(uint32_t, im2col1d_u32)
UPSAMPLE_NEAREST2D_OP(float, upsample_nearest2d_f32)
UPSAMPLE_NEAREST2D_OP(half, upsample_nearest2d_f16)
UPSAMPLE_NEAREST2D_OP(uint8_t, upsample_nearest2d_u8)
UPSAMPLE_NEAREST2D_OP(uint32_t, upsample_nearest2d_u32)
#if defined(__HAVE_BFLOAT__)
UPSAMPLE_NEAREST2D_OP(bfloat, upsample_nearest2d_bf16)
#endif
MAXPOOL2D_OP(float, max_pool2d_f32)
MAXPOOL2D_OP(half, max_pool2d_f16)