已合并
[feat]support shmemput/get #36222
pengqihw创建于 5月20日
[feat]support shmemput/get #36222
已合并
pengqihw成员
5月20日
5月20日
# 【合入来源】
> **如有社区issue，请关联issue链接**\
> **请勿携带内部流程信息（需求链接、问题单、内部issue等）**
- [x] 需求
- [ ] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】
1.SHMEM put/get 功能支持（对齐上游 NVSHMEM 后端实现方式）
- NPUSHMEMInterface.h ：Shmem_putmem_on_stream / Shmem_getmem_on_stream 声明
- NPUSHMEMInterface.cpp ： aclshmemx_putmem_on_stream / aclshmemx_getmem_on_stream 动态加载
- NPUSHMEMExtension.cpp ： nvshmem_put / nvshmem_get 接口实现和注册，使用 on_stream 接口并传入 c10_npu::getCurrentNPUStream()
- shmem_host_def.h：同步shmem库的结构体更新
- test_shmem.py: get/put测试用例
2. 线程安全保护（同步上游 commit 4ae3a4e）
- NPUSHMEMSymmetricMemory.hpp ：添加 #include ， NPUSHMEMSymmetricMemoryAllocator 新增 std::mutex mutex_
- NPUSHMEMSymmetricMemory.cpp ： alloc() 、 free() 、 get_alloc_size() 、 rendezvous() 四个方法中对 allocations_ 和 symm_mems_ 的访问加 std::lock_guard 保护
# 【资料变更】
不涉及
# 【接口变更】
不涉及
PyTorch 上游已提供了对称内存的抽象框架，并定义了 nvshmem_put / nvshmem_get 算子接口。
# 【功能验证】
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]
- [ ] 代码注释完备，正确记录错误日志
- [ ] 代码实现进行了返回值、空指针等校验
- [ ] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
- [ ] PR持续集成流水线（CI）执行通过，代码检查无异常
Pull Request已成功合入, 合并人@ascend-robot
（感谢 pengqihw 的贡献）
ascend-robot
5月20日 评论：
5月20日 评论：
5月20日 添加了label：stat/needs-squash
ascend-robot
5月20日 评论：
5月20日 评论：
Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-Ascend/pytorch | ✅ liujunzhu, hbhu_bin (2/2) | ✅ liujunzhu (1/1) |
| test | ✅ zhenyu10, liujunzhu, hbhu_bin (3/2) | ✅ liujunzhu (1/1) |
| torch_npu/csrc/distributed | ✅ liujunzhu, zhenyu10, hbhu_bin (3/2) | ✅ liujunzhu (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
pengqihw, thanks for your pull request. All authors of the commits have signed the CLA. 👍
5月20日 添加了label：ascend-cla/yes
此处折叠了64条消息 查看更多
5月21日 添加了label：lgtm
ascend-robot
5月21日 评论：
5月21日 评论：
5月21日 关闭了关联的issue
5月21日 合入了pull request