Pull Request已成功合入, 合并人@ascend-robot
(感谢 liuyutong 的贡献)变更摘要
此 PR 为 NPU 后端新增了一条绕过算子分发器(dispatcher)的快速分配路径 torch_npu._C._empty_strided_npu,对标上游 CUDA/XPU/MTIA 的 _empty_strided_<device> 设计。该快速路径主要用于 inductor 生成的 wrapper 在分配 strided NPU 张量时直接调用原生工厂函数,避免了 at::empty_strided 分发路径中 RECORD_FUNCTION、NPURecordFunction profiler guard 以及 0 字节 storage 的 resize 往返开销,从而降低 inductor 反向图大量 buffer 分配时的 host 侧延迟。
主要改动
-
新增
_empty_strided_npuPython 绑定: 在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 大小一致性、确定性模式一致性及重复性等场景,确保快速路径与正常路径行为一致。


代码审查
审查总结
本次审查针对 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 |
⛔ 需要修改




【合入来源】
【修改方案】
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.
parses the (sizes, strides, dtype) 3-tuple directly and calls the native
factory; registered in TorchNpuMethods.
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_stridedAPI兼容性测试
test/npu/test_npu.py::TestNpu::test_function_torch_empty_stridedtorch.empty_stridedAPI、不同dtype支持test/test_tensor_creation_ops.py::TestTensorCreationPRIVATEUSE1::test_empty_strided_npu新增测试用例验证
新增
test/npu/test_empty_strided_npu.py测试文件,包含9个测试用例全面验证_empty_strided_npufast path 功能。测试覆盖:
测试结果: 9/9 passed ✅
文件信息:
test/npu/test_empty_strided_npu.pypytest test/npu/test_empty_strided_npu.py -v该测试套件确保了inductor NPU内存分配fast path的功能正确性和兼容性。
【CheckList】