Star5
12
代码
代码
Issues10
Pull Requests2
项目讨论
Wiki
分析
Star5
12
已合并
feat(grad_minmax_bwd_cmp): adapt UpdateGradMinMax_hetero and BackwardSegmentCmp for Ascend NPU #23
lhp_lhp创建于 25 天前
feat(grad_minmax_bwd_cmp): adapt UpdateGradMinMax_hetero and BackwardSegmentCmp for Ascend NPU #23
已合并
lhp_lhp创建于 25 天前
从lhp_lhp/dgl-ascendfeat/grad-minmax-bwd-cmp-ascend合入到AI4Science/dgl-ascendmaster
L
lhp_lhp成员
25 天前

Description

为 Ascend 910B3 (dav-2201) 适配 2 个 segment reduce 反向算子的 NPU 原生路径，替换 CPU 回退和 LOG(FATAL)。

适配算子

算子 数学语义 实现方式 操作类型
BackwardSegmentCmp out[arg[i,k], k] = feat[i,k] (arg>=0) 直接赋值（无原子，arg 列内唯一无写冲突） scatter write
UpdateGradMinMax_hetero if (type==idx_type[r,c]) out[idx[r,c], c] += feat[r,c] 条件性原子累加（SetAtomicAdd） conditional scatter add

实现方案

  • 架构：SIMD / MemBase，参照 scatter_add_kernel.cpp
  • 核心设计：单元素散列写（scratchBuf + GetValue/SetValue + DataCopyPad blockLen=4）
  • 同步机制：复用 scatter_add 验证过的范式（PipeBarrier + SetFlag/WaitFlag + SetAtomicNone）
  • 多核切分：按 N 切分，blockDim=min(N,40)
  • host launcher：hetero 遍历 etype 串行 launch（与 CUDA/CPU 一致）

支持的数据类型

IdType Dtype 状态
int32 float32 ✅ 原生支持
int64 float32 ✅ 支持（int64 idx/arg 转 int32，同 segment_reduce.cc 范式）
int32 half (float16) ❌ LOG(FATAL) 占位
int64 half (float16) ❌ LOG(FATAL) 占位
int32 bfloat16 ❌ LOG(FATAL) 占位
int64 bfloat16 ❌ LOG(FATAL) 占位
int32 double (float64) ❌ LOG(FATAL) 占位
int64 double (float64) ❌ LOG(FATAL) 占位

两个算子同时涉及 IdType（索引类型 int32/int64）和 DType（特征数据类型 float32）。IdType 用于 arg/idx/idx_etype 数组，DType 用于 feat/out 张量。

Checklist

Changes

Test Results

  • 58/59 passed, 1 xfail（DGL 绑定层 use-after-free，非算子缺陷）
  • QA 独立探针 43/43 pass
  • BackwardSegmentCmp 全部 bit-exact
  • UpdateGradMinMax_hetero max_abs_error=3.81e-6（远低于 1e-2 阈值）
likedislike
Pull Request已成功合入, 合并人@knightgd
（感谢 lhp_lhp 的贡献）
Llhp_lhp成员
25 天前 修改了pull request 的描述
Llhp_lhp成员
25 天前 修改了pull request 的描述
Llhp_lhp成员
25 天前 修改了pull request 的描述
Llhp_lhp成员
25 天前 强制推送  1 个提交:e7482e6f-feat(grad_minmax_bwd_cmp): adapt UpdateGradMinMax_hetero and BackwardSegmentCmp for Ascend NPU
Llhp_lhp成员
20 天前 关联了issue：UpdateGradMinMax_hetero算子支持NPU，BackwardSegmentCmp算子支持NPU
Llhp_lhp成员
20 天前 删除了关联的issue：BackwardSegmentCmp算子支持NPU
Llhp_lhp成员
20 天前 将 lhp_lhp 设为合并人
Llhp_lhp成员
20 天前 移除了合并人lhp_lhp
Llhp_lhp成员
20 天前 将 lhp_lhp 设为合并人
Llhp_lhp成员
20 天前 将 knightgd 设为合并人
Llhp_lhp成员
20 天前 移除了合并人lhp_lhp
K
knightgd成员19 天前进行代码检视3
src/array/ascend/backward_segment_cmp.cc
已过期
@@ -0,0 +85,4 @@
85+template <typename DType>
86+void BackwardSegmentCmpAscendImpl(NDArray feat, NDArray arg, NDArray out) {
87+ DGLContext ctx = feat->ctx;
88+ ASCEND_CALL(aclrtSynchronizeDevice());
K
knightgd19 天前评论：

host侧的线程需要等待device侧的所有stream流任务执行完才能往下执行，阻塞的范围过大，改成stream粒度的同步

likedislike
System
系统消息系统
7 天前 评论：

changed this line on f7d7475c view diff detail

L
lhp_lhp成员
6 天前 评论：

aclrtSynchronizeDevice()（设备级、等待所有流）已替换为流粒度同步：入口和 int64→int32 转换前通过 c10_npu::getCurrentNPUStream() 获取 torch 当前流并调用 torch 自己的 stream_synchronize（torch_npu 在独立 ACL context 上执行，DGL 默认 context 的流同步无法可靠等待它）；无 torch_npu 时回退 getCurrentAscendStream()。

排查中发现并顺带修复了当前 torch_npu 2.7.1.post4 环境下 int64 路径的两个问题：①框架阻塞式 D2H 拷贝跨 context 读到残缺数据；②AsNumBits AscendC kernel 仅转换前 n/GetBlockNum() 个元素的潜在 bug（将另立 issue）。int64→int32 窄化改为 torch 流上的流序化异步拷贝 + host 转换（ConvertIdx64To32）。

K
knightgd成员19 天前进行代码检视3
src/array/ascend/backward_segment_cmp.cc
已过期
@@ -0,0 +121,4 @@
121+#ifdef DGL_USE_ASCEND
122+ // int64 arg -> int32 conversion.
123+ // Synchronize the device before D2H CopyTo to avoid stream race.
124+ ASCEND_CALL(aclrtSynchronizeDevice());
K
knightgd19 天前评论：

host侧的线程需要等待device侧的所有stream流任务执行完才能往下执行，阻塞的范围过大，改成stream粒度的同步

likedislike
System
系统消息系统
7 天前 评论：

changed this line on f7d7475c view diff detail

L
lhp_lhp成员
6 天前 评论：

aclrtSynchronizeDevice()（设备级、等待所有流）已替换为流粒度同步：入口和 int64→int32 转换前通过 c10_npu::getCurrentNPUStream() 获取 torch 当前流并调用 torch 自己的 stream_synchronize（torch_npu 在独立 ACL context 上执行，DGL 默认 context 的流同步无法可靠等待它）；无 torch_npu 时回退 getCurrentAscendStream()。

排查中发现并顺带修复了当前 torch_npu 2.7.1.post4 环境下 int64 路径的两个问题：①框架阻塞式 D2H 拷贝跨 context 读到残缺数据；②AsNumBits AscendC kernel 仅转换前 n/GetBlockNum() 个元素的潜在 bug（将另立 issue）。int64→int32 窄化改为 torch 流上的流序化异步拷贝 + host 转换（ConvertIdx64To32）。

K
knightgd成员18 天前进行代码检视3
src/array/ascend/update_grad_minmax_hetero.cc
已过期
@@ -0,0 +104,4 @@
104+ }
105+ 
106+ DGLContext ctx = feat.front()->ctx;
107+ ASCEND_CALL(aclrtSynchronizeDevice());
K
knightgd18 天前评论：

host侧的线程需要等待device侧的所有stream流任务执行完才能往下执行，阻塞的范围过大，改成stream粒度的同步

likedislike
System
系统消息系统
7 天前 评论：

changed this line on f7d7475c view diff detail

L
lhp_lhp成员
6 天前 评论：

aclrtSynchronizeDevice()（设备级、等待所有流）已替换为流粒度同步：入口和 int64→int32 转换前通过 c10_npu::getCurrentNPUStream() 获取 torch 当前流并调用 torch 自己的 stream_synchronize（torch_npu 在独立 ACL context 上执行，DGL 默认 context 的流同步无法可靠等待它）；无 torch_npu 时回退 getCurrentAscendStream()。

排查中发现并顺带修复了当前 torch_npu 2.7.1.post4 环境下 int64 路径的两个问题：①框架阻塞式 D2H 拷贝跨 context 读到残缺数据；②AsNumBits AscendC kernel 仅转换前 n/GetBlockNum() 个元素的潜在 bug（将另立 issue）。int64→int32 窄化改为 torch 流上的流序化异步拷贝 + host 转换（ConvertIdx64To32）。

K
knightgd成员18 天前进行代码检视3
src/array/ascend/update_grad_minmax_hetero.cc
已过期
@@ -0,0 +166,4 @@
166+ // Synchronize the device before D2H CopyTo to avoid stream race.
167+ DGLContext ctx = feat.front()->ctx;
168+ DGLContext cpu_ctx{kDGLCPU, 0};
169+ ASCEND_CALL(aclrtSynchronizeDevice());
K
knightgd18 天前评论：

host侧的线程需要等待device侧的所有stream流任务执行完才能往下执行，阻塞的范围过大，改成stream粒度的同步

likedislike
System
系统消息系统
7 天前 评论：

changed this line on f7d7475c view diff detail

L
lhp_lhp成员
6 天前 评论：

aclrtSynchronizeDevice()（设备级、等待所有流）已替换为流粒度同步：入口和 int64→int32 转换前通过 c10_npu::getCurrentNPUStream() 获取 torch 当前流并调用 torch 自己的 stream_synchronize（torch_npu 在独立 ACL context 上执行，DGL 默认 context 的流同步无法可靠等待它）；无 torch_npu 时回退 getCurrentAscendStream()。

排查中发现并顺带修复了当前 torch_npu 2.7.1.post4 环境下 int64 路径的两个问题：①框架阻塞式 D2H 拷贝跨 context 读到残缺数据；②AsNumBits AscendC kernel 仅转换前 n/GetBlockNum() 个元素的潜在 bug（将另立 issue）。int64→int32 窄化改为 torch 流上的流序化异步拷贝 + host 转换（ConvertIdx64To32）。

Llhp_lhp成员
7 天前 推送  1 个提交:f7d7475c-fix(grad_minmax_bwd_cmp): replace device-wide sync with stream-granular sync
Llhp_lhp成员
6 天前 推送  1 个提交:ce47399d-fix(grad_minmax_bwd_cmp): reliable int64->int32 narrowing on current env
Kknightgd成员
6 天前 合入了pull request，合并节点 SHA：57b37d809fe1f1776a0798f5a6daebacfc556ce9
Llhp_lhp成员
11 小时前 关联了issue：BackwardSegmentCmp算子支持NPU