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

Use Map1 on unary ops.

Browse Source
This commit is contained in:
laurent
2023-06-29 09:37:38 +01:00
parent fff13dbb4e
commit 8ad03a5fb6
2 changed files with 33 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
candle-core/src/cuda_backend.rs
View File

@ -350,6 +350,29 @@ impl<'a> Map1 for Sum<'a> {
} }
} }
impl<U: crate::op::UnaryOp> Map1 for U {
fn f<T: DeviceRepr + WithDType + ValidAsZeroBits>(
&self,
src: &CudaSlice<T>,
dev: &CudaDevice,
layout: &Layout,
) -> Result<CudaSlice<T>> {
let shape = layout.shape();
let dims = shape.dims();
let el_count = shape.elem_count();
let cfg = LaunchConfig::for_num_elems(el_count as u32);
let ds = dev.htod_copy([dims, layout.stride()].concat())?;
let src = &src.slice(layout.start_offset()..);
let func = dev.get_or_load_func(&kernel_name::<T>(U::KERNEL), kernels::UNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<T>(el_count) }?;
let params = (el_count, dims.len(), &ds, src, &out);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }?;
Ok(out)
}
}
fn slice_src_and_dst<'a, T>( fn slice_src_and_dst<'a, T>(
src: &'a CudaSlice<T>, src: &'a CudaSlice<T>,
src_l: &Layout, src_l: &Layout,
@ -539,58 +562,8 @@ impl CudaStorage {
} }
pub(crate) fn unary_impl<U: crate::op::UnaryOp>(&self, layout: &Layout) -> Result<Self> { pub(crate) fn unary_impl<U: crate::op::UnaryOp>(&self, layout: &Layout) -> Result<Self> {
let shape = layout.shape(); let device = self.device().clone();
let dims = shape.dims(); let slice = U::V.map(&self.slice, &device, layout)?;
let el_count = shape.elem_count();
let cfg = LaunchConfig::for_num_elems(el_count as u32);
let dev = &self.device;
let ds = dev.htod_copy([dims, layout.stride()].concat())?;
let slice = match &self.slice {
CudaStorageSlice::U32(_arg) => {
todo!("No unary kernels for u32");
}
CudaStorageSlice::BF16(arg) => {
let arg = &arg.slice(layout.start_offset()..);
let func = dev.get_or_load_func(U::KERNEL_BF16, kernels::UNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<bf16>(el_count) }?;
let params = (el_count, dims.len(), &ds, arg, &out);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }?;
CudaStorageSlice::BF16(out)
}
CudaStorageSlice::F16(arg) => {
let arg = &arg.slice(layout.start_offset()..);
let func = dev.get_or_load_func(U::KERNEL_F16, kernels::UNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<f16>(el_count) }?;
let params = (el_count, dims.len(), &ds, arg, &out);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }?;
CudaStorageSlice::F16(out)
}
CudaStorageSlice::F32(arg) => {
let arg = &arg.slice(layout.start_offset()..);
let func = dev.get_or_load_func(U::KERNEL_F32, kernels::UNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<f32>(el_count) }?;
let params = (el_count, dims.len(), &ds, arg, &out);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }?;
CudaStorageSlice::F32(out)
}
CudaStorageSlice::F64(arg) => {
let arg = &arg.slice(layout.start_offset()..);
let func = dev.get_or_load_func(U::KERNEL_F64, kernels::UNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<f64>(el_count) }?;
let params = (el_count, dims.len(), &ds, arg, &out);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }?;
CudaStorageSlice::F64(out)
}
};
let device = dev.clone();
Ok(Self { slice, device }) Ok(Self { slice, device })
} }

28
candle-core/src/op.rs
View File

@ -43,11 +43,8 @@ pub(crate) enum Op {
pub(crate) trait UnaryOp { pub(crate) trait UnaryOp {
const NAME: &'static str; const NAME: &'static str;
const KERNEL_BF16: &'static str; const KERNEL: &'static str;
const KERNEL_F16: &'static str; const V: Self;
const KERNEL_F32: &'static str;
const KERNEL_F64: &'static str;
const KERNEL_U32: &'static str;
fn bf16(v1: bf16) -> bf16; fn bf16(v1: bf16) -> bf16;
fn f16(v1: f16) -> f16; fn f16(v1: f16) -> f16;
fn f32(v1: f32) -> f32; fn f32(v1: f32) -> f32;
@ -121,11 +118,8 @@ macro_rules! unary_op {
($op: ident, $name: literal, $a: ident, $e: expr) => { ($op: ident, $name: literal, $a: ident, $e: expr) => {
impl UnaryOp for $op { impl UnaryOp for $op {
const NAME: &'static str = $name; const NAME: &'static str = $name;
const KERNEL_BF16: &'static str = concat!("u", $name, "_bf16"); const KERNEL: &'static str = concat!("u", $name);
const KERNEL_F16: &'static str = concat!("u", $name, "_f16"); const V: Self = $op;
const KERNEL_F32: &'static str = concat!("u", $name, "_f32");
const KERNEL_F64: &'static str = concat!("u", $name, "_f64");
const KERNEL_U32: &'static str = concat!("u", $name, "_u32");
fn bf16($a: bf16) -> bf16 { fn bf16($a: bf16) -> bf16 {
$e $e
} }
@ -158,6 +152,7 @@ unary_op!(Sqrt, "sqrt", v, v.sqrt());
/// <https://en.wikipedia.org/wiki/Activation_function#Comparison_of_activation_functions> /// <https://en.wikipedia.org/wiki/Activation_function#Comparison_of_activation_functions>
impl UnaryOp for Gelu { impl UnaryOp for Gelu {
const NAME: &'static str = "gelu"; const NAME: &'static str = "gelu";
const V: Self = Gelu;
fn bf16(v: bf16) -> bf16 { fn bf16(v: bf16) -> bf16 {
bf16::from_f32_const(0.5) bf16::from_f32_const(0.5)
* v * v
@ -191,20 +186,13 @@ impl UnaryOp for Gelu {
fn u32(_: u32) -> u32 { fn u32(_: u32) -> u32 {
0 0
} }
const KERNEL_BF16: &'static str = "ugelu_bf16"; const KERNEL: &'static str = "ugelu";
const KERNEL_F16: &'static str = "ugelu_f16";
const KERNEL_F32: &'static str = "ugelu_f32";
const KERNEL_F64: &'static str = "ugelu_f64";
const KERNEL_U32: &'static str = "ugelu_u32";
} }
impl UnaryOp for Relu { impl UnaryOp for Relu {
const NAME: &'static str = "relu"; const NAME: &'static str = "relu";
const KERNEL_BF16: &'static str = "urelu_bf16"; const KERNEL: &'static str = "urelu";
const KERNEL_F16: &'static str = "urelu_f16"; const V: Self = Relu;
const KERNEL_F32: &'static str = "urelu_f32";
const KERNEL_F64: &'static str = "urelu_f64";
const KERNEL_U32: &'static str = "urelu_u32";
fn bf16(v: bf16) -> bf16 { fn bf16(v: bf16) -> bf16 {
v.max(bf16::ZERO) v.max(bf16::ZERO)
} }