cann-robotReplace blockDim with numBlocks

文件	最后提交记录	最后更新时间
CMakeLists.txt	add feature sample	4 个月前
README.md	add feature sample	4 个月前
add_custom.asc	Replace blockDim with numBlocks	4 个月前
add_custom_test.py	add feature sample	4 个月前

torch.library注册自定义算子直调样例

概述

本样例展示了如何使用PyTorch的torch.library机制注册自定义算子，并通过<<<>>>内核调用符调用核函数，以简单的Add算子为例，实现两个向量的逐元素相加。

支持的产品

Ascend 950PR/Ascend 950DT
Atlas A3 训练系列产品/Atlas A3 推理系列产品
Atlas A2 训练系列产品/Atlas A2 推理系列产品

目录结构介绍

├── torch_library
│   ├── CMakeLists.txt          // 编译工程文件
│   ├── add_custom_test.py      // PyTorch调用脚本
│   └── add_custom.asc          // Ascend C算子实现 & 自torch.library注册

算子描述

算子功能：

Add算子实现了两个数据相加，返回相加结果的功能。对应的数学表达式为：
```
z = x + y
```
算子规格：

算子类型(OpType) AddCustom

算子输入 name shape data type format

x 8 * 2048 float16 ND

y 8 * 2048 float16 ND

算子输出 z 8 * 2048 float16 ND

核函数名 add_custom
算子实现：

Ascend C提供的矢量计算接口Add的操作元素都为LocalTensor，输入数据需要先搬运进片上存储，然后使用计算接口完成两个输入参数相加，得到最终结果，再搬出到外部存储上。

Add算子的实现流程分为3个基本任务：CopyIn，Compute，CopyOut。CopyIn任务负责将Global Memory上的输入Tensor xGm和yGm搬运到Local Memory，分别存储在xLocal、yLocal，Compute任务负责对xLocal、yLocal执行加法操作，计算结果存储在zLocal中，CopyOut任务负责将输出数据从zLocal搬运至Global Memory上的输出Tensor zGm中。
自定义算子注册：

本样例在add_custom.asc中定义了一个名为ascendc_ops的命名空间，并在其中注册了ascendc_add函数。

PyTorch提供TORCH_LIBRARY宏作为自定义算子注册的核心接口，用于创建并初始化自定义算子库，注册后在Python侧可以通过torch.ops.namespace.op_name方式进行调用，例如：
```
TORCH_LIBRARY(ascendc_ops, m) {
    m.def(ascendc_add"(Tensor x, Tensor y) -> Tensor");
}
```
TORCH_LIBRARY_IMPL用于将算子逻辑绑定到特定的DispatchKey（PyTorch设备调度标识）。针对NPU设备，需要将算子实现注册到PrivateUse1这一专属的DispatchKey上，例如：
```
TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
{
    m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
}
```
在ascendc_add函数中通过c10_npu::getCurrentNPUStream()函数获取当前NPU上的流，并通过内核调用符<<<>>>调用自定义的Kernel函数add_custom，在NPU上执行算子。
Python测试脚本

在add_custom_test.py调用脚本中，通过torch.ops.load_library加载生成的自定义算子库，调用注册的ascendc_add函数，并通过对比NPU输出与CPU标准加法结果来验证自定义算子的数值正确性。

编译运行

安装PyTorch以及Ascend Extension for PyTorch插件

请参考pytorch: Ascend Extension for PyTorch开源代码仓或Ascend Extension for PyTorch昇腾社区的安装说明，选取支持的Python版本配套发行版，完成torch和torch-npu的安装。
安装前置依赖
```
pip3 install expecttest
```
配置环境变量

请根据当前环境上CANN开发套件包的安装方式，选择对应配置环境变量的命令。
- 默认路径，root用户安装CANN软件包
```
source /usr/local/Ascend/cann/set_env.sh
```
- 默认路径，非root用户安装CANN软件包
```
source $HOME/Ascend/cann/set_env.sh
```
- 指定路径install_path，安装CANN软件包
```
source ${install_path}/cann/set_env.sh
```
样例执行

在本样例根目录下执行如下步骤，运行该样例。
```
mkdir -p build; cd build
cmake ..; make -j
python3 ../add_custom_test.py
```
执行结果如下，说明精度对比成功。
```
Ran 1 test in **s.
OK
```

算子类型(OpType)	AddCustom
算子输入	name	shape	data type	format
	x	8 * 2048	float16	ND
	y	8 * 2048	float16	ND
算子输出	z	8 * 2048	float16	ND
核函数名	add_custom