Nigel
/
jittor
forked from jittor/jittor
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update nccl_broadcast_op.cc 

support uint8 in nccl
This commit is contained in:
Xiang-Li Li 2021-09-25 20:48:18 +08:00 committed by GitHub
parent 764c8425af
commit beebddafd5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
python/jittor/extern/cuda/nccl/ops/nccl_broadcast_op.cc vendored
View File

@ -47,6 +47,7 @@ void NcclBroadcastOp::jit_run() {
@if(@strcmp(@Tx,int)==0 || @strcmp(@Tx,int32)==0, ncclInt)
@if(@strcmp(@Tx,float64)==0, ncclFloat64)
@if(@strcmp(@Tx,int64)==0, ncclInt64)
@if(@strcmp(@Tx,uint8)==0, ncclUint8)
)
auto* __restrict__ xp = x->ptr<Tx>();
auto* __restrict__ yp = y->ptr<Tx>();