本 Issue 对应 MR #44:https://gitcode.com/AI4Science/dgl-ascend/merge_requests/44
当前仓库中 Range(int32)与 CSRToCOO 在 Ascend 上存在 CPU 兜底实现,且多处使用 aclrtSynchronizeDevice() 做全设备同步,存在正确性与性能问题。
Range
CSRToCOO
aclrtSynchronizeDevice()
range.cc
tiling.dtype_bytes
GetTilingDev()
dtype_bytes
Range(0,3)
[0,0,1]
csr_to_coo_kernel.cpp
GlobalTensor::SetValue
DataCopyPad
csr_get_row_nnz
csr_to_coo.cc
heterograph_capi.cc
graph_apis.cc
randomwalks.cc
_ffi/base.py
torch.npu.synchronize()
aclrtSynchronizeStream(nullptr)
unit_graph.cc
array.cc
CopyTo
network_emb.py
torch.optim.SparseAdam.step
__init__
DGL_NPU_SPARSE_ADAM_COMPAT=0
tensor.py
test_csr_to_coo_npu.py
node_subgraph
test_nn.py
aclrtSynchronizeStream
关联 PR
本 Issue 对应 MR #44:https://gitcode.com/AI4Science/dgl-ascend/merge_requests/44
背景
当前仓库中
Range(int32)与CSRToCOO在 Ascend 上存在 CPU 兜底实现,且多处使用aclrtSynchronizeDevice()做全设备同步,存在正确性与性能问题。本次改动内容
1. Range int32 / CSRToCOO 原生算子修复(去掉 CPU 兜底)
range.cc:原 bug 在 host 侧——tiling.dtype_bytes在GetTilingDev()(host→device 拷贝)之后才赋值,导致 kernel 读到未初始化的dtype_bytes,恒走 int64 分支,int32 缓冲被 8 字节写坏(Range(0,3)输出[0,0,1])。修复:在拷贝前赋值。csr_to_coo_kernel.cpp:原实现用GlobalTensor::SetValue逐字写 GM,非零 block 的写会被静默丢弃。重写为 UB LocalTensor staging +DataCopyPad(参考csr_get_row_nnz的成熟模式)。csr_to_coo.cc:恢复原生多 block launch(按行划分 vector core),去掉 CPU 兜底。2. 冗余全设备同步优化
aclrtSynchronizeDevice()(heterograph_capi.ccInEdges/OutEdges/InDegrees/OutDegrees/EdgeSubgraph、graph_apis.ccVertexSubgraph/EdgeSubgraph、randomwalks.cc):python 流桥(_ffi/base.py每次 FFI 前torch.npu.synchronize())已保证 torch→DGL 输入排序。aclrtSynchronizeDevice()降级为aclrtSynchronizeStream(nullptr)(unit_graph.ccCSR/CSC 物化、heterograph_capi.cc结果可见性):DGL kernel 全在 ACL 默认流上,宿主等到默认流完成即可,无需等待 torch_npu 无关流。array.cc中 D2H 前的aclrtSynchronizeDevice()(CopyTo内部已同步默认流)。3. 其他
network_emb.py:去掉 import 时全局劫持torch.optim.SparseAdam.step,改为 DeepWalk/MetaPath2Vec__init__内惰性安装(DGL_NPU_SPARSE_ADAM_COMPAT=0可关)。tensor.py中 int32 min/max 的 NPU 注释(该兜底在 CANN 9.0.0 上仍然必要)。验证
test_csr_to_coo_npu.py:8/8 通过。node_subgraph触发,int32/int64,2~49820 条边含多 block)与 CPU 结果一致。test_nn.py:2969 passed / 0 failed(910B3,CANN 9.0.0)。后续
aclrtSynchronizeStreamhost-barrier(约 77 处)的移除与 AscendC launch 开销优化,进一步降低单算子 ~1.3ms 固定开销。