为 Ascend 910B3 (dav-2201) 实现原生 ScatterAdd kernel,替换原有的 CPU 回退路径(D2H → CPU ScatterAdd → H2D)。
参照 CUDA 实现 src/array/cuda/segment_reduce.cuh:54 的 ScatterAddKernel,采用 AscendC MemBase 范式实现。
src/array/cuda/segment_reduce.cuh:54
ScatterAddKernel
将 ScatterAddDispatch(src/array/kernel.cc)中 Ascend 设备的调度路径从 CPU 回退改为原生 NPU kernel 调用(ATEN_XPU_SWITCH_CUDA_ASCEND),与仓内 SegmentReduce/SpMM/SDDMM 等 Ascend 算子的 dispatch 范式一致。
ScatterAddDispatch
src/array/kernel.cc
ATEN_XPU_SWITCH_CUDA_ASCEND
SegmentReduce
SpMM
SDDMM
ScatterAdd(scatter add on first dimension)
out[idx[i], *] += feat[i, *]
F.zeros
segment_reduce_sum_kernel.cpp
SetAtomicAdd<float>()
DataCopyPad
atomicAdd
blockDim = min(N, 40)
src/array/ascend/scatter_add_kernel.cpp
src/array/ascend/scatter_add.cc
tests/ascend/
docs/ascend/operators/scatter_add.md
Description
为 Ascend 910B3 (dav-2201) 实现原生 ScatterAdd kernel,替换原有的 CPU 回退路径(D2H → CPU ScatterAdd → H2D)。
参照 CUDA 实现
src/array/cuda/segment_reduce.cuh:54的ScatterAddKernel,采用 AscendC MemBase 范式实现。修改功能
将
ScatterAddDispatch(src/array/kernel.cc)中 Ascend 设备的调度路径从 CPU 回退改为原生 NPU kernel 调用(ATEN_XPU_SWITCH_CUDA_ASCEND),与仓内SegmentReduce/SpMM/SDDMM等 Ascend 算子的 dispatch 范式一致。适配算子
ScatterAdd(scatter add on first dimension)
out[idx[i], *] += feat[i, *]F.zeros),与 CUDA/CPU 实现一致支持的数据类型
实现方案
segment_reduce_sum_kernel.cppSetAtomicAdd<float>()+DataCopyPadUB→GM 原子累加(与 CUDAatomicAdd语义一致)blockDim = min(N, 40)验收结果
Checklist
Changes