Star181
848
代码
代码
Issues149
Pull Requests329
流水线
项目讨论
Wiki
分析
Star181
848
已合并
添加aclgraph+<<<>>> 示例代码 #4019
mihudan创建于 1月9日
添加aclgraph+<<<>>> 示例代码 #4019
已合并
mihudan创建于 1月9日
从mihudan/op-pluginmaster合入到Ascend/op-pluginmaster
共 16 个文件变更+1420-0
Aexamples/kernel_extension_aclgraph/pybind/README.md+170-0
@@ -0,0 +1,170 @@
1+# 自定义算子直调并适配aclgraph
2+ 
3+## 概述
4+ 
5+本样例展示了如何使用Pybind注册自定义算子，通过<<<>>>内核调用符调用核函数，并适配aclgraph使用该自定义算子，以简单的Add算子和三角函数计算的原地算子为例，实现aclgraph下自定义算子的调用。
6+ 
7+## 支持的产品
8+ 
9+- Atlas A3 训练系列产品/Atlas A3 推理系列产品
10+- Atlas A2 训练系列产品/Atlas A2 推理系列产品
11+ 
12+## 目录结构介绍
13+ 
14+```
15+├── README.md // 示例介绍
16+├── setup.py // setup文件
17+├── csrc
18+│ ├── add_custom.asc // Add算子实现 & 自定义算子注册
19+│ └── trig_inplace_custom.asc // 原地三角函数算子实现 & 自定义算子注册
20+├── op_extension
21+│ ├── __init__.py // python初始化文件
22+└── test
23+ ├── add_aclgraph_test.py // Add算子aclgraph测试demo
24+ └── trig_aclgraph_test.py // 原地三角函数aclgraph测试demo
25+```
26+ 
27+## 算子描述
28+### Add算子
29+- 算子功能：
30+ Add算子实现了两个数据相加，返回相加结果的功能。对应的算子原型为：
31+
32+ ```
33+ ascendc_add(Tensor x, Tensor y) -> Tensor
34+ ```
35+- 算子规格：
36+
37+ <table>
38+ <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">add_custom</td></tr>
39+ </tr>
40+ <tr><td rowspan="3" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr>
41+ <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
Y
Yyang-binrong1月17日

需要注释说明2048的含义，增加可读性

likedislike
42+ <tr><td align="center">y</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
43+ </tr>
44+ </tr>
45+ <tr><td rowspan="1" align="center">算子输出</td><td align="center">z</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
46+ </tr>
47+
48+ </table>
49+ 
50+### 原地三角函数算子
51+- 算子功能：
52+ 该算子入参为x, out_sin ,out_cos, 算子调用后，out_sin会被原地修改为sin(x)计算结果，out_cos会被原地修改为cos(x)计算结果，返回值tan(x)计算结果。对应的算子原型为：
53+
54+ ```
55+ ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor
56+ ```
57+- 算子规格：
58+ 
59+ <table>
60+ <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">trig_inplace_custom</td></tr>
61+ </tr>
62+ <tr><td rowspan="4" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr>
63+ <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
64+ <tr><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
65+ <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
66+
67+ </tr>
68+ </tr>
69+ <tr><td rowspan="3" align="center">算子输出</td><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
70+ <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
71+ <tr><td align="center">out_tan</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
72+ </tr>
73+
74+ </table>
75+ 
76+## 代码实现介绍
77+ 
78+- 以Add算子为例，样例在*.asc文件中定义了一个名为ascendc_ops的命名空间，并在其中注册了ascendc_add函数。在ascendc_add函数中通过`c10_npu::getCurrentNPUStream()`函数获取当前NPU上的流，并通过内核调用符<<<>>>调用自定义的Kernel函数add_custom，在NPU上执行算子。
79+ ```c++
80+ add_custom<<<blockDim, nullptr, aclStream>>>(xGm, yGm, zGm, totalLength);
81+ ```
82+
83+- 在pybind11.asc文件中使用了pybind11库来将C++代码封装成Python模块，在Python侧可以通过`import`方式进行调用。例如：
84+
85+ ```c++
86+ PYBIND11_MODULE(custom_ops, m)
87+ {
88+ m.def("run_add_custom", &ascendc_ops::run_add_custom, "");
89+ m.def("run_trig_custom", &ascendc_ops::run_trig_custom, "");
90+ }
91+ ```
92+ 
93+- python侧通过`torch.library`将算子逻辑绑定到特定的DispatchKey（PyTorch设备调度标识）。针对NPU设备，需要将算子实现注册到PrivateUse1这一专属的DispatchKey上，例如：
94+
95+ ```python
96+ ascendc_ops = library.Library("ascendc_ops", "DEF")
97+ 
98+ ascendc_ops.define("ascendc_add(Tensor a, Tensor b) -> Tensor")
99+ 
100+ @library.impl(ascendc_ops, "ascendc_add", "PrivateUse1")
101+ def add_custom_ops(a, b):
102+ return custom_ops.run_add_custom(a, b)
103+ ```
104+ 
105+- 注册Meta函数：
106+ 注册Meta函数使faketensor流程正常工作，在使用fx, compile等功能涉及，本示例在add_aclgraph_test.py开头注册代码如下：
107+ 
108+ ```python
109+ @library.impl(ascendc_ops, "ascendc_add", "Meta")
110+ def ascendc_add_meta(a, b):
111+ return torch.empty_like(a)
112+ ```
113+ 
114+- aclgraph的调用：
115+ [示例代码](./test/add_aclgraph_test.py)中，展示了3种aclgraph的使能方式，通过对比NPU输出与CPU标准加法结果来验证自定义算子的数值正确性。
116+ 
117+1. torch.npu.NPUGraph()
118+2. torch.npu.make_graphed_callables
119+3. backend="npugraph_ex"
120+ 
121+## 编译运行
122+ 
123+在本样例根目录下执行如下步骤，编译并执行算子。
124+ 
125+- 环境安装
126+
127+1. 请参考与您当前使用的版本配套的[《Ascend Extension for PyTorch
128+ 软件安装指南》](https://www.hiascend.com/document/detail/zh/Pytorch/720/configandinstg/instg/insg_0001.html)，获取PyTorch和torch_npu详细的安装步骤。
129+
130+ 本样例需torch2.6.0版本以上，支持`backend="npugraph_ex"`需7.3.0版本以上。
131+2. 根据实际环境安装CANN toolkit包，安装指导详见《[CANN 软件安装指南](https://www.hiascend.com/document/redirect/CannCommunityInstSoftware)》。
132+3. 根据实际环境安装CANN ops包。根据产品型号和环境架构，下载对应安装包，可参考[下载链接](https://ascend.devcloud.huaweicloud.com/cann/run/software/8.5.0-beta.1)并执行如下命令安装：
133+
134+ ```bash
135+ # 确保安装包具有可执行权限
136+ chmod +x Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run
137+ # 安装命令
138+ ./Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run --install --quiet --install-path=${install_path}
139+ ```
140+
141+ - \$\{soc\_name\}：表示NPU型号名称，即\$\{soc\_version\}删除“ascend”后剩余的内容。
142+ - \$\{install\_path\}：表示指定安装路径，需要与toolkit包安装在相同路径，默认安装在`/usr/local/Ascend`目录。
143+ 
144+- 配置环境变量
145+
146+ 请根据当前环境上CANN开发套件包的安装位置，执行如下配置环境变量的命令。
147+
148+ ```bash
149+ source ${install_path}/ascend-toolkit/set_env.sh
150+ ```
151+ 
152+ 
153+ 
154+- 样例执行
155+
156+ ```bash
157+ python setup.py bdist_wheel
158+ pip install dist/*.whl --force-reinstall
159+ cd test
160+ python ./add_aclgraph_test.py
161+ ```
162+ 
163+执行结果如下，说明精度对比成功。
164+ 
165+```bash
166+Ran * test in **s.
167+OK
168+```
169+ 
170+ 
Aexamples/kernel_extension_aclgraph/pybind/csrc/add_custom.asc+109-0
@@ -0,0 +1,109 @@
1+// Copyright (c) 2026 Huawei Technologies Co., Ltd
Y
Yyang-binrong1月17日

文件头错误

likedislike
2+// All rights reserved.
3+//
4+// Licensed under the BSD 3-Clause License (the "License");
5+// you may not use this file except in compliance with the License.
6+// You may obtain a copy of the License at
7+//
8+// Unless required by applicable law or agreed to in writing, software
9+// distributed under the License is distributed on an "AS IS" BASIS,
10+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+// See the License for the specific language governing permissions and
12+// limitations under the License.
13+#include <pybind11/pybind11.h>
14+#include <torch/extension.h>
15+ 
16+#include "torch_npu/csrc/core/npu/NPUStream.h"
17+#include "kernel_operator.h"
18+ 
19+constexpr uint32_t BUFFER_NUM = 2; //tensor num for each queue
20+class KernelAdd {
21+public:
22+ __aicore__ inline KernelAdd() {}
23+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
24+ {
25+ this->blockLength = totalLength / AscendC::GetBlockNum();
Y
Yyang-binrong1月17日

AscendC::GetBlockNum()未判断是否为0，可能导致除0

likedislike
26+ this->tileNum = 8;
Y
Yyang-binrong1月17日

需要说明含义，增加可读性

likedislike
27+ this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
28+ xGm.SetGlobalBuffer((__gm__ int32_t *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
29+ yGm.SetGlobalBuffer((__gm__ int32_t *)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
30+ zGm.SetGlobalBuffer((__gm__ int32_t *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
31+ pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int32_t));
32+ pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(int32_t));
33+ pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(int32_t));
34+ }
35+ __aicore__ inline void Process()
36+ {
37+ int32_t loopCount = this->tileNum * BUFFER_NUM;
38+ for (int32_t i = 0; i < loopCount; i++) {
39+ CopyIn(i);
40+ Compute(i);
41+ CopyOut(i);
42+ }
43+ }
44+ 
45+private:
46+ __aicore__ inline void CopyIn(int32_t progress)
47+ {
48+ AscendC::LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
49+ AscendC::LocalTensor<int32_t> yLocal = inQueueY.AllocTensor<int32_t>();
50+ AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
51+ AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
52+ inQueueX.EnQue(xLocal);
53+ inQueueY.EnQue(yLocal);
54+ }
55+ __aicore__ inline void Compute(int32_t progress)
56+ {
57+ AscendC::LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
58+ AscendC::LocalTensor<int32_t> yLocal = inQueueY.DeQue<int32_t>();
59+ AscendC::LocalTensor<int32_t> zLocal = outQueueZ.AllocTensor<int32_t>();
60+ AscendC::Add(zLocal, xLocal, yLocal, this->tileLength);
61+ outQueueZ.EnQue<int32_t>(zLocal);
62+ inQueueX.FreeTensor(xLocal);
63+ inQueueY.FreeTensor(yLocal);
64+ }
65+ __aicore__ inline void CopyOut(int32_t progress)
66+ {
67+ AscendC::LocalTensor<int32_t> zLocal = outQueueZ.DeQue<int32_t>();
68+ AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
69+ outQueueZ.FreeTensor(zLocal);
70+ }
71+ 
72+private:
73+ AscendC::TPipe pipe;
74+ AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
75+ AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
76+ AscendC::GlobalTensor<int32_t> xGm;
77+ AscendC::GlobalTensor<int32_t> yGm;
78+ AscendC::GlobalTensor<int32_t> zGm;
79+ uint32_t blockLength;
80+ uint32_t tileNum;
81+ uint32_t tileLength;
82+};
83+ 
84+__global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
85+{
86+ KernelAdd op;
87+ op.Init(x, y, z, totalLength);
88+ op.Process();
89+}
90+ 
91+namespace ascendc_ops {
92+at::Tensor run_add_custom(const at::Tensor &x, const at::Tensor &y)
93+{
94+ auto acl_stream = c10_npu::getCurrentNPUStream().stream(true);
Y
Yyang-binrong1月17日

acl_stream 与getCurrentNPUStream等风格不一致，需要统一为小驼峰或者其他规格

likedislike
95+ at::Tensor z = at::empty_like(x);
96+ uint32_t blockDim = 8;
Y
Yyang-binrong1月17日

需要以dtype相关sizeof计算替换数字,防止功能拓展时散弹式修改

likedislike
97+ uint32_t totalLength = 1;
98+ for (uint32_t size : x.sizes()) {
99+ totalLength *= size;
100+ }
101+ // Launch the custom kernel use <<<>>>
102+ auto xGm = static_cast<uint8_t *>(const_cast<void *>(x.storage().data()));
103+ auto yGm = static_cast<uint8_t *>(const_cast<void *>(y.storage().data()));
104+ auto zGm = static_cast<uint8_t *>(const_cast<void *>(z.storage().data()));
105+ add_custom<<<blockDim, nullptr, acl_stream>>>(xGm, yGm, zGm, totalLength);
106+ return z;
107+}
108+ 
109+} // namespace ascendc_ops
Aexamples/kernel_extension_aclgraph/pybind/csrc/pybind11.asc+14-0
@@ -0,0 +1,14 @@
1+#include <pybind11/pybind11.h>
2+#include <torch/extension.h>
3+ 
4+namespace ascendc_ops {
5+at::Tensor run_add_custom(const at::Tensor &x, const at::Tensor &y);
Y
Yyang-binrong1月17日

缺少空格

likedislike
6+at::Tensor run_trig_custom(const at::Tensor &x, const at::Tensor &out_sin, const at::Tensor &out_cos);
7+}
8+ 
9+// expose Ascend custom ops to Python
10+PYBIND11_MODULE(custom_ops_lib, m)
11+{
12+ m.def("run_add_custom", &ascendc_ops::run_add_custom, "");
13+ m.def("run_trig_custom", &ascendc_ops::run_trig_custom, "");
14+}
Aexamples/kernel_extension_aclgraph/pybind/csrc/trig_inplace_custom.asc+126-0
@@ -0,0 +1,126 @@
1+// Copyright (c) 2026 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+//
4+// Licensed under the BSD 3-Clause License (the "License");
Y
Yyang-binrong1月17日

文件头错误，需要改为正确内容

likedislike
5+// you may not use this file except in compliance with the License.
6+// You may obtain a copy of the License at
7+//
8+// Unless required by applicable law or agreed to in writing, software
9+// distributed under the License is distributed on an "AS IS" BASIS,
10+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+// See the License for the specific language governing permissions and
12+// limitations under the License.
13+#include <pybind11/pybind11.h>
14+#include <torch/extension.h>
15+ 
16+#include "torch_npu/csrc/core/npu/NPUStream.h"
17+#include "kernel_operator.h"
18+ 
19+constexpr uint32_t BUFFER_NUM = 2; //tensor num for each queue
20+class KernelTrig {
21+public:
22+ __aicore__ inline KernelTrig() {}
23+ 
24+ // Initialize the global memory and buffer queues
25+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan, uint32_t totalLength)
26+ {
27+ this->blockLength = totalLength / AscendC::GetBlockNum();
Y
Yyang-binrong1月17日

【致命】AscendC::GetBlockNum()未校验是否为0,可能导致除0

likedislike
28+ this->tileNum = 8;
29+ this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
Y
Yyang-binrong1月17日

【致命】this->tileNum未校验是否为0,可能导致除0

likedislike
30+ xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
31+ sinGm.SetGlobalBuffer((__gm__ float *)out_sin + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
32+ cosGm.SetGlobalBuffer((__gm__ float *)out_cos + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
33+ tanGm.SetGlobalBuffer((__gm__ float *)out_tan + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
34+ pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
35+ pipe.InitBuffer(outQueueSin, BUFFER_NUM, this->tileLength * sizeof(float));
36+ pipe.InitBuffer(outQueueCos, BUFFER_NUM, this->tileLength * sizeof(float));
37+ pipe.InitBuffer(outQueueTan, BUFFER_NUM, this->tileLength * sizeof(float));
38+ }
39+ 
40+ __aicore__ inline void Process()
41+ {
42+ int32_t loopCount = this->tileNum * BUFFER_NUM;
43+ for (int32_t i = 0; i < loopCount; i++) {
44+ CopyIn(i);
45+ Compute(i);
46+ CopyOut(i);
47+ }
48+ }
49+ 
50+private:
51+ __aicore__ inline void CopyIn(int32_t progress)
52+ {
53+ AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
54+ AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
55+ inQueueX.EnQue(xLocal);
56+ }
57+ __aicore__ inline void Compute(int32_t progress)
58+ {
59+ AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
60+ AscendC::LocalTensor<float> sinLocal = outQueueSin.AllocTensor<float>();
61+ AscendC::LocalTensor<float> cosLocal = outQueueCos.AllocTensor<float>();
62+ AscendC::LocalTensor<float> tanLocal = outQueueTan.AllocTensor<float>();
63+ 
64+ AscendC::Sin(sinLocal, xLocal, this->tileLength);
65+ AscendC::Cos(cosLocal, xLocal, this->tileLength);
66+ AscendC::Tan(tanLocal, xLocal, this->tileLength);
67+ 
68+ outQueueSin.EnQue<float>(sinLocal);
69+ outQueueCos.EnQue<float>(cosLocal);
70+ outQueueTan.EnQue<float>(tanLocal);
71+ inQueueX.FreeTensor(xLocal);
72+ }
73+ __aicore__ inline void CopyOut(int32_t progress)
74+ {
75+ // Copy the sin, cos, and tan values from local memory to global memory (inplace modification)
76+ AscendC::LocalTensor<float> sinLocal = outQueueSin.DeQue<float>();
77+ AscendC::LocalTensor<float> cosLocal = outQueueCos.DeQue<float>();
78+ AscendC::LocalTensor<float> tanLocal = outQueueTan.DeQue<float>();
79+ AscendC::DataCopy(sinGm[progress * this->tileLength], sinLocal, this->tileLength);
80+ AscendC::DataCopy(cosGm[progress * this->tileLength], cosLocal, this->tileLength);
81+ AscendC::DataCopy(tanGm[progress * this->tileLength], tanLocal, this->tileLength);
82+ outQueueSin.FreeTensor(sinLocal);
83+ outQueueCos.FreeTensor(cosLocal);
84+ outQueueTan.FreeTensor(tanLocal);
85+ }
86+ 
87+private:
88+ AscendC::TPipe pipe;
89+ AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
90+ AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueSin, outQueueCos, outQueueTan;
91+ AscendC::GlobalTensor<float> xGm;
92+ AscendC::GlobalTensor<float> sinGm;
93+ AscendC::GlobalTensor<float> cosGm;
94+ AscendC::GlobalTensor<float> tanGm;
95+ uint32_t blockLength;
96+ uint32_t tileNum;
97+ uint32_t tileLength;
98+};
99+ 
100+__global__ __aicore__ void trig_inplace_custom(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan,
101+ uint32_t totalLength)
102+{
103+ KernelTrig op;
104+ op.Init(x, out_sin, out_cos, out_tan, totalLength);
105+ op.Process();
106+}
107+ 
108+namespace ascendc_ops {
109+at::Tensor run_trig_custom(const at::Tensor &x, const at::Tensor &out_sin, const at::Tensor &out_cos)
110+{
111+ auto acl_stream = c10_npu::getCurrentNPUStream().stream(true);
112+ at::Tensor out_tan = at::empty_like(x);
113+ uint32_t blockDim = 8;
Y
Yyang-binrong1月17日

数字8需要改为宏或者全局变量或加注释，防止功能拓展时散弹式修改

likedislike
114+ uint32_t totalLength = 1;
115+ for (uint32_t size : x.sizes()) {
116+ totalLength *= size;
117+ }
118+ auto xGm = static_cast<uint8_t *>(const_cast<void *>(x.storage().data()));
119+ auto sinGm = static_cast<uint8_t *>(const_cast<void *>(out_sin.storage().data()));
120+ auto cosGm = static_cast<uint8_t *>(const_cast<void *>(out_cos.storage().data()));
121+ auto tanGm = static_cast<uint8_t *>(const_cast<void *>(out_tan.storage().data()));
122+ // Launch the custom kernel using <<<>>>
123+ trig_inplace_custom<<<blockDim, nullptr, acl_stream>>>(xGm, sinGm, cosGm, tanGm, totalLength);
124+ return out_tan;
125+}
126+} // namespace ascendc_ops
Aexamples/kernel_extension_aclgraph/pybind/op_extension/__init__.py+3-0
@@ -0,0 +1,3 @@
1+from op_extension.custom_ops_lib import run_add_custom, run_trig_custom
2+ 
3+__all__ = ["run_add_custom", "run_trig_custom"]
Aexamples/kernel_extension_aclgraph/pybind/setup.py+97-0
@@ -0,0 +1,97 @@
1+import os
Y
Yyang-binrong1月17日

缺少文件头

likedislike
2+import glob
3+import sysconfig
4+from distutils.errors import CompileError
5+from distutils.spawn import find_executable
6+import torch
7+import torch_npu
8+import torch.utils.cpp_extension as cpp_extension
9+from setuptools import setup, Extension, find_packages
10+from setuptools.command.build_ext import build_ext
11+ 
12+BASE_DIR = os.path.dirname(os.path.realpath(__file__))
13+source_files = glob.glob(os.path.join(BASE_DIR, "csrc", "*.asc"), recursive=True)
14+ 
15+ 
16+def get_dependency_paths():
17+ python_include = sysconfig.get_config_var("INCLUDEPY")
18+ python_lib = sysconfig.get_config_var("LIBDIR")
19+ 
20+ torch_include_paths = cpp_extension.include_paths()
21+ torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib")
22+ 
23+ torch_npu_path = os.path.dirname(torch_npu.__file__)
24+ torch_npu_include = os.path.join(torch_npu_path, "include")
25+ torch_npu_lib = os.path.join(torch_npu_path, "lib")
26+ 
27+ all_include_paths = list([
28+ *torch_include_paths,
29+ python_include,
30+ torch_npu_include,
31+ ])
32+ 
33+ all_libs = list([
34+ python_lib,
35+ torch_lib,
36+ torch_npu_lib,
37+ ])
38+ 
39+ return {
40+ "all_includes": all_include_paths,
41+ "all_libs": all_libs
42+ }
43+ 
44+ 
45+class AscendBuildExtension(build_ext):
46+ def __init__(self, *args, **kwargs):
47+ super().__init__(*args, **kwargs)
48+ 
49+ def _check_bisheng_compiler(self):
50+ bisheng_compiler = find_executable('bisheng')
51+ if not bisheng_compiler:
52+ raise RuntimeError("bisheng command not found!")
53+ 
54+ def build_extension(self, ext):
55+ self._check_bisheng_compiler()
56+ dep_paths = get_dependency_paths()
57+ 
58+ ext_fullpath = self.get_ext_fullpath(ext.name)
59+ os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True)
60+ 
61+ compile_cmd = [
62+ "bisheng",
63+ "-x", "asc",
64+ "--npu-arch=dav-2201",
65+ "-shared",
66+ "-fPIC",
67+ "-std=c++17",
68+ "-ltorch_npu", "-ltorch", "-lc10",
69+ *ext.sources,
70+ "-o", ext_fullpath,
71+ ]
72+ 
73+ for include_dir in dep_paths["all_includes"]:
74+ compile_cmd.append(f"-I{include_dir}")
75+ 
76+ for lib_dir in dep_paths["all_libs"]:
77+ compile_cmd.append(f"-L{lib_dir}")
78+ 
79+ try:
80+ self.spawn(compile_cmd)
81+ except Exception as e:
82+ raise CompileError(f"{str(e)}") from e
83+ 
84+ 
85+your_ext = Extension(
86+ name="op_extension.custom_ops_lib",
87+ sources=source_files,
88+ language="asc",
89+)
90+ 
91+setup(
92+ name="op_extension",
93+ version="0.1",
94+ ext_modules=[your_ext],
95+ packages=find_packages(),
96+ cmdclass={"build_ext": AscendBuildExtension},
97+)
Aexamples/kernel_extension_aclgraph/pybind/test/add_aclgraph_test.py+92-0
@@ -0,0 +1,92 @@
1+import sys
2+import os
3+import torch
4+import torch_npu
5+import torch.library as library
6+from torch_npu.testing.testcase import TestCase, run_tests
7+import op_extension
8+ 
9+# Define Ascend custom operator library
10+ascendc_ops = library.Library("ascendc_ops", "DEF") # "DEF" means defining new operators
11+ 
12+# Define a new operator
13+ascendc_ops.define("ascendc_add(Tensor a, Tensor b) -> Tensor")
14+ 
15+ 
16+# Register a meta function
17+@library.impl(ascendc_ops, "ascendc_add", "Meta")
18+def ascendc_add_meta(a, b):
19+ return torch.empty_like(a)
20+ 
21+ 
22+# Register implementation for the "PrivateUse1" backend
23+@library.impl(ascendc_ops, "ascendc_add", "PrivateUse1")
24+def add_custom_ops(a, b):
25+ return op_extension.run_add_custom(a, b)
26+ 
27+ 
28+# Define a simple model using the custom operation
29+class Model(torch.nn.Module):
30+ def forward(self, x, y):
31+ return torch.ops.ascendc_ops.ascendc_add(x, y)
32+ 
33+ 
34+length = [8, 2048]
Y
Yyang-binrong1月17日

需要说明length属于哪个轴，并改用对应宏，避免散弹式修改

likedislike
35+ 
36+ 
37+class TestCustomAdd(TestCase):
38+ 
39+ def get_rand_input(self):
40+ x = torch.randint(low=1, high=100, size=length, device='npu', dtype=torch.int)
Y
Yyang-binrong1月17日

high=100格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
41+ y = torch.randint(low=1, high=100, size=length, device='npu', dtype=torch.int)
Y
Yyang-binrong1月17日

low=1格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
42+ return x, y
43+ 
44+ # Test using torch.npu.NPUGraph
45+ def test_npugraph(self):
46+ static_x, static_y = self.get_rand_input()
47+ static_target = torch.randint(low=1, high=100, size=length, device='npu:0', dtype=torch.int)
Y
Yyang-binrong1月17日

high=100格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
48+ 
49+ g = torch.npu.NPUGraph()
50+ model = Model()
51+ with torch.npu.graph(g):
52+ static_target = model(static_x, static_y)
53+ 
54+ real_x, real_y = self.get_rand_input()
55+ static_x.copy_(real_x)
56+ static_y.copy_(real_y)
57+ # replay
58+ g.replay()
59+ cpuout = torch.add(real_x, real_y)
60+ self.assertEqual(static_target, cpuout)
61+ 
62+ # Test using make_graphed_callables
63+ def test_make_graphed_callables(self):
64+ model = Model().npu()
65+ x, y = self.get_rand_input()
66+ model = torch.npu.make_graphed_callables(model, (x, y))
67+ 
68+ real_x = torch.randint_like(x, low=1, high=100)
Y
Yyang-binrong1月17日

high=100格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
69+ real_y = torch.randint_like(y, low=1, high=100)
Y
Yyang-binrong1月17日

格式有问题，缺少空格，且需要改用宏

likedislike
70+ output = model(real_x, real_y)
71+ cpuout = torch.add(real_x, real_y)
72+ self.assertEqual(output, cpuout)
73+ 
74+ # Test using the npugraph_ex backend for model compilation
75+ def test_npugraph_ex_backend(self):
76+ model = Model().npu()
77+ compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
Y
Yyang-binrong1月17日

fullgraph=True需要修正格式

likedislike
78+ x, y = self.get_rand_input()
79+ output = compiled_model(x, y)
80+ cpuout = torch.add(x, y)
81+ self.assertEqual(output, cpuout)
82+ 
83+ # Test single custom operator call
84+ def test_add_custom_ops(self):
85+ x, y = self.get_rand_input()
86+ output = torch.ops.ascendc_ops.ascendc_add(x.npu(), y.npu()).cpu()
87+ cpuout = torch.add(x, y)
88+ self.assertEqual(output, cpuout)
Y
Yyang-binrong1月17日

cpuout需要改为cpu_out

likedislike
89+ 
90+ 
91+if __name__ == "__main__":
92+ run_tests()
Aexamples/kernel_extension_aclgraph/pybind/test/trig_aclgraph_test.py+98-0
@@ -0,0 +1,98 @@
1+import torch
Y
Yyang-binrong1月17日

无文件头，开源需要文件头

likedislike
2+import torch_npu
3+from torch_npu.testing.testcase import TestCase, run_tests
4+import torch.library as library
5+import op_extension
6+ 
7+# Define Ascend custom operator library
8+ascendc_ops = library.Library("ascendc_ops", "DEF") # "DEF" means defining new operators
9+ 
10+# Define a new operator
11+ascendc_ops.define("ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor")
M
Mmihudan1月14日

检查一下tensor move有没有消掉

likedislike
12+ 
13+ 
14+# Register a meta function
15+@library.impl(ascendc_ops, "ascendc_trig", "Meta")
16+def ascendc_trig_meta(x, out_sin, out_cos):
17+ return torch.empty_like(x)
18+ 
19+ 
20+# Register implementation for the "PrivateUse1" backend
21+@library.impl(ascendc_ops, "ascendc_trig", "PrivateUse1")
Y
Yyang-binrong1月17日

ascendc_trig与PrivateUse1风格不统一

likedislike
22+def trig_custom_ops(x, out_sin, out_cos):
23+ return op_extension.run_trig_custom(x, out_sin, out_cos)
24+ 
25+ 
26+# Define a simple model using the custom operation
27+class Model(torch.nn.Module):
28+ def forward(self, x, out_sin, out_cos):
29+ out_tan = torch.ops.ascendc_ops.ascendc_trig(x, out_sin, out_cos)
30+ return out_tan
31+ 
32+ 
33+length = [8, 2048]
Y
Yyang-binrong1月17日

需要体现是哪个轴的长度

likedislike
34+ 
35+ 
36+class TestCustomTrig(TestCase):
37+ 
38+ def get_rand_input(self):
39+ x = torch.rand(length, device='npu', dtype=torch.float32)
40+ out_sin = torch.empty_like(x)
41+ out_cos = torch.empty_like(x)
42+ return x, out_sin, out_cos
43+ 
44+ # Test using torch.npu.NPUGraph
45+ def test_npugraph(self):
46+ static_x, static_out_sin, static_out_cos = self.get_rand_input()
47+ static_out_tan = torch.rand(length, device='npu', dtype=torch.float32)
48+ 
49+ g = torch.npu.NPUGraph()
50+ model = Model()
51+ with torch.npu.graph(g):
52+ static_out_tan = model(static_x, static_out_sin, static_out_cos)
53+ 
54+ real_x, real_out_sin, real_out_cos = self.get_rand_input()
55+ 
56+ static_x.copy_(real_x)
57+ static_out_sin.copy_(real_out_sin)
58+ static_out_cos.copy_(real_out_cos)
59+ # replay
60+ g.replay()
61+ self.check_res(real_x, static_out_sin, static_out_cos, static_out_tan)
62+ 
63+ # Test using torch.npu.NPUGraph
64+ def test_make_graphed_callables(self):
65+ model = Model().npu()
66+ x, out_sin, out_cos = self.get_rand_input()
67+ model = torch.npu.make_graphed_callables(model, (x, out_sin, out_cos))
68+ 
69+ real_x = torch.rand_like(x)
70+ real_out_tan = model(real_x, out_sin, out_cos)
71+ self.check_res(real_x, out_sin, out_cos, real_out_tan)
72+ 
73+ # Test using make_graphed_callables
74+ def test_npugraph_ex_backend(self):
75+ model = Model().npu()
76+ compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
Y
Yyang-binrong1月17日

fullgraph=True的格式需要修正

likedislike
77+ x, out_sin, out_cos = self.get_rand_input()
78+ out_tan = compiled_model(x, out_sin, out_cos)
79+ self.check_res(x, out_sin, out_cos, out_tan)
80+ 
81+ # Test single custom operator call
82+ def test_trig_inplace_ops(self):
83+ x, out_sin, out_cos = self.get_rand_input()
84+ out_tan = torch.ops.ascendc_ops.ascendc_trig(x, out_sin, out_cos)
85+ self.check_res(x, out_sin, out_cos, out_tan)
86+ 
87+ # Test using the npugraph_ex backend for model compilation
88+ def check_res(self, x, out_sin, out_cos, out_tan):
89+ cpu_out_sin = torch.sin(x)
90+ cpu_out_cos = torch.cos(x)
91+ cpu_out_tan = torch.tan(x)
92+ self.assertRtolEqual(out_sin, cpu_out_sin)
93+ self.assertRtolEqual(out_cos, cpu_out_cos)
94+ self.assertRtolEqual(out_tan, cpu_out_tan)
95+ 
96+ 
97+if __name__ == "__main__":
98+ run_tests()
Aexamples/kernel_extension_aclgraph/torch_library/README.md+169-0
@@ -0,0 +1,169 @@
1+# 自定义算子直调并适配aclgraph
2+ 
3+## 概述
4+ 
5+本样例展示了如何使用PyTorch的torch.library注册自定义算子，通过<<<>>>内核调用符调用核函数，并适配aclgraph使用该自定义算子，以简单的Add算子和三角函数计算的原地算子为例，实现aclgraph下自定义算子的调用。
6+ 
7+## 支持的产品
8+ 
9+- Atlas A3 训练系列产品/Atlas A3 推理系列产品
10+- Atlas A2 训练系列产品/Atlas A2 推理系列产品
11+ 
12+## 目录结构介绍
13+ 
14+```
15+├── README.md // 示例介绍
16+├── setup.py // setup文件
17+├── csrc
18+│ ├── add_custom.asc // Add算子实现 & 自定义算子注册
19+│ └── trig_inplace_custom.asc // 原地三角函数算子实现 & 自定义算子注册
20+├── op_extension
21+│ ├── __init__.py // python初始化文件
22+│ └── _load.py // 加载模块
23+└── test
24+ ├── add_aclgraph_test.py // Add算子aclgraph测试demo
25+ └── trig_aclgraph_test.py // 原地三角函数aclgraph测试demo
26+```
27+ 
28+## 算子描述
29+### Add算子
30+- 算子功能：
31+ Add算子实现了两个数据相加，返回相加结果的功能。对应的算子原型为：
32+
33+ ```
34+ ascendc_add(Tensor x, Tensor y) -> Tensor
35+ ```
36+- 算子规格：
37+
38+ <table>
39+ <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">add_custom</td></tr>
40+ </tr>
41+ <tr><td rowspan="3" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr>
42+ <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
43+ <tr><td align="center">y</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
44+ </tr>
45+ </tr>
46+ <tr><td rowspan="1" align="center">算子输出</td><td align="center">z</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr>
47+ </tr>
48+
49+ </table>
50+ 
51+### 原地三角函数算子
52+- 算子功能：
53+ 该算子入参为x, out_sin ,out_cos, 算子调用后，out_sin会被原地修改为sin(x)计算结果，out_cos会被原地修改为cos(x)计算结果，返回值tan(x)计算结果。对应的算子原型为：
54+
55+ ```
56+ ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor
57+ ```
58+- 算子规格：
59+ 
60+ <table>
61+ <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">trig_inplace_custom</td></tr>
62+ </tr>
63+ <tr><td rowspan="4" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr>
64+ <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
65+ <tr><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
66+ <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
67+
68+ </tr>
69+ </tr>
70+ <tr><td rowspan="3" align="center">算子输出</td><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
71+ <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
72+ <tr><td align="center">out_tan</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr>
73+ </tr>
74+
75+ </table>
76+ 
77+## 代码实现介绍
78+ 
79+ - 以Add算子为例，样例在*.asc文件中定义了一个名为ascendc_ops的命名空间，并在其中注册了ascendc_add函数。在ascendc_add函数中通过`c10_npu::getCurrentNPUStream()`函数获取当前NPU上的流，并通过内核调用符<<<>>>调用自定义的Kernel函数add_custom，在NPU上执行算子。
80+ ```c++
81+ add_custom<<<blockDim, nullptr, aclStream>>>(xGm, yGm, zGm, totalLength);
82+ ```
83+
84+ - PyTorch提供`TORCH_LIBRARY_FRAGMENT`宏作为自定义算子注册的核心接口，用于创建并初始化自定义算子库，注册后在Python侧可以通过`torch.ops.namespace.op_name`方式进行调用，例如：
85+
86+ ```c++
87+ TORCH_LIBRARY_FRAGMENT(ascendc_ops, m)
88+ {
89+ m.def(ascendc_add"(Tensor x, Tensor y) -> Tensor");
90+ }
91+ ```
92+
93+ - `TORCH_LIBRARY_IMPL`用于将算子逻辑绑定到特定的DispatchKey（PyTorch设备调度标识）。针对NPU设备，需要将算子实现注册到PrivateUse1这一专属的DispatchKey上，例如：
94+
95+ ```c++
96+ TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
97+ {
98+ m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
99+ }
100+ ```
101+ 
102+- 注册Meta函数：
103+
104+ 注册Meta函数使faketensor流程正常工作，在使用fx, compile等功能涉及，注册代码如下：
105+ 
106+ ```c++
107+ TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m)
108+ {
109+ m.impl("ascendc_add", &add_impl_meta);
110+ }
111+ ```
112+ 
113+- aclgraph的调用：
114+ [示例代码](./test/add_aclgraph_test.py)中，通过`torch.ops.load_library`加载生成的自定义算子库，并展示了3种aclgraph的使能方式，通过对比NPU输出与CPU标准加法结果来验证自定义算子的数值正确性。
115+ 
116+1. torch.npu.NPUGraph()
117+2. torch.npu.make_graphed_callables
118+3. backend="npugraph_ex"
119+ 
120+## 编译运行
121+ 
122+在本样例根目录下执行如下步骤，编译并执行算子。
123+ 
124+- 环境安装
125+
126+1. 请参考与您当前使用的版本配套的[《Ascend Extension for PyTorch
127+ 软件安装指南》](https://www.hiascend.com/document/detail/zh/Pytorch/720/configandinstg/instg/insg_0001.html)，获取PyTorch和torch_npu详细的安装步骤。
128+
129+ 本样例需torch2.6.0版本以上，支持`backend="npugraph_ex"`需7.3.0版本以上。
130+2. 根据实际环境安装CANN toolkit包，安装指导详见《[CANN 软件安装指南](https://www.hiascend.com/document/redirect/CannCommunityInstSoftware)》。
131+3. 根据实际环境安装CANN ops包。根据产品型号和环境架构，下载对应安装包，可参考[下载链接](https://ascend.devcloud.huaweicloud.com/cann/run/software/8.5.0-beta.1)并执行如下命令安装：
132+
133+ ```bash
134+ # 确保安装包具有可执行权限
135+ chmod +x Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run
136+ # 安装命令
137+ ./Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run --install --quiet --install-path=${install_path}
138+ ```
139+
140+ - \$\{soc\_name\}：表示NPU型号名称，即\$\{soc\_version\}删除“ascend”后剩余的内容。
141+ - \$\{install\_path\}：表示指定安装路径，需要与toolkit包安装在相同路径，默认安装在`/usr/local/Ascend`目录。
142+ 
143+- 配置环境变量
144+
145+ 请根据当前环境上CANN开发套件包的安装位置，执行如下配置环境变量的命令。
146+
147+ ```bash
148+ source ${install_path}/ascend-toolkit/set_env.sh
149+ ```
150+ 
151+ 
152+ 
153+- 样例执行
154+
155+ ```bash
156+ python setup.py bdist_wheel
157+ pip install dist/*.whl --force-reinstall
158+ cd test
159+ python ./add_aclgraph_test.py
160+ ```
161+ 
162+执行结果如下，说明精度对比成功。
163+ 
164+```bash
165+Ran * test in **s.
166+OK
167+```
168+ 
169+ 
Aexamples/kernel_extension_aclgraph/torch_library/csrc/add_custom.asc+129-0
@@ -0,0 +1,129 @@
1+// Copyright (c) 2026 Huawei Technologies Co., Ltd
2+// All rights reserved.
Y
Yyang-binrong1月17日

文件头内容需要修正

likedislike
3+//
4+// Licensed under the BSD 3-Clause License (the "License");
5+// you may not use this file except in compliance with the License.
6+// You may obtain a copy of the License at
7+//
8+// Unless required by applicable law or agreed to in writing, software
9+// distributed under the License is distributed on an "AS IS" BASIS,
10+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+// See the License for the specific language governing permissions and
12+// limitations under the License.
13+ 
14+#include <torch/extension.h>
15+#include "torch_npu/csrc/core/npu/NPUStream.h"
16+#include "kernel_operator.h"
17+ 
18+constexpr uint32_t BUFFER_NUM = 2; //tensor num for each queue
19+class KernelAdd {
20+public:
21+ __aicore__ inline KernelAdd() {}
22+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
23+ {
24+ this->blockLength = totalLength / AscendC::GetBlockNum();
Y
Yyang-binrong1月17日

【致命】没有校验除数是否为0，可能导致除0

likedislike
25+ this->tileNum = 8;
26+ this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
Y
Yyang-binrong1月17日

【致命】没有校验tileNum是否为0，可能导致除0

likedislike
27+ xGm.SetGlobalBuffer((__gm__ int32_t *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
28+ yGm.SetGlobalBuffer((__gm__ int32_t *)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
29+ zGm.SetGlobalBuffer((__gm__ int32_t *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
30+ pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int32_t));
31+ pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(int32_t));
32+ pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(int32_t));
33+ }
34+ __aicore__ inline void Process()
35+ {
36+ int32_t loopCount = this->tileNum * BUFFER_NUM;
37+ for (int32_t i = 0; i < loopCount; i++) {
38+ CopyIn(i);
39+ Compute(i);
40+ CopyOut(i);
41+ }
42+ }
43+ 
44+private:
45+ __aicore__ inline void CopyIn(int32_t progress)
46+ {
47+ AscendC::LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
48+ AscendC::LocalTensor<int32_t> yLocal = inQueueY.AllocTensor<int32_t>();
49+ AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
50+ AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
51+ inQueueX.EnQue(xLocal);
52+ inQueueY.EnQue(yLocal);
53+ }
54+ __aicore__ inline void Compute(int32_t progress)
55+ {
56+ AscendC::LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
57+ AscendC::LocalTensor<int32_t> yLocal = inQueueY.DeQue<int32_t>();
58+ AscendC::LocalTensor<int32_t> zLocal = outQueueZ.AllocTensor<int32_t>();
59+ AscendC::Add(zLocal, xLocal, yLocal, this->tileLength);
60+ outQueueZ.EnQue<int32_t>(zLocal);
61+ inQueueX.FreeTensor(xLocal);
62+ inQueueY.FreeTensor(yLocal);
63+ }
64+ __aicore__ inline void CopyOut(int32_t progress)
65+ {
66+ AscendC::LocalTensor<int32_t> zLocal = outQueueZ.DeQue<int32_t>();
67+ AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
68+ outQueueZ.FreeTensor(zLocal);
69+ }
70+ 
71+private:
72+ AscendC::TPipe pipe;
73+ AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
74+ AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
75+ AscendC::GlobalTensor<int32_t> xGm;
76+ AscendC::GlobalTensor<int32_t> yGm;
77+ AscendC::GlobalTensor<int32_t> zGm;
78+ uint32_t blockLength;
79+ uint32_t tileNum;
80+ uint32_t tileLength;
81+};
82+ 
83+__global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
84+{
85+ KernelAdd op;
86+ op.Init(x, y, z, totalLength);
87+ op.Process();
88+}
89+ 
90+namespace ascendc_ops {
91+at::Tensor ascendc_add(const at::Tensor &x, const at::Tensor &y)
92+{
93+ auto acl_stream = c10_npu::getCurrentNPUStream().stream(true);
Y
Yyang-binrong1月17日

acl_stream 与下面的blockDim 风格不一致

likedislike
94+ at::Tensor z = at::empty_like(x);
95+ uint32_t blockDim = 8;
Y
Yyang-binrong1月17日

需要以dtype相关sizeof计算替换8,防止功能拓展时散弹式修改

likedislike
96+ uint32_t totalLength = 1;
97+ for (uint32_t size : x.sizes()) {
98+ totalLength *= size;
99+ }
100+ // Launch the custom kernel use <<<>>>
101+ add_custom<<<blockDim, nullptr, acl_stream>>>((uint8_t *)(x.mutable_data_ptr()), (uint8_t *)(y.mutable_data_ptr()),
102+ (uint8_t *)(z.mutable_data_ptr()), totalLength);
103+ return z;
104+}
105+ 
106+} // namespace ascendc_ops
107+ 
108+at::Tensor add_impl_meta(const at::Tensor& x, const at::Tensor& y)
109+{
110+ return at::empty_like(x);
111+}
112+ 
113+// Define a new operator
114+TORCH_LIBRARY_FRAGMENT(ascendc_ops, m)
M
Mmihudan1月14日

补个注释

likedislike
115+{
116+ m.def("ascendc_add(Tensor x, Tensor y) -> Tensor");
117+}
118+ 
119+// Register implementation for the "PrivateUse1" backend
120+TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
Y
Yyang-binrong1月17日

PrivateUse1与ascendc_ops风格不一致

likedislike
121+{
122+ m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
123+}
124+ 
125+// Define a simple model using the custom operation
126+TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m)
Y
Yyang-binrong1月17日

Meta与ascendc_ops风格不一致

likedislike
127+{
128+ m.impl("ascendc_add", &add_impl_meta);
129+}
Aexamples/kernel_extension_aclgraph/torch_library/csrc/trig_inplace_custom.asc+147-0
@@ -0,0 +1,147 @@
1+// Copyright (c) 2026 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+//
4+// Licensed under the BSD 3-Clause License (the "License");
Y
Yyang-binrong1月17日

文件头内容需要修正

likedislike
5+// you may not use this file except in compliance with the License.
6+// You may obtain a copy of the License at
7+//
8+// Unless required by applicable law or agreed to in writing, software
9+// distributed under the License is distributed on an "AS IS" BASIS,
10+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+// See the License for the specific language governing permissions and
12+// limitations under the License.
13+ 
14+#include <torch/extension.h>
15+#include "torch_npu/csrc/core/npu/NPUStream.h"
16+#include "kernel_operator.h"
17+ 
18+constexpr uint32_t BUFFER_NUM = 2; //tensor num for each queue
19+class KernelTrig {
20+public:
21+ __aicore__ inline KernelTrig() {}
22+ 
23+ // Initialize the global memory and buffer queues
24+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan, uint32_t totalLength)
25+ {
26+ this->blockLength = totalLength / AscendC::GetBlockNum();
Y
Yyang-binrong1月17日

【致命】需要校验GetBlockNum 是否为0，防止除0

likedislike
27+ this->tileNum = 8;
28+ this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
Y
Yyang-binrong1月17日

【致命】需要校验tileNum 是否为0，防止除0

likedislike
29+ xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
30+ sinGm.SetGlobalBuffer((__gm__ float *)out_sin + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
31+ cosGm.SetGlobalBuffer((__gm__ float *)out_cos + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
32+ tanGm.SetGlobalBuffer((__gm__ float *)out_tan + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
33+ pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
34+ pipe.InitBuffer(outQueueSin, BUFFER_NUM, this->tileLength * sizeof(float));
35+ pipe.InitBuffer(outQueueCos, BUFFER_NUM, this->tileLength * sizeof(float));
36+ pipe.InitBuffer(outQueueTan, BUFFER_NUM, this->tileLength * sizeof(float));
37+ }
38+ 
39+ __aicore__ inline void Process()
40+ {
41+ int32_t loopCount = this->tileNum * BUFFER_NUM;
42+ for (int32_t i = 0; i < loopCount; i++) {
43+ CopyIn(i);
44+ Compute(i);
45+ CopyOut(i);
46+ }
47+ }
48+ 
49+private:
50+ __aicore__ inline void CopyIn(int32_t progress)
51+ {
52+ AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
53+ AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
54+ inQueueX.EnQue(xLocal);
55+ }
56+ __aicore__ inline void Compute(int32_t progress)
57+ {
58+ AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
59+ AscendC::LocalTensor<float> sinLocal = outQueueSin.AllocTensor<float>();
60+ AscendC::LocalTensor<float> cosLocal = outQueueCos.AllocTensor<float>();
61+ AscendC::LocalTensor<float> tanLocal = outQueueTan.AllocTensor<float>();
62+ 
63+ AscendC::Sin(sinLocal, xLocal, this->tileLength);
64+ AscendC::Cos(cosLocal, xLocal, this->tileLength);
65+ AscendC::Tan(tanLocal, xLocal, this->tileLength);
66+ 
67+ outQueueSin.EnQue<float>(sinLocal);
68+ outQueueCos.EnQue<float>(cosLocal);
69+ outQueueTan.EnQue<float>(tanLocal);
70+ inQueueX.FreeTensor(xLocal);
71+ }
72+ __aicore__ inline void CopyOut(int32_t progress)
73+ {
74+ // Copy the sin, cos, and tan values from local memory to global memory (inplace modification)
75+ AscendC::LocalTensor<float> sinLocal = outQueueSin.DeQue<float>();
76+ AscendC::LocalTensor<float> cosLocal = outQueueCos.DeQue<float>();
77+ AscendC::LocalTensor<float> tanLocal = outQueueTan.DeQue<float>();
78+ AscendC::DataCopy(sinGm[progress * this->tileLength], sinLocal, this->tileLength);
79+ AscendC::DataCopy(cosGm[progress * this->tileLength], cosLocal, this->tileLength);
80+ AscendC::DataCopy(tanGm[progress * this->tileLength], tanLocal, this->tileLength);
81+ outQueueSin.FreeTensor(sinLocal);
82+ outQueueCos.FreeTensor(cosLocal);
83+ outQueueTan.FreeTensor(tanLocal);
84+ }
85+ 
86+private:
87+ AscendC::TPipe pipe;
88+ AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
89+ AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueSin, outQueueCos, outQueueTan;
90+ AscendC::GlobalTensor<float> xGm;
91+ AscendC::GlobalTensor<float> sinGm;
92+ AscendC::GlobalTensor<float> cosGm;
93+ AscendC::GlobalTensor<float> tanGm;
94+ uint32_t blockLength;
95+ uint32_t tileNum;
96+ uint32_t tileLength;
97+};
98+ 
99+__global__ __aicore__ void trig_inplace_custom(GM_ADDR x, GM_ADDR out_sin, GM_ADDR out_cos, GM_ADDR out_tan,
100+ uint32_t totalLength)
101+{
102+ KernelTrig op;
103+ op.Init(x, out_sin, out_cos, out_tan, totalLength);
104+ op.Process();
105+}
106+ 
107+namespace ascendc_ops {
108+at::Tensor ascendc_trig(const at::Tensor &x, const at::Tensor &out_sin, const at::Tensor &out_cos)
109+{
110+ auto acl_stream = c10_npu::getCurrentNPUStream().stream(true);
111+ at::Tensor out_tan = at::empty_like(x);
112+ uint32_t blockDim = 8;
Y
Yyang-binrong1月17日

需要以dtype相关sizeof计算替换8与1,防止功能拓展时散弹式修改

likedislike
113+ uint32_t totalLength = 1;
114+ for (uint32_t size : x.sizes()) {
115+ totalLength *= size;
116+ }
117+ // Launch the custom kernel using <<<>>>
118+ trig_inplace_custom<<<blockDim, nullptr, acl_stream>>>(
119+ (uint8_t *)(x.mutable_data_ptr()), (uint8_t *)(out_sin.mutable_data_ptr()),
120+ (uint8_t *)(out_cos.mutable_data_ptr()), (uint8_t *)(out_tan.mutable_data_ptr()), totalLength);
121+ return out_tan;
122+}
123+ 
124+} // namespace ascendc_ops
125+ 
126+at::Tensor trig_impl_meta(const at::Tensor& x, const at::Tensor& out_sin, const at::Tensor& out_cos)
127+{
128+ return at::empty_like(x);
129+}
130+ 
131+// Define a new operator
132+TORCH_LIBRARY_FRAGMENT(ascendc_ops, m)
133+{
134+ m.def("ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor");
135+}
136+ 
137+// Register implementation for the "PrivateUse1" backend
138+TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
Y
Yyang-binrong1月17日

PrivateUse1与ascendc_ops风格不一致

likedislike
139+{
140+ m.impl("ascendc_trig", TORCH_FN(ascendc_ops::ascendc_trig));
141+}
142+ 
143+// Define a simple model using the custom operation
144+TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m)
Y
Yyang-binrong1月17日

Meta与ascendc_ops风格不一致

likedislike
145+{
146+ m.impl("ascendc_trig", &trig_impl_meta);
147+}
Aexamples/kernel_extension_aclgraph/torch_library/op_extension/__init__.py+3-0
@@ -0,0 +1,3 @@
1+from ._load import _load_opextension_so
Y
Yyang-binrong1月17日

缺少文件头，开源后需要文件头

likedislike
2+ 
3+_load_opextension_so()
Aexamples/kernel_extension_aclgraph/torch_library/op_extension/_load.py+14-0
@@ -0,0 +1,14 @@
1+import pathlib
2+import torch
3+ 
4+ 
5+# Load the custom operator library
6+def _load_opextension_so():
7+ so_dir = pathlib.Path(__file__).parents[0]
8+ so_files = list(so_dir.glob('custom_ops_lib*.so'))
9+ 
10+ if not so_files:
11+ raise FileNotFoundError(f"not find custom_ops_lib*.so in {so_dir}")
12+ 
13+ atb_so_path = str(so_files[0])
14+ torch.ops.load_library(atb_so_path)
Aexamples/kernel_extension_aclgraph/torch_library/setup.py+97-0
@@ -0,0 +1,97 @@
1+import os
2+import glob
3+import sysconfig
4+from distutils.errors import CompileError
5+from distutils.spawn import find_executable
6+import torch
7+import torch_npu
8+import torch.utils.cpp_extension as cpp_extension
9+from setuptools import setup, Extension, find_packages
10+from setuptools.command.build_ext import build_ext
11+ 
12+BASE_DIR = os.path.dirname(os.path.realpath(__file__))
13+source_files = glob.glob(os.path.join(BASE_DIR, "csrc", "*.asc"), recursive=True)
14+ 
15+ 
16+def get_dependency_paths():
17+ python_include = sysconfig.get_config_var("INCLUDEPY")
18+ python_lib = sysconfig.get_config_var("LIBDIR")
19+ 
20+ torch_include_paths = cpp_extension.include_paths()
21+ torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib")
22+ 
23+ torch_npu_path = os.path.dirname(torch_npu.__file__)
24+ torch_npu_include = os.path.join(torch_npu_path, "include")
25+ torch_npu_lib = os.path.join(torch_npu_path, "lib")
26+ 
27+ all_include_paths = list([
28+ *torch_include_paths,
29+ python_include,
30+ torch_npu_include,
31+ ])
32+ 
33+ all_libs = list([
34+ python_lib,
35+ torch_lib,
36+ torch_npu_lib,
37+ ])
38+ 
39+ return {
40+ "all_includes": all_include_paths,
41+ "all_libs": all_libs
42+ }
43+ 
44+ 
45+class AscendBuildExtension(build_ext):
46+ def __init__(self, *args, **kwargs):
47+ super().__init__(*args, **kwargs)
48+ 
49+ def _check_bisheng_compiler(self):
50+ bisheng_compiler = find_executable('bisheng')
51+ if not bisheng_compiler:
52+ raise RuntimeError("bisheng command not found!")
53+ 
54+ def build_extension(self, ext):
55+ self._check_bisheng_compiler()
56+ dep_paths = get_dependency_paths()
57+ 
58+ ext_fullpath = self.get_ext_fullpath(ext.name)
59+ os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True)
60+ 
61+ compile_cmd = [
62+ "bisheng",
63+ "-x", "asc",
64+ "--npu-arch=dav-2201",
65+ "-shared",
66+ "-fPIC",
67+ "-std=c++17",
68+ "-ltorch_npu", "-ltorch", "-lc10",
69+ *ext.sources,
70+ "-o", ext_fullpath,
71+ ]
72+ 
73+ for include_dir in dep_paths["all_includes"]:
74+ compile_cmd.append(f"-I{include_dir}")
75+ 
76+ for lib_dir in dep_paths["all_libs"]:
77+ compile_cmd.append(f"-L{lib_dir}")
78+ 
79+ try:
80+ self.spawn(compile_cmd)
81+ except Exception as e:
82+ raise CompileError(f"{str(e)}") from e
83+ 
84+ 
85+your_ext = Extension(
86+ name="op_extension.custom_ops_lib",
87+ sources=source_files,
88+ language="asc",
89+)
90+ 
91+setup(
92+ name="op_extension",
93+ version="0.1",
94+ ext_modules=[your_ext],
95+ packages=find_packages(),
96+ cmdclass={"build_ext": AscendBuildExtension},
97+)
Aexamples/kernel_extension_aclgraph/torch_library/test/add_aclgraph_test.py+73-0
@@ -0,0 +1,73 @@
1+import torch
Y
Yyang-binrong1月17日

缺少文件头

likedislike
2+import torch_npu
3+import op_extension
4+from torch_npu.testing.testcase import TestCase, run_tests
5+from torch.library import Library
6+ 
7+ 
8+# Define a simple model using the custom operation
9+class Model(torch.nn.Module):
10+ def forward(self, x, y):
11+ return torch.ops.ascendc_ops.ascendc_add(x, y)
12+ 
13+ 
14+length = [8, 2048]
Y
Yyang-binrong1月17日

数字8需要改为宏或者全局变量或加注释，防止功能拓展时散弹式修改

likedislike
15+ 
16+ 
17+class TestCustomAdd(TestCase):
18+ 
19+ def get_rand_input(self):
20+ x = torch.randint(low=1, high=100, size=length, device='npu', dtype=torch.int)
Y
Yyang-binrong1月17日

high=100格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
21+ y = torch.randint(low=1, high=100, size=length, device='npu', dtype=torch.int)
Y
Yyang-binrong1月17日

low=1格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
22+ return x, y
23+ 
24+ # Test using torch.npu.NPUGraph
25+ def test_npugraph(self):
26+ static_x, static_y = self.get_rand_input()
27+ static_target = torch.randint(low=1, high=100, size=length, device='npu:0', dtype=torch.int)
Y
Yyang-binrong1月17日

high=100格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
28+ 
29+ g = torch.npu.NPUGraph()
30+ model = Model()
31+ with torch.npu.graph(g):
32+ static_target = model(static_x, static_y)
33+ 
34+ real_x, real_y = self.get_rand_input()
35+ static_x.copy_(real_x)
36+ static_y.copy_(real_y)
37+ # replay
38+ g.replay()
39+ cpuout = torch.add(real_x, real_y)
40+ self.assertEqual(static_target, cpuout)
41+ 
42+ # Test using make_graphed_callables
43+ def test_make_graphed_callables(self):
44+ model = Model().npu()
45+ x, y = self.get_rand_input()
46+ model = torch.npu.make_graphed_callables(model, (x, y))
47+ 
48+ real_x = torch.randint_like(x, low=1, high=100)
Y
Yyang-binrong1月17日

low=1格式有问题，缺少空格，且需要改用宏，避免散弹式修改

likedislike
49+ real_y = torch.randint_like(y, low=1, high=100)
50+ output = model(real_x, real_y)
51+ cpuout = torch.add(real_x, real_y)
52+ self.assertEqual(output, cpuout)
Y
Yyang-binrong1月17日

cpuout需要改成cpu_out

likedislike
53+ 
54+ # Test using the npugraph_ex backend for model compilation
55+ def test_npugraph_ex_backend(self):
56+ model = Model().npu()
57+ compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
58+ x, y = self.get_rand_input()
59+ output = compiled_model(x, y)
60+ cpuout = torch.add(x, y)
61+ self.assertEqual(output, cpuout)
62+ 
63+ # Test single custom operator call
64+ def test_add_custom_ops(self):
65+ x, y = self.get_rand_input()
66+ output = torch.ops.ascendc_ops.ascendc_add(x.npu(), y.npu()).cpu()
67+ cpuout = torch.add(x, y)
68+ self.assertEqual(output, cpuout)
69+ 
70+ 
71+ 
72+if __name__ == "__main__":
73+ run_tests()
Aexamples/kernel_extension_aclgraph/torch_library/test/trig_aclgraph_test.py+79-0
@@ -0,0 +1,79 @@
1+import torch
2+import torch_npu
Y
Yyang-binrong1月17日

缺少文件头，开源后需要文件头

likedislike
3+import op_extension
4+from torch_npu.testing.testcase import TestCase, run_tests
5+ 
6+ 
7+# Define a simple model using the custom operation
8+class Model(torch.nn.Module):
9+ def forward(self, x, out_sin, out_cos):
10+ out_tan = torch.ops.ascendc_ops.ascendc_trig(x, out_sin, out_cos)
11+ return out_tan
12+ 
13+ 
14+length = [8, 2048]
Y
Yyang-binrong1月17日

需要以dtype相关sizeof计算替换2048,防止功能拓展时散弹式修改

likedislike
15+ 
16+ 
17+class TestCustomTrig(TestCase):
18+ 
19+ def get_rand_input(self):
20+ x = torch.rand(length, device='npu', dtype=torch.float32)
21+ out_sin = torch.empty_like(x)
22+ out_cos = torch.empty_like(x)
23+ return x, out_sin, out_cos
24+ 
25+ # Test using torch.npu.NPUGraph
26+ def test_npugraph(self):
27+ static_x, static_out_sin, static_out_cos = self.get_rand_input()
28+ static_out_tan = torch.rand(length, device='npu', dtype=torch.float32)
29+ 
30+ g = torch.npu.NPUGraph()
31+ model = Model()
32+ with torch.npu.graph(g):
33+ static_out_tan = model(static_x, static_out_sin, static_out_cos)
34+ 
35+ real_x, real_out_sin, real_out_cos = self.get_rand_input()
36+ 
37+ static_x.copy_(real_x)
38+ static_out_sin.copy_(real_out_sin)
39+ static_out_cos.copy_(real_out_cos)
40+ # replay
41+ g.replay()
42+ self.check_res(real_x, static_out_sin, static_out_cos, static_out_tan)
43+ 
44+ # Test using torch.npu.NPUGraph
45+ def test_make_graphed_callables(self):
46+ model = Model().npu()
47+ x, out_sin, out_cos = self.get_rand_input()
48+ model = torch.npu.make_graphed_callables(model, (x, out_sin, out_cos))
49+ 
50+ real_x = torch.rand_like(x)
51+ real_out_tan = model(real_x, out_sin, out_cos)
52+ self.check_res(real_x, out_sin, out_cos, real_out_tan)
53+ 
54+ # Test using make_graphed_callables
55+ def test_npugraph_ex_backend(self):
56+ model = Model().npu()
57+ compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
Y
Yyang-binrong1月17日

fullgrap需要与out_tan风格一致

likedislike
58+ x, out_sin, out_cos = self.get_rand_input()
59+ out_tan = compiled_model(x, out_sin, out_cos)
60+ self.check_res(x, out_sin, out_cos, out_tan)
61+ 
62+ # Test single custom operator call
63+ def test_trig_inplace_ops(self):
64+ x, out_sin, out_cos = self.get_rand_input()
65+ out_tan = torch.ops.ascendc_ops.ascendc_trig(x, out_sin, out_cos)
66+ self.check_res(x, out_sin, out_cos, out_tan)
67+ 
68+ # Test using the npugraph_ex backend for model compilation
69+ def check_res(self, x, out_sin, out_cos, out_tan):
70+ cpu_out_sin = torch.sin(x)
71+ cpu_out_cos = torch.cos(x)
72+ cpu_out_tan = torch.tan(x)
73+ self.assertRtolEqual(out_sin, cpu_out_sin)
74+ self.assertRtolEqual(out_cos, cpu_out_cos)
75+ self.assertRtolEqual(out_tan, cpu_out_tan)
76+ 
77+ 
78+if __name__ == "__main__":
79+ run_tests()