| [feat]support shmemput/get
Co-authored-by: pengqi<pengqi33@huawei.com>
# message auto-generated for no-merge-commit merge:
!36241 merge master_symm into master
[feat]support shmemput/get
Created-by: pengqihw
Commit-by: pengqihw;pengqi
Merged-by: ascend-robot
Description: <!--
PR描述模板更新日期:20260203
-->
# 【合入来源】
> <font color="red">**如有社区issue,请关联issue链接**</font>\
> <font color="red">**请勿携带内部流程信息(需求链接、问题单、内部issue等)**</font>
- [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 <mutex> , NPUSHMEMSymmetricMemoryAllocator 新增 std::mutex mutex_
- NPUSHMEMSymmetricMemory.cpp : alloc() 、 free() 、 get_alloc_size() 、 rendezvous() 四个方法中对 allocations_ 和 symm_mems_ 的访问加 std::lock_guard<std::mutex> 保护
# 【资料变更】
不涉及
# 【接口变更】
不涉及
PyTorch 上游已提供了对称内存的抽象框架,并定义了 nvshmem_put / nvshmem_get 算子接口。
# 【功能验证】

# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]
- [ ] 代码注释完备,正确记录错误日志
- [ ] 代码实现进行了返回值、空指针等校验
- [ ] PR标题正确使用类型标签,如:feat、fix、refactor、docs、test等
- [ ] PR持续集成流水线(CI)执行通过,代码检查无异常
See merge request: Ascend/pytorch!36241 | 12 天前 |