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

Add a test for scatter add. (#238)

Browse Source 
* Add a test for scatter add (segfaults on gpus for now).

* Bugfix for the scatter add cuda kernel.
This commit is contained in:
Laurent Mazare
2023-07-25 09:12:14 +01:00
committed by GitHub
parent 18cc73954a
commit 944d70bd9a
3 changed files with 47 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -846,7 +846,7 @@ impl<'a> Map2InPlace for ScatterAdd<'a> {
fn f<T: DeviceRepr + WithDType + ValidAsZeroBits>(
&self,
dst: &mut CudaSlice<T>,
_dst_shape: &Shape,
dst_shape: &Shape,
src: &CudaSlice<T>,
src_l: &Layout,
dev: &CudaDevice,
@ -874,11 +874,11 @@ impl<'a> Map2InPlace for ScatterAdd<'a> {
let left_sz: usize = src_l.dims()[..dim].iter().product();
let right_sz: usize = src_l.dims()[dim + 1..].iter().product();
let src_dim_sz = src_l.dims()[dim];
let ids_dim_sz = ids_l.dims()[dim];
let dst_dim_sz = dst_shape.dims()[dim];
let cfg = LaunchConfig::for_num_elems((left_sz * right_sz) as u32);
let func = dev.get_or_load_func(&kernel_name::<T>(name), kernels::INDEXING)?;
// SAFETY: Set later by running the kernel.
let params = (ids, &src, dst, left_sz, src_dim_sz, ids_dim_sz, right_sz);
let params = (ids, &src, dst, left_sz, src_dim_sz, dst_dim_sz, right_sz);
// SAFETY: ffi.
unsafe { func.launch(cfg, params) }.w()?;
Ok(())