已关闭
fix(ascend): Range int32 与 CSRToCOO 原生算子修复,以及同步性能优化(PR #44) #38
xuejikan11创建于  11 天前关闭于  9 天前
xuejikan11成员
11 天前 创建

关联 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.cc InEdges/OutEdges/InDegrees/OutDegrees/EdgeSubgraph、graph_apis.cc VertexSubgraph/EdgeSubgraph、randomwalks.cc):python 流桥(_ffi/base.py 每次 FFI 前 torch.npu.synchronize())已保证 torch→DGL 输入排序。
  • 出口处 aclrtSynchronizeDevice() 降级为 aclrtSynchronizeStream(nullptr)(unit_graph.cc CSR/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 通过。
  • Range(经 node_subgraph 触发,int32/int64,2~49820 条边含多 block)与 CPU 结果一致。
  • 全量 test_nn.py:2969 passed / 0 failed(910B3,CANN 9.0.0)。
  • 性能:删除冗余 device sync 后 D2H 重路径约提速 9.5%,其余 1-5%(共享宿主机噪声内);单算子固定开销仍以 AscendC 算子执行路径为主。

后续

  • 探索算子内 aclrtSynchronizeStream host-barrier(约 77 处)的移除与 AscendC launch 开销优化,进一步降低单算子 ~1.3ms 固定开销。
likedislike
Xxuejikan11成员
9 天前 关闭了 issue