Star854
1342
代码
代码
Issues1643
Pull Requests2204
项目讨论
Wiki
分析
Star854
1342
已合并
delete cpp_extensions examples #15927
zhanhao创建于 2024年11月12日
delete cpp_extensions examples #15927
已合并
zhanhao创建于 2024年11月12日
从refs/pull/15927/head合入到master
共 1 个文件变更+0-302
Dexamples/cpp_extension/README.md+0-302
@@ -1,302 +0,0 @@
1-# 构建自定义算子wheel包工程
2- 
3-## 简介
4- 
5-本样例介绍自定义API样例打包工程，包括如何编写自定义API，绑定前反向自动求导，编译生成wheel包。具体代码样例详见[样例工程代码](https://gitee.com/ascend/samples/tree/master/operator/AddCustomSample/FrameworkLaunch/CppExtensions/setup)，具体使用方法详见[使用说明](https://gitee.com/ascend/samples/tree/master/operator/AddCustomSample/FrameworkLaunch#%E4%BD%BF%E7%94%A8%E7%BC%96%E8%AF%91wheel%E5%8C%85%E7%9A%84%E6%96%B9%E5%BC%8F%E8%B0%83%E7%94%A8)。
6- 
7-> 注：使用该样例工程之前需要先编译并部署底层的自定义算子包，详见[编译部署自定义算子包](https://gitee.com/ascend/samples/tree/master/operator/AddCustomSample/FrameworkLaunch#%E7%BC%96%E8%AF%91%E7%AE%97%E5%AD%90%E5%B7%A5%E7%A8%8B%E9%83%A8%E7%BD%B2%E7%AE%97%E5%AD%90%E5%8C%85)。
8- 
9-## 代码介绍
10- 
11-### 目录结构
12- 
13-```
14-setup
15-├── csrc # 算子适配层及绑定代码
16-│   ├── extension_add1.cpp # 自定义算子适配层
17-│   ├── extension_add.cpp # 自定义算子适配层
18-│   ├── function.h # 自定义算子函数声明
19-│   ├── pytorch_npu_helper.hpp # 算子下发框架（无需关注）
20-│   └── register.cpp # 自定义算子注册绑定等
21-├── custom_ops # python打包目录
22-│   ├── add_custom.py # 定义python侧接口
23-│   └── __init__.py # 初始化代码
24-├── graph # 图模式相关代码
25-│   ├── CMakeLists.txt # 编译文件（无需关注）
26-│   ├── codegen.cpp # 编译文件（无需关注）
27-│   ├── custom_reg_op.h # 存放算子原型REG_OP
28-│   └── operator_reg.h # 编译所需头文件（无需关注）
29-├── setup.py # 编译打包文件
30-└── test # 测试用例
31- ├── test_add_custom_graph.py # 图模式测试用例
32- └── test_add_custom.py # 自定义算子测试用例
33- 
34-```
35- 
36-### 注册自定义算子
37- 
38-#### 注册自定义算子schema
39- 
40-首先通过TORCH_LIBRARY宏注册一个名为`myops`的命名空间，注意命名空间名字必须是唯一的。在`myops`命名空间里注册两个自定义schema，分别为前向和反向。如果需要注册多个schema，只需在同一个命名空间里继续添加即可。
41- 
42-```c++
43-// register.cpp
44-#include <torch/library.h>
45- 
46-TORCH_LIBRARY(myops, m) {
47- m.def("my_op(Tensor self, Tensor other) -> Tensor");
48- m.def("my_op_backward(Tensor self) -> (Tensor, Tensor)");
49-}
50-```
51- 
52-#### 适配层kernel编写
53- 
54-以add算子为例，编写自定义add的kernel，通过aclnn方式调用底层npu算子实现。
55- 
56-```c++
57-// extension_add.cpp
58-// 为NPU设备注册前向实现
59-at::Tensor my_op_impl_npu(const at::Tensor& self, const at::Tensor& other) {
60- // 创建输出内存
61- at::Tensor result = at::Tensor(self);
62- 
63- // 调用底层aclnn接口计算
64- EXEC_NPU_CMD(aclnnAddCustom, self, other, result);
65- return result;
66-}
67- 
68-// 为NPU设备注册反向实现
69-std::tuple<at::Tensor, at::Tensor> my_op_backward_impl_npu(const at::Tensor& self) {
70- // 创建输出内存
71- at::Tensor result = at::Tensor(self);
72- 
73- return {result, result};
74-}
75-```
76- 
77-#### 为NPU设备绑定前向和反向的对应实现
78- 
79-给新注册的schema添加了PrivateUse1设备的前反向对应实现，npu在Pytorch2.1及以上版本用的dispatch key是`PrivateUse1`，2.1以下是`XLA`。
80- 
81-```c++
82-// extension_add.cpp
83-TORCH_LIBRARY_IMPL(myops, PrivateUse1, m) {
84- m.impl("my_op", &my_op_impl_npu);
85- m.impl("my_op_backward", &my_op_backward_impl_npu);
86-}
87-```
88- 
89-#### 添加前反向绑定
90- 
91-调用my_op的时候通过注册的schema和对应的disptach key来找到对应的实现。
92- 
93-```c++
94-// extension_add.cpp
95-// 寻找注册在my_op上的不同设备的实现
96-at::Tensor my_op_impl(const at::Tensor& self, const at::Tensor& other) {
97- static auto op = torch::Dispatcher::singleton()
98- .findSchemaOrThrow("myops::my_op", "")
99- .typed<decltype(my_op_impl)>();
100- return op.call(self, other);
101-}
102-// 寻找注册my_op_backward上的不同设备的实现
103-std::tuple<at::Tensor, at::Tensor> my_op_backward_impl(const at::Tensor& self) {
104- static auto op = torch::Dispatcher::singleton()
105- .findSchemaOrThrow("myops::my_op_backward", "")
106- .typed<decltype(my_op_backward_impl)>();
107- return op.call(self);
108-}
109-```
110- 
111-Pytorch提供了`torch::autograd::Function`方法实现前反向绑定，在里面定义`forward`和`backward`函数，通过`apply`方法调用，同时把自动求导的方法实现注册到`AutogradPrivateUse1`，实现自动求导，如果pytorch版本是2.1以下，需要将`AutogradPrivateUse1`换成`AutogradXLA`。
112- 
113-```c++
114-// extension_add.cpp
115-class MyAddFunction : public torch::autograd::Function<MyAddFunction> {
116- public:
117- static at::Tensor forward(AutogradContext *ctx, at::Tensor self, at::Tensor other) {
118- at::AutoDispatchBelowADInplaceOrView guard;
119- return my_op_impl(self, other);
120- }
121- 
122- static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) {
123- auto grad_output = grad_outputs[0];
124- auto result = my_op_backward_impl(grad_output);
125- return {std::get<0>(result), std::get<1>(result)};
126- }
127-};
128- 
129-at::Tensor my_op_impl_autograd(const at::Tensor& self, const at::Tensor& other) {
130- return MyAddFunction::apply(self, other);
131-}
132- 
133-TORCH_LIBRARY_IMPL(myops, PrivateUse1, m) {
134- m.impl("my_op", &my_op_impl_npu);
135- m.impl("my_op_backward", &my_op_backward_impl_npu);
136-}
137- 
138-// 给自定义op绑定NPU的自动求导实现
139-// 如果是pytorch 2.1以下的版本，AutogradPrivateUse1需要改成AutogradXLA
140-TORCH_LIBRARY_IMPL(myops, AutogradPrivateUse1, m) {
141- m.impl("my_op", &my_op_impl_autograd);
142-}
143-```
144- 
145-#### 通过pybind绑定C++和Python接口
146- 
147-通过pybind11提供的接口将C++侧接口和Python侧接口绑定，这样在python可以用`add_custom`调用。
148- 
149-```c++
150-// register.cpp
151-#include <torch/extension.h>
152-#include "function.h"
153- 
154-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
155- m.def("add_custom", &my_op_impl_autograd, "x + y");
156-}
157-```
158- 
159-`add_custom`：python侧调用的接口名，可以自定义名称。
160- 
161-`my_op_impl_autograd`：c++侧的函数名，需要是上面已经定义的函数。
162- 
163-`"x + y"`：接口描述，无实际用途，可选参数。
164- 
165-### 添加python侧调用接口
166- 
167-c++侧代码最终会编译生成so（具体编译步骤在下面章节[编译打包工程](#编译打包工程)会介绍），这里假设生成的so名称为`custom_ops_lib`，需要在python提供对外接口以供调用。在custom_ops目录下的add_custom.py里面新增python侧接口，并在__init__.py里面import该接口。
168- 
169-```python
170-# add_custom.py
171-import custom_ops_lib
172- 
173-def add_custom(self, other):
174- return custom_ops_lib.add_custom(self, other)
175- 
176-# __init__.py
177-import custom_ops_lib
178-from .add_custom import add_custom
179-```
180- 
181-### 图模式适配（可选）
182- 
183-此步骤可以让用户注册的自定义算子增加入图能力，实现图模式的功能，如不需要入图则可以跳过。
184- 
185-> 注：入图前需要完成前面步骤的自定义算子注册。
186- 
187-#### 添加meta设备的实现
188- 
189-meta tensor用于图模式中的infershape，是一种没有具体值的tensor，只有形状，数据类型，内存结构等信息。注册方法与npu设备一样，需要实现前反向，并且注册到`Meta`这个dispatch key上。
190- 
191-```c++
192-// extension_add.cpp
193- 
194-// 为Meta设备注册前向实现
195-at::Tensor my_op_impl_meta(const at::Tensor& self, const at::Tensor& other) {
196- return empty_like(self);
197-}
198- 
199-// 为Meta设备注册反向实现
200-std::tuple<at::Tensor, at::Tensor> my_op_backward_impl_meta(const at::Tensor& self) {
201- auto result = empty_like(self);
202- return std::make_tuple(result, result);
203-}
204- 
205-TORCH_LIBRARY_IMPL(myops, Meta, m) {
206- m.impl("my_op", &my_op_impl_meta);
207- m.impl("my_op_backward", &my_op_backward_impl_meta);
208-}
209-```
210- 
211-#### 编译生成GE构图API
212- 
213-通过自动生成脚本将算子原型REG_OP转换为GE构图API，并通过torchair仓提供的接口完成converter注册，从而使能自定义算子的入图能力。
214- 
215-> 注：此步骤为NPU设备上入图特有的操作。
216- 
217-将自定义算子的算子原型REG_OP添加到graph目录下的custom_reg_op.h文件中。算子原型来源于算子工程编译结果。
218- 
219-```c++
220-// custom_reg_op.h
221-#include "operator_reg.h"
222- 
223-namespace ge {
224-REG_OP(AddCustom)
225- .INPUT(x, ge::TensorType::ALL())
226- .INPUT(y, ge::TensorType::ALL())
227- .OUTPUT(z, ge::TensorType::ALL())
228- .OP_END_FACTORY_REG(AddCustom);
229-}
230- 
231-```
232- 
233-执行编译命令，会在graph目录下生成文件auto_generated_ge_raw_custom_ops.py，将其拷贝到自己的工程或者拷贝源码至自己的调用文件里，保证能够调用到即可。
234- 
235-```bash
236-mkdir build
237-cd build
238-cmake ..
239-make generate_ge_raw_custom_ops
240-```
241- 
242-通过torchair提供的装饰器@register_fx_node_ge_converter完成自定义aten ir的converter注册。
243- 
244-```python
245-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, Callable
246-import torchair
247-from torchair.ge_concrete_graph.utils import dtype_promote
248-from torchair.ge_concrete_graph.fx2ge_converter import register_fx_node_ge_converter
249-from torchair.ge_concrete_graph.ge_graph import Tensor, TensorSpec
250-from auto_generated_ge_raw_custom_ops import AddCustom
251- 
252-@register_fx_node_ge_converter(torch.ops.myops.my_op.default)
253-def conveter_custom_op(
254- input1: Tensor,
255- input2: Tensor,
256- out: Tensor = None,
257- meta_outputs: Any = None):
258- input1, input2 = dtype_promote(input1, input2, target_dtype=meta_outputs.dtype)
259- return AddCustom(input1, input2)
260-```
261- 
262-至此完成自定义算子入图适配工作，用户可以运行参考用例中的示例验证。
263- 
264-### 编译打包工程
265- 
266-完成自定义算子注册以后，需要通过setuptools工具来同时编译c++代码，并将生成的so和custom_ops目录下的python代码一起生成wheel包。setuptools工具通过执行setup.py来编译打包，里面会指定编译的一些参数。
267- 
268-```python
269-# setup.py
270-import torch
271-from setuptools import setup, find_packages
272-from torch.utils.cpp_extension import BuildExtension
273-import torch_npu
274-from torch_npu.utils.cpp_extension import NpuExtension
275- 
276-USE_NINJA = os.getenv('USE_NINJA') == '1'
277-exts = []
278-ext1 = NpuExtension(
279- name="custom_ops_lib",
280- # 如果还有其他cpp文件参与编译，需要在这里添加
281- sources=["./csrc/extension_add.cpp", "./csrc/extension_add1.cpp", "./csrc/register.cpp"],
282- extra_compile_args = [
283- '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/acl/inc"),
284- ],
285-)
286-exts.append(ext1)
287- 
288-setup(
289- name="custom_ops",
290- version='1.0',
291- keywords='custom_ops',
292- ext_modules=exts,
293- packages=find_packages(),
294- cmdclass={"build_ext": BuildExtension.with_options(use_ninja=USE_NINJA)},
295-)
296-```
297- 
298-setup里面指定了wheel包的名称、版本号、需要编译的扩展、打包的文件、编译命令等，同时ext_modules里面指定了c++编译源文件、头文件路径、编译选项等参数。通过执行`python3 setup.py build bdist_wheel`命令会编译打包，在dist目录下会生成wheel包，通过pip install的方式可以安装该wheel包，具体使用方法参考测试用例。
299- 
300-### 测试用例
301- 
302-在test目录下提供了两个测试用例，test_add_custom.py是自定义API调用的测试用例，里面包含了whl包具体调用方法和API调用方法；test_add_custom_graph.py则是图模式的相关测试用例