README.md

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 路径支持 float32 equal-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 迁移完成。