torch.add Ascend C SIMT 迁移说明
概述
本目录迁移 torch.add / aten::add.Tensor 的非 double、非 complex 路径到 Ascend C SIMT 平台,输出形态为 out-of-tree torch_npu 扩展。
用户指定的 pytorch/aten/src/ATen/native/cuda/PointwiseOpsKernel.cu 经源码检查不包含 torch.add 主实现。torch.add 的真实链路来自 native_functions.yaml 的 ufunc 配置、BinaryOps.cpp 中的 add_stub(device_type(), *this, alpha) 调用,以及 BinaryOps.h 中 codegen 声明的 DECLARE_DISPATCH(structured_binary_fn_alpha, add_stub)。
本迁移坚持一对一迁移原则,不把完整 TensorIterator 点算子降级成最终单一路径 kernel。当前 NPU 侧仅完成 add_op_alias::add.Tensor 的 equal-shape contiguous smoke 验证;原生 torch.add / aten::add 替换和完整 SIMT TensorIterator descriptor 仍标记为阻塞/待补齐。
目录结构
ported-ops/add/
├── add_op/
│ ├── __init__.py
│ ├── ops.py
│ └── csrc/
│ ├── add.asc
│ ├── ascend/
│ │ ├── cuda_loops_compat.h
│ │ ├── dispatch_stub_compat.h
│ │ ├── loops_compat.h
│ │ ├── memory_access_compat.h
│ │ └── tensor_iterator_adapter.h
│ └── simt/
│ └── add.asc
├── test/
│ ├── test_contract.py
│ └── test_extension.py
├── plan.md
├── pyproject.toml
├── requirements.txt
└── setup.py
已迁移内容
- 保留
TensorIteratorBase -> add_stub -> add_kernel_cuda(iter, alpha)的抽象形态。 - 保留
alpha参数语义:out = self + alpha * other。 - CPU counterpart 使用
TensorIterator验证 broadcast、non-contiguous 和 dtype 分派。 - NPU SIMT smoke 路径支持
float32equal-shape contiguous 输入。 - 提供
torch.ops.add_op_alias.add.Tensor作为 aten-like 别名,用于验证注册链路。 - double 和 complex 按用户要求不迁移,运行时显式报错。
未完成/阻塞项
- 默认不覆盖当前环境中的
aten::add.Tensor/torch.add注册,避免和已安装 torch_npu 发生重复注册冲突。 ADD_OP_REGISTER_ATEN=1仅保留为集成开关,尚未在可控 torch_npu 集成分支完成原生aten::add替换验证。- NPU 侧完整 TensorIterator descriptor、broadcast、stride、non-contiguous、out/inplace、多 dtype、多 kernel-selection 路径尚未补齐。
- half/bfloat16 需要继续补齐 Ascend 原生 dtype mapping。
- double、complex 明确不迁移。
编译
cd /home/q00954432/simt-buddy/ported-ops/add
source /usr/local/Ascend/cann-9.0.0/set_env.sh
python setup.py build_ext --inplace
本次验证结果:编译通过,生成 add_op/_C.cpython-313-x86_64-linux-gnu.so。编译期间仅出现 torch_npu 头文件 warning。
测试
cd /home/q00954432/simt-buddy/ported-ops/add
python -m unittest discover -s test
本次验证结果:
Ran 7 tests in 0.009s
OK
Ascend 950 PR Smoke
验证入口:
import torch
import torch_npu
import add_op
x = torch.tensor([1.0, 2.0, 3.0], device="npu")
y = torch.tensor([7.0, 8.0, 9.0], device="npu")
got = torch.ops.add_op_alias.add.Tensor(x, y, 2)
expected = torch.add(x.cpu(), y.cpu(), alpha=2)
torch.testing.assert_close(got.cpu(), expected)
本次实际输出:
npu_available= True
device= Ascend950PR_9599
got= tensor([15., 18., 21.])
expected= tensor([15., 18., 21.])
alias_npu_smoke=OK
当前最终状态
状态标签:blocked by environment or unsupported feature
原因:别名注册路径已在 Ascend 950 PR 上验证通过,但原生 torch.add / aten::add 一对一替换尚未在可控 torch_npu 集成环境完成;完整 SIMT TensorIterator 多路径也尚未补齐。按“不接受降级”的用户决策,本迁移不把当前 equal-shape contiguous smoke 路径声明为完整 torch.add 迁移完成。