已合并
feat(inductor): Add dispatcher-free _empty_strided_npu fast path allocation #41836
feat(inductor): Add dispatcher-free _empty_strided_npu fast path allocation #41836
已合并
liuyutong创建于 7月16日
liuyutong
liuyutong成员
7月16日

【合入来源】

如有社区issue,请关联issue链接
请勿携带内部流程信息(需求链接、问题单、内部issue等)
https://gitcode.com/Ascend/pytorch/issues/2888

【修改方案】

Add a torch_npu._C._empty_strided_npu binding that mirrors upstream's
empty_strided (CUDA/XPU/MTIA) fast path for inductor-generated
wrappers, bypassing the operator dispatcher.

  • InitNpuBindings.cpp: THPModule_empty_strided_npu METH_VARARGS stub that
    parses the (sizes, strides, dtype) 3-tuple directly and calls the native
    factory; registered in TorchNpuMethods.
  • TensorFactories.{h,cpp}: TORCH_NPU_API empty_strided_npu — inlines the
    essential empty_strided steps (storage byte-size from size/stride, single
    allocate, set sizes/strides, one SetDesc) without empty({0})'s
    RECORD_FUNCTION / NPURecordFunction guard / 0-byte resize round-trip.

at::empty_strided(device='npu') is dispatched (~2us/alloc). Inductor backward
graphs allocate many buffers per step, so this host overhead dominates; the
fast path removes it while keeping the NPU storage-descriptor setup.

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

基础功能测试

测试用例 状态 验证功能
test/npu/test_tensor.py::TestTensor::test_empty_strided ✅ PASSED 基本empty_strided功能、各种shape/stride组合、空tensor处理、storage大小计算正确性

API兼容性测试

测试用例 状态 验证功能
test/npu/test_npu.py::TestNpu::test_function_torch_empty_strided ✅ PASSED torch.empty_strided API、不同dtype支持
test/test_tensor_creation_ops.py::TestTensorCreationPRIVATEUSE1::test_empty_strided_npu ✅ PASSED 通用tensor创建、跨设备一致性

新增测试用例验证

新增 test/npu/test_empty_strided_npu.py 测试文件,包含9个测试用例全面验证 _empty_strided_npu fast path 功能。

测试覆盖:

  • 基础功能、空tensor、多dtype支持 (float32/16, int32/8, uint8, bool)
  • 复杂stride模式、广播stride
  • 确定性计算场景torch.empty_strided一致性验证、Storage大小计算正确性
  • new_empty_strided API集成、@Dtypes装饰器测试

测试结果: 9/9 passed ✅

文件信息:

  • 路径: test/npu/test_empty_strided_npu.py
  • 行数: 200行
  • 运行: pytest test/npu/test_empty_strided_npu.py -v

该测试套件确保了inductor NPU内存分配fast path的功能正确性和兼容性。

【CheckList】

PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 liuyutong 的贡献)
liuyutongliuyutong成员
7月16日 创建了 pull request,commit 20105222
liuyutongliuyutong成员
7月16日 关联了issue:[Feature]: Add dispatcher-free `_empty_strided_npu` fast path for Inductor performance
atomgit-bot
atomgit-bot
7月16日 评论:

变更摘要

此 PR 为 NPU 后端新增了一条绕过算子分发器(dispatcher)的快速分配路径 torch_npu._C._empty_strided_npu,对标上游 CUDA/XPU/MTIA 的 _empty_strided_<device> 设计。该快速路径主要用于 inductor 生成的 wrapper 在分配 strided NPU 张量时直接调用原生工厂函数,避免了 at::empty_strided 分发路径中 RECORD_FUNCTIONNPURecordFunction profiler guard 以及 0 字节 storage 的 resize 往返开销,从而降低 inductor 反向图大量 buffer 分配时的 host 侧延迟。

主要改动

  • 新增 _empty_strided_npu Python 绑定: 在 InitNpuBindings.cpp 中通过 TorchNpuMethods 注册了 THPModule_empty_strided_npu,直接解析 (sizes, strides, dtype) 三元组并调用 at_npu::native::empty_strided_npu,完全绕过算子分发器。

  • 新增 empty_strided_npu 快速分配实现: 在 TensorFactories.cpp 中实现了 TORCH_NPU_API 导出的 empty_strided_npu 函数,内联了 strided NPU 分配的核心步骤(从 size/stride 计算 storage 字节大小、单次 allocate、设置 sizes/strides、一次 SetDesc),去除了 empty({0}) 引入的冗余开销。

  • 新增头文件导出声明: 在 TensorFactories.h 中声明了 empty_strided_npu,通过 TORCH_NPU_API 宏暴露给 torch_npu._C 调用,保证该符号在 libtorch_npu.so 中对外可见。

  • 新增完整测试覆盖: 新增 test/npu/test_empty_strided_npu.py,覆盖空张量、多种 dtype、复杂 stride 模式、广播 stride、storage 大小一致性、确定性模式一致性及重复性等场景,确保快速路径与正常路径行为一致。

likedislike
atomgit-bot
atomgit-bot
7月16日 评论:

代码审查

审查总结

本次审查针对 feat(inductor): Add dispatcher-free _empty_strided_npu fast path allocation 变更的 4 个文件进行了全面审查。

各文件审查结果

文件 结果
test/npu/test_empty_strided_npu.py 无问题
torch_npu/csrc/InitNpuBindings.cpp 发现 2 个问题 (P0, P1)
torch_npu/csrc/aten/common/TensorFactories.cpp 发现 2 个问题 (P2, P2)
torch_npu/csrc/aten/common/TensorFactories.h 无问题

发现问题统计

  • P0 (1): _npu_unwrap_size_tuple 声明为 static void 但包含 return false;,导致编译错误,阻塞所有使用。
  • P1 (1): THPModule_empty_strided_npu 调用 _npu_unwrap_size_tuple 后未检查返回值,解析失败后仍继续执行,可能导致数据损坏和 Python 异常丢失
  • P2 (2): empty_strided_npu 缺少 size/stride 长度一致性和 stride 非负性校验,在边界输入下可能产生越界访问内存分配错误

总体风险判定

风险较高。P0 编译错误将使该功能完全无法使用。P1 和 P2 问题在正确的使用场景(inductor 内部调用)下触发概率较低,但作为公开接口缺少防御性校验,存在健壮性隐患。建议优先修复 P0 编译错误和 P1 调用方返回值检查问题,然后补充 P2 的输入校验。

类型 数量
🔴 阻塞 4
🟡 建议 1

⛔ 需要修改

likedislike
此处折叠了85条消息 查看更多
ascend-robotascend-robot成员
28 天前 添加了label:lgtm
ascend-robotascend-robot成员
28 天前 删除了label:ci-pipeline-passed
ascend-robotascend-robot成员
28 天前 解决了最后一个问题
ascend-robotascend-robot成员
28 天前 合入了pull request
ascend-robot
ascend-robot成员
28 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#12976 [ commitID:1cba57de ] 已完成
likedislike