Enable the affine kernel for u8/u32.

huggingface · Aug 1, 2024 · 00c675b · 00c675b
1 parent 4a52aeb
commit 00c675b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/candle-core/src/metal_backend/mod.rs b/candle-core/src/metal_backend/mod.rs
@@ -119,6 +119,8 @@ impl BackendStorage for MetalStorage {
  DType::F32 => "affine_f32",
  DType::F16 => "affine_f16",
  DType::BF16 => "affine_bf16",
+ DType::U8 => "affine_u8",
+ DType::U32 => "affine_u32",
  dtype => crate::bail!("Metal contiguous affine {dtype:?} not implemented"),
  };
  candle_metal_kernels::call_affine(

diff --git a/candle-metal-kernels/src/affine.metal b/candle-metal-kernels/src/affine.metal
@@ -109,6 +109,8 @@ kernel void FN_NAME##_strided( \
 } \
 
 
+AFFINE(affine_u8, uint8_t)
+AFFINE(affine_u32, uint32_t)
 AFFINE(affine_f32, float)
 AFFINE(affine_f16, half)
 POWF(powf_f32, float)