Skip to content

feat(edge_softmax): add edge_softmax and CSRTranspose NPU support for Ascend - #15

Closed
xuejiakn wants to merge 10 commits into
BUPT-GAMMA:masterfrom
xuejiakn:feat-edge-softmax-ascend
Closed

feat(edge_softmax): add edge_softmax and CSRTranspose NPU support for Ascend#15
xuejiakn wants to merge 10 commits into
BUPT-GAMMA:masterfrom
xuejiakn:feat-edge-softmax-ascend

Conversation

@xuejiakn

@xuejiakn xuejiakn commented Jul 22, 2026

Copy link
Copy Markdown

Summary

Add Ascend NPU adaptation for edge_softmax operator, enabling GAT/PAGTN graph attention network training on Ascend 910B NPU.

Changes

New Operators

Operator Files Description
edge_softmax (forward + backward) edge_softmax_kernel.cpp, edge_softmax_tiling.h, edge_softmax.cc Ascend C kernel with FullLoad/RowSplit dual-mode, FP32/FP16, AR/ARA (num_heads=1/>1) branches. Backward adapted to DGL interface (sds = out * grad_out).
CSRTranspose csr_transpose.cc Ascend implementation via COOToCSR(COOTranspose(CSRToCOO)) chain, reusing existing Ascend kernels.

Framework Integration

  • kernel.cc: Dispatch edge_softmax by data tensor device (not graph context) to handle NPU tensors with CPU-resident graph topology. Add CPU fallback for BackwardSegmentCmp and ScatterAdd (needed by GAT WeightedSumAndMax readout).
  • array.cc: Add Ascend dispatch branch for CSRTranspose.

Host Code (edge_softmax.cc)

  • ACL host with forward/backward template specializations (FP32/FP16 x int32/int64)
  • CSC edge ID remapping via CPU gather/scatter (NPU torch index ops unreliable on DGL blob tensors)
  • Backward interface adapted to DGL convention: back_out = sds - out * sum(sds)

Test Coverage

Unit Tests (tests/ascend/test_edge_softmax_npu.py)

78 tests, 100% pass rate:

Category Cases Description
Forward 36 FP32/FP16 x heads=1/4/8 x 6 graph types
Backward 36 FP32/FP16 x heads=1/4/8 x 6 graph types
Autograd 1 Forward + backward chain with weight multiplication
Sum-to-one 1 NPU vs CPU consistency
Single edge 1 Edge case: single edge -> output = 1.0
Zero-degree 1 Edge case: node with no incoming edges
Large degree 1 RowSplit path (clique graph, degree=300)
Performance 1 Latency baseline (no assert)

End-to-End Validation

Test Result Notes
pytest test_gat.py::test_gat_classification PASS GAT model forward + backward
pytest test_gat.py::test_gat_reload PASS Model save/load/predict consistency
pytest test_gat.py::test_gat_regression PASS (800ep) Needs more epochs than default 500 (CPU also needs more)
pytest test_pagtn.py No crash PAGTN also uses edge_softmax; precision is test-epoch limit

Dependencies

Co-authored-by

CANNBot (Ascend C operator development tool)

  • Complete design docs: DESIGN.md, PLAN.md, REVIEW.md, precision/perf reports
  • Operator project: operators/edge_softmax/

lizengyong and others added 10 commits June 12, 2026 16:27
@lizengyong
图神经网络分布式训练中图数据划分、分布式图存储、分布式邻居采样、跨设备通信、前向传播、反向传播、梯度同步和参数更新关键步骤适配昇腾NPU
700239c
@lizengyong
修复GFORTRAN_LIB和GOMP_LIB路径问题
efe10fe
@lizengyong
修复nd2nzA1Params.nValue和nd2nzA1Params.dValue赋值错误问题
b1a0730
@xuejiakn
feat(sddmm): add SDDMM operator with NPU-native kernels for Ascend
173ec02 
- Add sddmm_dot_kernel.cpp: Ascend C kernel for dot op (FP32/FP16, COO format)
- Add sddmm_copy_lhs_kernel.cpp: NPU-native gather kernel for copy_lhs/copy_rhs
  (avoids CPU roundtrip that caused PyTorch NPU stream sync issues with autograd)
- Add sddmm.cc: host dispatch supporting all ops (dot/add/sub/mul/div/copy_lhs/copy_rhs)
  - dot: NPU kernel (COO + CSR via CSRToCOO)
  - copy_lhs/copy_rhs: NPU gather kernel (no CPU roundtrip)
  - add/sub/mul/div: CPU fallback (FP32 only)
  - int64 index: cast to int32 then delegate to int32 path
- Modify kernel.cc: SDDMM/CSRSDDMM/COOSDDMM dispatch macros to ATEN_XPU_SWITCH_CUDA_ASCEND
@xuejiakn
feat(spmm): add SPMM FP32 NPU-native kernel and fix stream management
da4bb7d 
- Add spmm_kernel.cpp: Ascend C kernel for SpMM (copy_lhs + sum/max/min, FP32/FP16)
  - Multi-core: split by dst rows, dynamic core count
  - Double buffer feature queue (BUFFER_NUM=2)
  - FP16: cast to FP32 for accumulation, cast back
  - Tail core safety: load indptr by actualRows (not rowsPerCore+1)
- Add spmm_tiling.h: tiling struct shared between kernel and host
- Modify spmm.cc:
  - FP32: use new NPU kernel (SpMMCSRNpuF32) instead of CPU fallback
  - int64 index: cast CSR to int32 via CastCSRToInt32SpMM
  - Non-copy_lhs ops (mul/add/copy_rhs): CPU fallback
  - Use default stream (nullptr) instead of custom stream to align with PyTorch NPU
  - Remove use_pytorch_stream() env var dependency
@xuejiakn
fix(npu): sync PyTorch NPU stream for graph-structure queries and seg…
007a5ba 
…ment_reduce

- streams.py: to_dgl_stream_handle() now supports torch.npu.Stream (.npu_stream)
  in addition to torch.cuda.Stream (.cuda_stream)
- heterograph.py:
  - DGLGraph.to() calls record_stream() with current PyTorch NPU stream
  - New _sync_npu() helper calls torch.npu.synchronize() before/after every
    structure query (in_degrees, out_degrees, edges, in_edges) on NPU graphs
- _sparse_ops.py: add torch.npu.synchronize() before _CAPI_DGLKernelSegmentReduce
  to prevent data races between PyTorch NPU stream and DGL default ACL stream
- core.py: fix assert to use int(deg) for numpy type compatibility
@xuejiakn
fix(npu): add Ascend CPU fallback for COO ops and fix segment_reduce …
58265b0 
…stream

- array.cc: add CPU fallback with aclrtSynchronizeDevice() for 5 COO operations
  that previously only dispatched to CPU/CUDA (not Ascend):
  - COOGetRowNNZ (scalar + array)
  - COOGetRowDataAndIndices
  - COOSliceRows(NDArray)
  - COOSort_
  - Also add aclrtSynchronizeDevice() to existing COOGetDataAndIndices/COOGetData
- segment_reduce.cc:
  - Use default stream (nullptr) instead of custom stream
  - Use synchronous aclrtMemcpy instead of aclrtMemcpyAsync
  - Add aclrtSynchronizeDevice() at entry to ensure PyTorch ops complete
- unit_graph.cc: InEdges() prefers COO format on Ascend to avoid CSR kernel
  stream sync issues (COO InEdges uses CPU fallback which is more reliable)
@xuejiakn
test: add SDDMM/SPMM Ascend NPU test suite
498d8b0 
- 51 test cases covering:
  - SDDMM: dot/add/mul/copy_lhs/copy_rhs × COO/CSR × FP32/FP16 × int32/int64
  - SPMM: copy_lhs + sum/max/min × int32/int64, update_all, backward
  - UDF reduce (degree bucketing) + backward
  - Stream sync: in_edges/in_degrees after update_all/gsddmm
  - Repeated update_all (multi-step stability)
  - segment_reduce (mean_nodes) forward + backward
  - End-to-end: GCN forward/backward, multi-layer GCN, training loop
- 4 skipped: sub/div CPU fallback (linker limitation)
- 3 xfailed: FP16 SDDMM (intermittent NPU stream race, passes in isolation)
@xuejiakn
fix(sddmm): fix add/sub/mul/div CPU fallback with NPU stream sync
14b4e38 
- sddmm.cc: use explicit aclrtSynchronizeDevice after each CopyTo in CPU fallback
  to ensure NPU data is fully transferred before CPU computation
- _sparse_ops.py: add torch.npu.synchronize() before _CAPI_DGLKernelSDDMM
  (same pattern as segment_reduce) to prevent PyTorch NPU stream data races
- test_sddmm_spmm.py: enable all SDDMM ops (add/sub/mul/div) with proper
  tolerance for div, add sync in test setup/teardown
- All 58 tests pass: 0 failed, 0 skipped, 0 xfailed
@xuejiakn CANNBot
feat(edge_softmax): add edge_softmax and CSRTranspose NPU support for…
ca600a8 
… Ascend

Add Ascend NPU adaptation for edge_softmax operator, enabling GAT/PAGTN
graph attention network training on Ascend 910B NPU.

New operators:
- edge_softmax forward/backward: Ascend C kernel with FullLoad/RowSplit
  dual-mode, FP32/FP16 support, AR/ARA (num_heads=1/>1) branches
- CSRTranspose: Ascend implementation via COOToSR(COOTranspose(CSRToCOO))
  chain, reusing existing Ascend kernels

Framework integration:
- kernel.cc: dispatch edge_softmax by data tensor device (not graph context)
  to handle NPU tensors with CPU-resident graph topology; add CPU fallback
  for BackwardSegmentCmp and ScatterAdd (needed by GAT WeightedSumAndMax)
- array.cc: add Ascend dispatch branch for CSRTranspose

Host code (edge_softmax.cc):
- ACL host with forward/backward template specializations (FP32/FP16 x
  int32/int64)
- CSC edge ID remapping via CPU gather/scatter (NPU torch index ops
  unreliable on DGL blob tensors)
- Backward interface adapted to DGL convention (sds = out * grad_out)

End-to-end validation:
- pytest test_gat.py: 2/3 passed, 1 regression needs 800 epochs (CPU also
  needs more epochs)
- pytest test_pagtn.py: no operator crashes (precision is test-epoch limit)
- test_edge_softmax_npu.py: 78 tests, 100% pass (FP32/FP16 x fwd/bwd x
  6 graph types x 3 num_heads + edge cases + perf)

Co-Authored-By: CANNBot <cannbot@ascend>
This was referenced Jul 28, 2026
Open
Open
Open
Open
Open
Open
Closed
@xuejiakn

xuejiakn commented Jul 28, 2026

Copy link
Copy Markdown
Author

PR 已由重组后的新 PR 替代

本 PR（10 commits）存在三个问题：① 巨型 commit 700239cc（65 文件）混杂 8 个功能模块；② 与 #13/#14 大面积重叠（stream sync / SDDMM / SPMM / COO fallback / 测试）；③ 含已合并到 master 的残留（#15[2] 编译修复、#15[3] nd2nz 修复）。

经核查，#15[1] 巨型 commit 的 65 个文件主体已在上游 master 合并，残留 6 文件均为旧版回退（README 回退为英文版、CMake/build.sh 含硬编码个人路径 /home/zqm1/、spmm_sum_kernel 是 nd2nz 修复的反向），无有价值新增，因此不再拆分。

重组后各 commit 去向：

原 commit 处置 新 PR
#15[1] 巨型commit (65文件) 关闭，主体已在 master —
#15[2] 编译修复 已在 master —
#15[3] nd2nz 修复 已在 master —
#15[4] SDDMM 实现 重组 #16 feat/sddmm-ascend
#15[5] SPMM 实现 重组 #17 feat/spmm-ascend
#15[6] stream sync 重组（与#13合并） #13 fix/npu-stream-sync
#15[7] COO fallback 重组 #18 fix/coo-cpu-fallback
#15[8] 测试 重组 #19 test/sddmm-spmm
#15[9] SDDMM fix 重组 #16 feat/sddmm-ascend
#15[10] edge_softmax 重组 #20 feat/edge-softmax-csr-transpose

关闭原因：内容已全部去重并重新组织到 #13/#16/#17/#18/#19/#20 中。

@xuejiakn xuejiakn closed this Jul 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@xuejiakn @lizengyong