feat: SDDMM/SPMM NPU-native kernels + stream sync fixes for Ascend - #14
Closed
xuejiakn wants to merge 7 commits into
Closed
feat: SDDMM/SPMM NPU-native kernels + stream sync fixes for Ascend#14xuejiakn wants to merge 7 commits into
xuejiakn wants to merge 7 commits into
Conversation
- 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 - Add sddmm.cc: host dispatch supporting all ops - Modify kernel.cc: SDDMM dispatch macros to ATEN_XPU_SWITCH_CUDA_ASCEND
- Add spmm_kernel.cpp: Ascend C kernel for SpMM (copy_lhs + sum/max/min) - Add spmm_tiling.h: tiling struct - Modify spmm.cc: FP32 uses NPU kernel, int64 cast, default stream
…ment_reduce - streams.py: support torch.npu.Stream - heterograph.py: _sync_npu() before/after structure queries - _sparse_ops.py: sync before segment_reduce - core.py: fix assert int(deg)
…stream - array.cc: CPU fallback for COOGetRowNNZ/COOGetRowDataAndIndices/COOSliceRows/COOSort_ - segment_reduce.cc: default stream, synchronous memcpy, aclrtSynchronizeDevice - unit_graph.cc: InEdges prefers COO on Ascend
xuejiakn
force-pushed
the
feat-sddmm-spmm-ascend
branch
from
July 21, 2026 13:42
70eabed to
28901c0
Compare
51 test cases covering SDDMM/SPMM ops, UDF reduce, stream sync, and end-to-end GCN.
51 test cases covering SDDMM/SPMM ops, UDF reduce, stream sync, and end-to-end GCN.
- sddmm.cc: explicit aclrtSynchronizeDevice after each CopyTo in CPU fallback - _sparse_ops.py: add torch.npu.synchronize() before SDDMM CAPI call - All 58 tests pass: 0 failed, 0 skipped, 0 xfailed
This was referenced Jul 22, 2026
Author
PR 已由重组后的新 PR 替代本 PR 的内容已按功能维度重新拆分并整合到以下新 PR 中，避免与 #13/#15 的重复：
关闭原因：本 PR 与 #13（stream sync）、#15（SDDMM/SPMM）存在大面积功能重叠，reviewer 难以判断以哪个为准。重组后每个功能对应一个独立 PR，便于逐项 review 和 cherry-pick。 如需查看重组详情，请见各新 PR 的 description。 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add NPU-native Ascend C kernels for SDDMM and SPMM operators, plus PyTorch NPU stream synchronization fixes to resolve data races between DGL kernels and PyTorch autograd.
Changes
1. SDDMM operator (4 new files)
2. SPMM operator (2 new files)
3. Stream sync fixes (PR #13 + segment_reduce)
4. COO operation Ascend support
Test Results