带你了解AKG正反向算子注册+关联流程【图文】

​​摘要：简要介绍一下akg正反向算子的注册和关联流程。

本文分享自华为云社区​​《AKG正反向算子注册+关联》​​，作者：木子_007 。

一、环境

硬件：eulerosv2r8.aarch64

mindspore：1.1

算子注册需要编译安装框架才能生效，所以默认环境中已经有了mindspore的源码，并且已经可以编译安装

二、正向算子制作及测试

这里制作一个计算向量平方的算子

正向：y = x**2

反向：y = 2*x

先介绍正向

2.1 定义正向算子

路径：mindspore/akg/python/akg/ms/cce/，创建cus_square.py

参照同级目录下计算逻辑的定义，定义向量平方的计算逻辑

"""cus_square"""
from akg.tvm.hybrid import script
from akg.ops.math import mul
import akg
def CusSquare(x):
    output_shape = x.shape
    k = output_shape[0]
    n = output_shape[1]

    @script
    def cus_square_compute(x):
        y = output_tensor(output_shape, dtype=x.dtype)
        for i in range(k):
            for j in range(n):
                y[i, j] = x[i, j] * x[i, j]
        return y

    output = cus_square_compute(x)

    attrs = {
        'enable_post_poly_loop_partition': False,
        'enable_double_buffer': False,
        'enable_feature_library': True,
        'RewriteVarTensorIdx': True
    }

    return output, attrs

更多文章请关注《万象专栏》