为 Ascend 910B3 (dav-2201) 适配 2 个 segment reduce 反向算子的 NPU 原生路径,替换 CPU 回退和 LOG(FATAL)。
out[arg[i,k], k] = feat[i,k]
if (type==idx_type[r,c]) out[idx[r,c], c] += feat[r,c]
两个算子同时涉及 IdType(索引类型 int32/int64)和 DType(特征数据类型 float32)。IdType 用于 arg/idx/idx_etype 数组,DType 用于 feat/out 张量。
src/array/ascend/segment_reduce_grad_kernel.cpp
src/array/ascend/segment_reduce_grad_hetero_kernel.cpp
src/array/ascend/segment_reduce_grad.cc
src/array/kernel.cc
tests/ascend/test_grad_minmax_bwd_cmp_npu.py
Description
为 Ascend 910B3 (dav-2201) 适配 2 个 segment reduce 反向算子的 NPU 原生路径,替换 CPU 回退和 LOG(FATAL)。
适配算子
out[arg[i,k], k] = feat[i,k](arg>=0)if (type==idx_type[r,c]) out[idx[r,c], c] += feat[r,c]实现方案
支持的数据类型
Checklist
Changes
Test Results