文档
首页
产品
解决方案
开发者
合作伙伴
高校
昇腾AI市场
产业资讯
支持与服务
注册
Atlas系列硬件
Atlas 200l DK A2 开发者套件
Atlas 200l DK 开发者套件（型号：3000）
加速模块、加速卡、智能小站、服务器、集群
异构计算架构
CANN社区版
CANN商用版
AI框架
MindSpore
TensorFlow框架模型（社区版）
TensorFlow框架模型（商用版）
PyTorch框架适配
昇腾推理引擎
MindIE
应用使能MindX
MindX DL
MindX Edge
MindX SDK
全流程开发工具
MindStudio
文档首页

算子代码实现

通过调用TBE DSL接口，在算子工程下的“tbe/impl/add.py”文件中进行Add算子的实现，包括算子函数定义、算子入参校验、compute过程实现及调度与编译。

代码模板介绍

MindStudio在“tbe/impl/add.py”中生成代码模板。 
# 导入依赖的Python模块
import tbe.dsl as tbe
from tbe import tvm
from tbe.common.register import register_op_compute
from topi import generic


# 算子计算函数
@register_op_compute(("Add")
def add_compute(x, y, z, kernel_name="add"):
    """
    To do: Implement the operator by referring to the
           TBE Operator Development Guide.
    """

    res = tbe.XXX(x, y)
    return res


# 算子定义函数
def add(x, y, z, kernel_name="add"):
    """
    To do: Implement the operator by referring to the
           TBE Operator Development Guide.
    """
    # 输入参数占位
    data_x = tvm.placeholder(x.get("shape"), dtype=x.get("dtype"), name="data_x")
    data_y = tvm.placeholder(y.get("shape"), dtype=y.get("dtype"), name="data_y")

    # 调用算子计算函数
    res = add_compute(data_x, data_y, z, kernel_name)

    # 自动调度
    with tvm.target.cce():
        schedule = tbe.auto_schedule(res)

    # 编译
    config = {"name": kernel_name,
              "tensor_list": [data_x, data_y, res]}
    tbe.build(schedule, config)
  • 引入算子开发时依赖的Python模块。
    • “tbe.dsl”：引入TBE支持的特定域语言接口，包括常见的运算vmuls、vadds、matmul等。

      具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/dsl目录下的Python函数。

    • “tbe.tvm”：引入TVM后端代码生成机制。

      具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/tvm目录下的Python函数，使用方法请参见https://docs.tvm.ai/。

    • “tbe.common.register.register_op_compute”：提供了实现算子的UB自动融合的接口。

      具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/common/register/register_api.py文件中的register_op_compute函数的定义。

  • 模板生成以算子名称_compute命名的计算函数声明。
    • 若创建算子工程时选择的“Sample Template”或者“Tensorflow Template”，输入输出参数及属性会根据原型定义自动生成。
    • 若创建算子工程时选择“Empty Template”，默认生成的参数为一个输入与一个输出，不带属性。
  • 模板生成以算子名称命名的定义函数的声明与部分实现，模板中提供的实现函数中的示例代码包含如下功能：
    • 获取输入tensor的shape与dtype。
    • 对输入参数进行校验。
    • 对输入tensor进行占位。
    • 调用算子的compute函数进行计算、调度与编译。

算子定义函数实现

开发者需要根据MindStudio生成的模板代码进行算子计算函数的实现，同时需要在算子定义函数中增加算子输入/输出/属性的校验代码。由于Add算子允许两个输入数据的shape不同，但算子计算接口tbe.dsl.vadd( )要求两个输入shape相同，因此需要对算子两个输入的shape进行广播，并对其进行校验。有助于在算子编译阶段，提前发现问题。修改后代码如下所示：

若开发者通过Windows上远程打开的MindStudio，复制以下代码时可能会出现无法复制的情况，解决方法请参见Windows上远程打开MindStudio时，复制的内容无法粘贴到编辑器窗口中。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
 
from __future__ import absolute_import
import tbe.dsl as tbe
from functools import reduce
from tbe import tvm
from tbe.common.register import register_op_compute
from tbe.common.utils import para_check
from tbe.common.utils import shape_util


SHAPE_SIZE_LIMIT = 2147483648

# 实现Add算子的计算逻辑
@register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    # 将shape转换为list
    shape_x = shape_util.shape_to_list(input_x.shape)    
    shape_y = shape_util.shape_to_list(input_y.shape) 

    # shape_max取shape_x与shape_y的每个维度的大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")
    shape_size = reduce(lambda x, y: x * y, shape_max[:])
    if shape_size > SHAPE_SIZE_LIMIT:
        raise RuntimeError("the shape is too large to calculate")

    # 将input_x的shape广播为shape_max
    input_x = tbe.broadcast(input_x, shape_max)
    input_y = tbe.broadcast(input_y, shape_max)

    # 执行input_x + input_y
    res = tbe.vadd(input_x, input_y)

    return res

# 算子定义函数
@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT,
                            para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
def add(input_x, input_y, output_z, kernel_name="add"):
    # 获取算子输入tensor的shape与dtype
    shape_x = input_x.get("shape")
    shape_y = input_y.get("shape")

    # 检验算子输入类型
    check_tuple = ("float16", "float32", "int32")
    input_data_type = input_x.get("dtype").lower()
    para_check.check_dtype(input_data_type, check_tuple, param_name="input_x")

    # shape_max取shape_x与shape_y的每个维度的最大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")

     # 如果shape的长度等于1，就直接赋值，如果shape的长度不等于1，做切片，将最后一个维度舍弃（按照内存排布，最后一个维度为1与没有最后一个维度的数据排布相同，例如2*3=2*3*1，将最后一个为1的维度舍弃，可提升后续的调度效率）
    if shape_x[-1] == 1 and shape_y[-1] == 1 and shape_max[-1] == 1:
        shape_x = shape_x if len(shape_x) == 1 else shape_x[:-1]
        shape_y = shape_y if len(shape_y) == 1 else shape_y[:-1]
        shape_max = shape_max if len(shape_max) == 1 else shape_max[:-1]

    # 使用TVM的placeholder接口对输入tensor进行占位，返回一个tensor对象
    data_x = tvm.placeholder(shape_x, name="data_1", dtype=input_data_type)
    data_y = tvm.placeholder(shape_y, name="data_2", dtype=input_data_type)

    # 调用compute实现函数
    res = add_compute(data_x, data_y, output_z, kernel_name)

    # 自动调度
    with tvm.target.cce():
        schedule = tbe.auto_schedule(res)
    # 编译配置
    config = {"name": kernel_name,
              "tensor_list": (data_x, data_y, res)}
    tbe.build(schedule, config)
  1. 算子定义函数声明包含算子输入信息、输出信息以及内核名称。
    def add(input_x, input_y, output_z, kernel_name="add"):
    • input_x, input_y：Add算子的两个输入tensor，每个tensor需要采用字典的形式进行定义，包含shape和dtype等信息。输入tensor的个数需要与算子信息定义文件“tbe/op_info_cfg/ai_core/add.ini”中的定义保持一致。
    • output_z：输出tensor，包含shape和dtype等信息，字典格式，此字段为预留位。

      输出tensor的个数需要与算子信息定义文件“tbe/op_info_cfg/ai_core/add.ini”中的定义保持一致。

    • kernel_name：算子在内核中的名称（即生成的二进制文件以及算子描述文件的名称），用户自定义，保持唯一，只能是大小写字母、数字、“_”的组合，且必须是字母或者“_”开头，长度小于或等于200个字符。
  2. 算子输入校验和输出形状推导。

    Add算子需要对两个输入tensor的shape进行校验，且仅支持数据类型float16, float32, int32，此外Add算子允许两个输入shape不同，因此需要调用shape_util.broadcast_shapes()生成广播后的shape并对其进行校验。

  3. 对输入tensor进行占位。

    调用TVM的placeholder接口分别对两个输入tensor进行占位，并分别返回一个tensor对象。

    5中的tensor_list的输入tensor需要是tvm.placeholder接口返回的tensor对象，所以此对象在后续计算过程实现中不能被替换。

  4. 调用add_compute计算函数。
    res = add_compute(data_x, data_y, output_z, kernel_name)

    data_x、data_y为3生成的占位tensor对象。

    计算函数的实现请参见算子Compute函数实现。

  5. 算子调度与编译实现。

    调度配置config中的tensor_list为：

    "tensor_list": (data_x, data_y, res)

    分别为两个输入tensor与一个输出tensor。

算子Compute函数实现

开发者需要根据算子计算逻辑自定义实现算子compute函数，Add算子的Compute函数实现如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
 
@register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    # 将shape转换为list
    shape_x = shape_util.shape_to_list(input_x.shape)    
    shape_y = shape_util.shape_to_list(input_y.shape) 

    # shape_max取shape_x与shape_y的每个维度的大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")
    shape_size = reduce(lambda x, y: x * y, shape_max[:])
    if shape_size > SHAPE_SIZE_LIMIT:
        raise RuntimeError("the shape is too large to calculate")

    # 将input_x的shape广播为shape_max
    input_x = tbe.broadcast(input_x, shape_max)
    input_y = tbe.broadcast(input_y, shape_max)

    # 执行input_x + input_y
    res = tbe.vadd(input_x, input_y)

    return res
  1. add_compute函数的声明如下所示。
    @register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add")

    装饰器@register_op_compute("Add", op_mode="dynamic", support_fusion=True)是DSL算子开发方式中必需的，其作用是整网运行时算子支持做UB自动融合，使得当前自定义算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装。

    其中：

    • input_x，input_y：compute函数的入参，为在3中声明的输入tensor对应的placeholder，包含shape和dtype等信息。
    • output_z：为1中算子接口函数透传过来的dict类型。
    • kernel_name：算子在内核中的名称。
  2. 进行Add算子的计算逻辑的实现。

    Add算子要求相加的两个tensor的shape相同，所以首先通过调用tbe.broadcast接口将两个输入tensor广播成相同的shape，然后调用tbe.vadd接口实现输入tensor的相加，并返回计算结果tensor。

算子编译验证

  1. 在算子python文件最下方添加main函数调用该算子，通过MindStudio编译算子实现文件，用于单算子代码的简单语法校验，代码示例如下所示：
    1
2
3
4
    		 
    # 算子调用
if __name__ == '__main__':
    input_output_dict = {"shape": (5, 6, 7),"format": "ND","ori_shape": (5, 6, 7),"ori_format": "ND", "dtype": "float16"}
    add(input_output_dict, input_output_dict, input_output_dict, kernel_name="add")
  2. 右键单击“tbe/impl/add.py”，选择“Run‘add'”，编译算子。
    如果编译没有报错，且在当前目录“tbe/impl”下生成kernel_meta文件夹，包括以下文件，则表示算子代码能够编译运行。
    • 算子二进制文件*.o。
    • 算子描述文件*.json：用于定义算子属性及运行时所需要的资源。
父主题： TBE算子开发
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容，请尝试其他搜索词