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

Support for where-cond on cuda for u8 and u32. (#274)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-29 11:48:58 +01:00
committed by GitHub
parent 4bf2ebf836
commit c0a8ed19eb
2 changed files with 26 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -940,16 +940,22 @@ impl<'a> Map2 for WhereCond<'a> {
dev: &CudaDevice,
) -> Result<CudaSlice<T>> {
let ids_l = &self.1;
let ids = match &self.0.slice {
CudaStorageSlice::U32(slice) => slice.slice(ids_l.start_offset()..),
let (ids, name) = match &self.0.slice {
CudaStorageSlice::U8(slice) => {
let ptr = *slice.slice(ids_l.start_offset()..).device_ptr();
(ptr, "where_u8")
}
CudaStorageSlice::U32(slice) => {
let ptr = *slice.slice(ids_l.start_offset()..).device_ptr();
(ptr, "where_u32")
}
_ => Err(CudaError::UnexpectedDType {
msg: "where conditions should be u32",
msg: "where conditions should be u8 or u32",
expected: DType::U32,
got: self.0.dtype(),
})
.w()?,
};
let ids = &ids;
let shape = ids_l.shape();
let dims = shape.dims();
let el = shape.elem_count();
@ -959,7 +965,7 @@ impl<'a> Map2 for WhereCond<'a> {
.w()?;
let t = &t.slice(layout_t.start_offset()..);
let f = &f.slice(layout_f.start_offset()..);
let func = dev.get_or_load_func(&kernel_name::<T>("where"), kernels::TERNARY)?;
let func = dev.get_or_load_func(&kernel_name::<T>(name), kernels::TERNARY)?;
// SAFETY: Set later by running the kernel.
let out = unsafe { dev.alloc::<T>(el) }.w()?;
let params = (el, dims.len(), &ds, ids, t, f, &out);