帶你了解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

然后在同級目錄下的__init__.py文件中添加內容

from .cus_square import CusSquare

2.2 注冊算子

到路徑：mindspore/ops/_op_impl/akg/ascend，創建cus_square.py，添加如下代碼

"""CusSquare op"""
from mindspore.ops.op_info_register import op_info_register, AkgAscendRegOp, DataType as DT

op_info = AkgAscendRegOp("CusSquare") \
    .fusion_type("ELEMWISE") \
    .input(0, "x") \
    .output(0, "output") \
    .dtype_format(DT.F32_Default, DT.F32_Default) \
    .get_op_info()
@op_info_register(op_info)
def _cus_square_akg():
    """CusSquare Akg register"""
    return

然后在同級目錄的__init__.py添加如下代碼

from .cus_square import _cus_square_akg

2.3 定義算子原語

到：mindspore/ops/operations，新創建一個_cus_ops.py，添加如下代碼

描述算子的輸入：x，輸出output

infer_shape：描述輸出數據的shape

infer_dtype：說明輸出數據的類型

x1_shape：指的是第一個輸入的shape

x1_dtype：指的是第一個輸入參數的dtype

import math

from ..primitive import prim_attr_register, PrimitiveWithInfer
from ...common import dtype as mstype
from ..._checkparam import Validator as validator
from ..._checkparam import Rel

class CusSquare(PrimitiveWithInfer):
    """CusSquare"""

    @prim_attr_register
    def __init__(self):
        self.init_prim_io_names(inputs=['x'], outputs=['output'])

    def infer_shape(self, x1_shape):
        return x1_shape

    def infer_dtype(self, x1_dtype):
        return x1_dtype

然后在同目錄下的__init__.py文件中添加原語信息

from ._cus_ops import CusSquare

2.4 在ccsrc中添加算子的查詢信息

在mindspore/ccsrc/backend/kernel_compiler/http://kernel_query.cc的KernelQuery函數中添加如下信息

// cus_square
const PrimitivePtr kPrimCusSquare = std::make_shared<Primitive>("CusSquare");
if (IsPrimitiveCNode(kernel_node, kPrimCusSquare)) {
    kernel_type = KernelType::AKG_KERNEL;
}

2.5 編譯安裝框架

回到mindspore根目錄

bash build.sh -e ascend -j4  
cd ./build/package
pip install mindspore_ascend-1.1.2-cp37-cp37m-linux_aarch64.whl --force-reinstall

2.6 測試

import numpy as np
import mindspore.nn as nn
import mindspore.context as context
from mindspore import Tensor
from mindspore.ops import operations as P

context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.square = P.CusSquare()

    def construct(self, data):
        return self.square(data)

def test_net():
    x = np.array([[1.0, 4.0, 9.0]]).astype(np.float32)
    net = Net()
    output = net(Tensor(x))
    print("x: ", x)
    print("output: ", output)
if __name__ == "__main__":
    test_net()

輸出

三、反向算子的制作和測試

3.1 制作流程

反向算子的計算邏輯：對向量元素進行求導，如 y = x^2，則求導之后 y` = 2x

實際例子就是輸入向量[1, 4, 9] 輸出就是 [2, 8, 18]

反向算子明明為CusSquareGrad，與前邊的計算平方的算子流程相同，這里只貼一下關鍵代碼，流程不再贅述

計算邏輯代碼cus_square_grad.py

"""cus_square_grad"""
from akg.tvm.hybrid import script
import akg

def CusSquareGrad(x):
    output_shape = x.shape
    k = output_shape[0]
    n = output_shape[1]

    @script
    def cus_square_compute_grad(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] * 2
        return y

    output = cus_square_compute_grad(x)

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

    return output, attrs

注冊原語

class CusSquareGrad(PrimitiveWithInfer):
    """
    CusSquareGrad
    """

    @prim_attr_register
    def __init__(self):
        self.init_prim_io_names(inputs=['x'], outputs=['output'])

    def infer_shape(self, x1_shape):
        return x1_shape

    def infer_dtype(self, x1_dtype):
        return x1_dtype  

3.2 測試

import numpy as np
import mindspore.nn as nn
import mindspore.context as context
from mindspore import Tensor
from mindspore.ops import operations as P

context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.square = P.CusSquareGrad()   # 替換為grad算子

    def construct(self, data):
        return self.square(data)

def test_net():
    x = np.array([[1.0, 4.0, 9.0]]).astype(np.float32)
    net = Net()
    output = net(Tensor(x))
    print("x: ", x)
    print("output: ", output)
if __name__ == "__main__":
    test_net()

輸出

四、正反向算子關聯及測試

在源碼 mindspore/mindspore/ops/_grad/grad_array_ops.py中添加如下代碼

@bprop_getters.register(P.CusSquare)
def get_bprop_cussquare(self):
    """Generate bprop of CusSquare"""
    cus_square_grad = P.CusSquareGrad()
    matmul = ops.Mul()
    def bprop(x, out, dout):
        gradient = cus_square_grad(x)
        dx = matmul(gradient, dout)
        return (dx,)
    return bprop

bprop函數的輸入是，正向的輸入x，正向的輸出out，反向的梯度輸入dout

上面代碼的意思是指定算子CusSquare的反向梯度的計算方法，CusSquareGrad作為其中的一個函數使用

gradient = cus_square_grad(x)計算的是本平方算子的梯度，但並不能直接返回這個梯度

反向網絡到該算子，最后返回的是dx，注意算子的反向梯度計算一定要放在整個網絡的反向鏈式梯度計算中

測試

import numpy as np
import mindspore.nn as nn
import mindspore.context as context
from mindspore import Tensor
from mindspore.ops import operations as P
from mindspore.ops import composite as C

context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.square = P.CusSquare()  

    def construct(self, data):
        return self.square(data)

def test_net():
    x = Tensor(np.array([[1.0, 4.0, 9.0]]).astype(np.float32))
    grad = C.GradOperation(get_all=True)    # 計算網絡梯度
    net = Net()
    output = grad(net)(x)
    print("x: ", x)
    print("output: ", output)
if __name__ == "__main__":
    test_net()

輸出

 

點擊關注，第一時間了解華為雲新鮮技術~