Star844
1305
代码
代码
Issues1386
Pull Requests2250
流水线
项目讨论
Wiki
分析
Star844
1305
已合并
feat add shim_npu #35590
zzhongmin创建于 5月13日
feat add shim_npu #35590
已合并
zzhongmin创建于 5月13日
从zzhongmin/pytorchv2.9.0_shim合入到Ascend/pytorchv2.9.0
zzhongmin
zzhongmin成员
5月13日

【合入来源】

【修改方案】

将 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 流
  1. 改进已有函数的错误处理（shim_npu.cpp）
  • 所有输出参数增加空指针检查，预置 nullptr 防止异常时残留野值
  • create_tensor_from_blob_npu_v2 修正错误码传播：throw → TORCH_CHECK，检查内部调用返回值
  1. 新增流查找底层支撑（NPUStream.cpp）
  • 新增 getNPUStreamFromManagedAclrtStream，在 torch_npu 管理的流池中查找 aclrtStream 对应的 NPUStream 对象，等价于 CUDA 的 getStreamFromExternal
  1. 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 内部使用
  1. 修复流传递 Bug（model_container_runner_npu.cpp）
  • run_impl：尊重调用者传入的 stream_handle，仅在为 nullptr 时回退到当前流
  • run_with_npu_stream：修复忽略传入流参数的 Bug
  1. 修正 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]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
（感谢 zzhongmin 的贡献）
zzhongminzzhongmin成员
5月13日 创建了 pull request，commit 9b45a96b
ascend-robot
ascend-robot成员
5月13日 评论：
O
openLiBingCI成员
5月13日 评论：

本PR中共发现代码检查告警抑制3处，请Committer检视合理性

本评论自动扫描PR中使用的开源代码检查工具（ruff、clang-tidy、CodeQL等）的屏蔽注释。
这些屏蔽注释会阻止开源代码检查工具对特定代码区域的检测，可能导致潜在问题被忽略。

点击下载完整报告

文件路径 行号 屏蔽类型 代码片段 工具名称
test/_inductor/test_npu_aoti_shim.py 21 行级屏蔽 15: sys.path.remove(path)
16: if BUILD_PACKAGES_DIR.exists():
17: sys.path.insert(0, str(BUILD_PACKAGES_DIR))
19: import torch
21: from torch_npu.testing.testcase import run_tests, TestCase # noqa: F401
24: CPP_EXTENSIONS_DIR = REPO_ROOT / "test" / "cpp_extensions"
25: sys.path.insert(0, str(CPP_EXTENSIONS_DIR))
27: from torch_test_cpp_extension.load_npu_aoti_shim import load_npu_aoti_shim_extension
30: class TestNpuAOTIShim(TestCase):
31: module = None		 flake8
test/cpp_extensions/torch_test_cpp_extension/
load_npu_aoti_shim.py		 39 行级屏蔽 33: check=True,
34: env=env,
35: )
38: def load_npu_aoti_shim_extension():
39: import torch # noqa: F401
41: import torch_npu # noqa: F401
43: try:
44: return importlib.import_module(MODULE_NAME)
45: except ImportError:
46: _build_extension_inplace()		 flake8
test/cpp_extensions/torch_test_cpp_extension/
load_npu_aoti_shim.py		 41 行级屏蔽 34: env=env,
35: )
38: def load_npu_aoti_shim_extension():
39: import torch # noqa: F401
41: import torch_npu # noqa: F401
43: try:
44: return importlib.import_module(MODULE_NAME)
45: except ImportError:
46: _build_extension_inplace()
47: sys.modules.pop(MODULE_NAME, None)		 flake8
likedislike
ascend-robot
ascend-robot成员
5月13日 评论：

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 ✅ chujinjin, liujunzhu, hbhu_bin (3/3) ✅ chujinjin (1/1)
test ✅ chujinjin, liujunzhu, hbhu_bin (3/3) ✅ chujinjin (1/1)
torch_npu/_inductor ✅ chujinjin, liujunzhu, hbhu_bin (3/3) ✅ chujinjin (1/1)
torch_npu/csrc/inductor ✅ hbhu_bin, chujinjin, liujunzhu (3/3) ✅ chujinjin (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

zzhongmin, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robotascend-robot成员
5月13日 添加了label：ascend-cla/yes
此处折叠了158条消息 查看更多
ascend-robotascend-robot成员
5月18日 添加了label：approved
H
hbhu_bin成员
5月18日 评论：

/lgtm

likedislike
ascend-robotascend-robot成员
5月18日 添加了label：lgtm
ascend-robot
ascend-robot成员
5月18日 评论：

Review Guide

This pull-request passes review.
Committers who wrote a comment of /approve are: chujinjin.
Reviewers who wrote a comment of /lgtm are: liujunzhu, chujinjin, hbhu_bin.

likedislike
ascend-robotascend-robot成员
5月18日 合入了pull request