java与python3混合开发_AscendPytorch算子适配层开发

Ascend Pytorch算子适配层开发

适配方法

找到和PyTorch算子功能对应的NPU TBE算子&＃xff0c;根据算子功能计算出输出Tensor的size&＃xff0c;再根据TBE算子原型构造对应的input/output/attr&＃xff0c;传递给ACL完成TBE算子的执行。

说明&＃xff1a;

TBE算子实现的源文件存放路径由开发套件包Toolkit的安装方式决定&＃xff1a;

· 若使用root用户安装&＃xff0c;则存放在&＃xff1a;/usr/local/Ascend/ascend-toolkit/latest/opp/op_impl/built-in/ai_core/tbe/impl/

· 若使用非root用户安装&＃xff0c;则存放在&＃xff1a;~/.local/Ascend/ascend-toolkit/latest/opp/op_impl/built-in/ai_core/tbe/impl/

开发者可以通过查看算子实现源文件&＃xff0c;确定算子的功能。

存放路径和命名格式

对NPU的TBE算子适配文件保存在pytorch/aten/src/ATen/native/npu目录下&＃xff0c;命名风格采用大驼峰&＃xff0c;命名格式&＃xff1a; &＃43; .cpp&＃xff0c;如&＃xff1a;AddKernelNpu.cpp。

适配步骤

须知&＃xff1a;

适配代码基于C&＃43;&＃43;开发。

1. 引入依赖头文件。

2. #include "ATen/native/npu/utils/CalcuOpUtil.h"

3. #include "ATen/native/npu/utils/KernelNpuOutputSize.h"

#include "ATen/native/npu/utils/NpuUtils.h"

说明&＃xff1a;

"CalcuOpUtil.h"文件中主要包含与ACL接口相关的函数。

"KernelNpuOutputSize.h"中主要包含算子输出shape的推导函数。

"NpuUtils.h"文件中主要包含公共能力的函数。

4. 定义Add算子适配主体函数。

结合native_functions.yaml 中 add算子的分发定义&＃xff0c;算子适配中应包含如下函数&＃xff1a;

o add_npu_input 构造输入的NPUTensorDesc对象

o add_npu_output 构造输出的NPUTensorDesc对象

o add_npu_attr 构造NPU TBE Add算子attr属性

o add_out_npu 算子适配函数(yaml中npu派发函数&＃xff0c;支持传入输出tensor)&＃xff0c;other参数支持 Tensor & Scalar

o add_npu 算子适配函数(yaml中npu派发函数)&＃xff0c;other参数支持 Tensor & Scalar

5. 实现函数 add_npu_input。

将NPU适配函数(add_npu_input)的输入构造成NPUTensorDesc对象。

// 输入参数为"self": "Tensor"和"other": "Tensor"时&＃xff0c;适配函数add_npu_input的实现

SmallVector add_npu_input(const Tensor& self,const Tensor& other) {

bool isSelfWrapped &＃61; CalcuOpUtil::is_scalar_wrapped_to_tensor(self);

bool isOtherWrapped &＃61; CalcuOpUtil::is_scalar_wrapped_to_tensor(other);

auto inputs &＃61; CalcuOpUtil::create_npu_input_tensor_desc({self, other});

// &＃39;t &＃43; 2&＃39; to work with any type of tensor, not just LongTensor (which is what

// integersin Python represent).

if (isSelfWrapped && (!isOtherWrapped)) {

inputs[0].scalarType &＃61; other.scalar_type();

} else if (isOtherWrapped && (!isSelfWrapped)) {

inputs[1].scalarType &＃61; self.scalar_type();

}

return inputs;

}

// 输入参数为"self": "Tensor"和"other": "Scalar"时&＃xff0c;适配函数add_npu_input的实现

SmallVector add_npu_input(const Tensor& self,const Scalar& other) {

return CalcuOpUtil::create_npu_input_tensor_desc({self});

}

6. 实现函数 add_npu_output。

将函数 add_npu_output的输出tensor对象构造成NPUTensorDesc对象。

// 输出参数为 "Tensor" 时&＃xff0c;适配函数add_npu_output的实现

SmallVector add_npu_output(const Tensor& result) {

return CalcuOpUtil::create_npu_output_tensor_desc({result});

}

说明&＃xff1a;

一般来说&＃xff0c;算子的输出不需要特殊处理&＃xff0c;直接调用CreateNpuOutputTensorDesc即可。

7. 实现函数 add_npu_attr。

根据NPU TBE算子原型中所需的attr规格&＃xff0c;将参数适配成NPU TBE算子原型所需要的attr属性。

// 输入参数为"other": "Tensor"和"alpha": "Scalar"时&＃xff0c;对应的适配函数add_npu_attr实现

SmallVector add_npu_attr(const Tensor& self, const Tensor& other, Scalar alpha) {

float value &＃61; CalcuOpUtil::get_scalar_float_value(alpha);

NPUAttrDesc npuAttrScalar &＃61; NPUAttrDesc("alpha", value);

SmallVector attrs &＃61; {npuAttrScalar};

return attrs;

}

// 输入参数为"other": "Scalar"和"alpha": "Scalar"时&＃xff0c;对应的适配函数adds_npu_attr实现

SmallVector adds_npu_attr(const Tensor& self,const Scalar& other,const Scalar& alpha) {

float otherValue &＃61; CalcuOpUtil::get_scalar_float_value(other);

float alphaValue &＃61; CalcuOpUtil::get_scalar_float_value(alpha);

float value &＃61; otherValue * alphaValue;

NPUAttrDesc npuAttrValue &＃61; NPUAttrDesc("value", value);

SmallVector attrs &＃61; {npuAttrValue};

return attrs;

}

8. 实现函数 add_out_npu。

9. Tensor& add_out_npu(Tensor& result, const Tensor& self, const Tensor& other, Scalar alpha) {

10. if (other.dim() &＃61;&＃61; 0 && !other.is_npu()) {

11. adds_out_npu(result, self, other.item(), alpha);

12. } else if (self.dim() &＃61;&＃61; 0 && !self.is_npu()) {

13. adds_out_npu(result, other, self.item(), alpha);

14. } else {

15. // constructs the input and output NPUTensorDesc

16. auto inputs &＃61; add_npu_input(self, other);

17. auto outputs &＃61; add_npu_output({result});

18.

19. // constructs the attr of the NPUAttrDesc

20. auto attrs &＃61; add_npu_attr(self, other, alpha);

21. // executing the NPU operator

22. CalcuOpUtil::execute_npu_operate("Axpy", inputs, outputs, attrs);

23. }

24.

25. return result;

}

说明&＃xff1a;

add_out_npu和add_npu的差别是add_out_npu支持显示指定输出tensor&＃xff0c;往输出tensor中写入结果。

26. 实现函数 add_npu。

a. 定义并实现算子的shape推导函数&＃xff0c;根据输入参数计算输出的size。

Shape推导函数定义规范&＃xff1a;

"NPU适配函数名称" &＃43; "_" &＃43; "output" &＃43; "_" &＃43; "size"&＃xff0c;如add_npu_output_size()&＃xff1b;

说明&＃xff1a;

§ Shape推导函数定义和实现存放在 pytorch/aten/src/ATen/native/npu/utils&＃xff0c;对应的头文件和实现在 KernelNpuOutPutSize.h 和 KernelNpuOutPutSize.cpp中。

§ 在KernelNpuOutPutSize.h中&＃xff0c;函数存放位置按照函数名字排序。

//输入参数为"self": "Tensor"和"other": "Tensor"时&＃xff0c;Shape推导该函数

SmallVector add_npu_output_size(const Tensor& self,const Tensor& other) {

return broadcast_ops_npu_output_size(self, other); //定义Shape推导函数

}

// 输入参数为"self": "Tensor"和"other": "Scalar"时&＃xff0c;Shape推导该函数

IntArrayRef add_npu_output_size(const Tensor& self, const Scalar& other) {

return input_same_output_size(self);

}

说明&＃xff1a;

broadcast_ops_npu_output_size函数的作用是&＃xff1a;当两个参数符合PyTorch广播机制时&＃xff0c;函数会将两个参数自动扩展为相等大小

b. 调用对应的shape推导函数计算输出的size。

c. 根据输出的size调用at::empty_with_ format创建输出Tensor&＃xff0c;函数支持指定输出Tensor的format&＃xff0c;默认为NCHW格式。

说明&＃xff1a;

当前制定的Format设置规则为重型算子锚点扩散&＃43;连续性法则混合规则。

§ 重型算子如卷积、Matmul&＃xff0c;只支持某种特定format&＃xff0c;适配时显示指定为其需要的format&＃xff0c;format向周边扩散。

§ 而连续性法则指的是算子对格式不敏感&＃xff0c;算子format指定为与第一个输入tensor的format相同即可。

§ NPU中的卷积只支持NC1HWC0格式&＃xff0c;所以需要显式指定为NC1HWC0格式

d. 将构造好的输出Tensor和其他参数传给add_out_npu进行运算

e. // 输入参数为"self": "Tensor"和"other": "Tensor"时&＃xff0c;对应的适配函数add_npu实现

f. //调用对应的Shape推导函数计算输出的size

g. Tensor add_npu(const Tensor& self, const Tensor& other, Scalar alpha) {

h. Tensor outputTensor &＃61; add_dest_output(self, other);

i. auto outputSize &＃61; add_npu_output_size(self, other);

j.

k. //根据输出的size调用at::empty_with_format创建输出Tensor&＃xff0c;函数支持指定输出Tensor的format&＃xff0c;默认为NCHW格式

l. Tensor result &＃61; at::empty_with_format(outputSize, outputTensor.options(), CalcuOpUtil::get_tensor_npu_format(outputTensor));

m.

n. //将构造好的输出Tensor和其他参数传给add_out_npu进行运算

o. add_out_npu(result, self, other, alpha);

p. return result;

q. }

r.

s. // 输入参数为"self": "Tensor"和"other": "Scalar"时&＃xff0c;对应的适配函数add_npu实现

t. //调用对应的Shape推导函数计算输出的size

u. Tensor add_npu(const Tensor& self, Scalar other, Scalar alpha) {

v. auto outputSize &＃61; add_npu_output_size(self, other);

w.

x. //根据输出的size调用at::empty_with_format创建输出Tensor&＃xff0c;函数支持指定输出Tensor的format&＃xff0c;默认为NCHW格式

y. Tensor result &＃61; at::empty_with_format(outputSize, self.options(), CalcuOpUtil::get_tensor_npu_format(self));

z.

aa. //将构造好的输出Tensor和其他参数传给add_out_npu进行运算

bb. adds_out_npu(result, self, other, alpha);

cc. return result;

}