| feat add shim_npu
Co-authored-by: zzhongmin<zhongmin23@huawei.com>
# message auto-generated for no-merge-commit merge:
!35592 merge master_shim into master
feat add shim_npu
Created-by: zzhongmin
Commit-by: zzhongmin
Merged-by: ascend-robot
Description: <!--
PR描述模板更新日期:20260203
-->
# 【合入来源】
- [x] 需求
- [ ] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】
将 NPU AOTI Shim 层对齐到上游 CUDA 实现(shim_cuda.cpp),补齐缺失的设备管理、流管理、内存分配等关键 shim 函数。
1. 新增 7 组 NPU Shim 函数(shim_npu.cpp + shim.h)
- 设备 Guard 三件套:create_npu_guard / delete_npu_guard / npu_guard_set_index,实现 AOTI 模型执行时的设备自动切换与恢复
- 流 Guard 两件套:create_npu_stream_guard / delete_npu_stream_guard,实现流状态的自动切换与恢复
- 获取当前流:get_current_npu_stream,查询指定设备的当前 NPU 流
- 缓存分配器两件套:npu_caching_allocator_raw_alloc / raw_delete,复用 PyTorch 显存池,替代低效的 aclrtMalloc/aclrtFree
2. 改进已有函数的错误处理(shim_npu.cpp)
- 所有输出参数增加空指针检查,预置 nullptr 防止异常时残留野值
- create_tensor_from_blob_npu_v2 修正错误码传播:throw → TORCH_CHECK,检查内部调用返回值
3. 新增流查找底层支撑(NPUStream.cpp)
- 新增 getNPUStreamFromManagedAclrtStream,在 torch_npu 管理的流池中查找 aclrtStream 对应的 NPUStream 对象,等价于 CUDA 的 getStreamFromExternal
4. AOTI 运行时改进(model.h + utils_npu.h)
- RAII_npuMalloc 新增 AOT_INDUCTOR_USE_CACHING_ALLOCATOR 分支,支持通过 shim 调用缓存分配器
- 新增 normalize_run_stream:当 stream 为 nullptr 时自动获取当前 NPU 流,应用于 run() / run_single_threaded() / run_const_fold()
- 新增 utils_npu.h:提供 AOTINpuGuard / AOTINpuStreamGuard RAII 封装类,供 model.so 内部使用
5. 修复流传递 Bug(model_container_runner_npu.cpp)
- run_impl:尊重调用者传入的 stream_handle,仅在为 nullptr 时回退到当前流
- run_with_npu_stream:修复忽略传入流参数的 Bug
6. 修正 Python 代码生成(npu_device.py + cpp_wrapper.py)
- cpp_aoti_device_guard:raise NotImplementedError → "AOTINpuGuard"
- cpp_aoti_stream_guard:"AOTICudaStreamGuard" → "AOTINpuStreamGuard"
- aoti_get_stream:"aoti_torch_get_current_cuda_stream" → "aoti_torch_get_current_npu_stream"
- cpp_wrapper.py:生成代码中 include utils_npu.h
# 【资料变更】
> 请确认是否涉及资料变更。如涉及,需要在PR中体现,并简要说明修改内容。如不涉及,需填写“不涉及”
# 【接口变更】
> 请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及,需要详细说明接口以及对应的变更内容,同时需要在资料中体现。如不涉及,需填写“不涉及”
# 【功能验证】
> 说明测试场景,测试方法。如果本次测试方式与常规单元测试不同,请详细说明您的测试步骤\
> 新增/变更内容是否已新增/适配UT测试用例看护,并补充测试自验证截图
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]
- [x] 代码注释完备,正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签,如:feat、fix、refactor、docs、test等
- [x] PR持续集成流水线(CI)执行通过,代码检查无异常
See merge request: Ascend/pytorch!35592 | 15 天前 |