自定义算子高级用法
算子信息注册
算子信息主要描述了算子实现函数所支持的输入输出类型、输入输出数据格式、属性和target（平台信息），它是后端做算子选择和映射时的依据。它通过CustomRegOp接口定义，通过custom_info_register装饰器或者Custom原语构造函数中的reg_info参数，实现算子信息与算子实现函数的绑定，并最终注册到MindSpore
C++侧的算子信息库。reg_info参数优先级高于custom_info_register装饰器。
算子信息中的target的值可以为“Ascend”或“GPU”或“CPU”，描述的是算子实现函数在当前target上所支持的输入输出类型、输入输出数据格式和属性等信息。对于同一个算子实现函数，其在不同target上支持的输入输出类型可能不一致，但算子信息在同一target下只会被注册一次，所以可以通过target进行区分。
算子信息中定义输入输出信息的个数和顺序、算子实现函数中的输入输出信息的个数和顺序，两者要完全一致。
对于akg类型的自定义算子，若算子存在属性输入，则必须注册算子信息，算子信息中的属性名称与算子实现函数中使用的属性名称要一致；对于tbe类型的自定义算子，当前必须注册算子信息；对于aot类型的自定义算子，由于算子实现函数需要预先编译成动态库，所以无法通过装饰器方式绑定算子信息，只能通过
reg_info参数传入算子信息。若自定义算子只支持特定的输入输出数据类型或数据格式，则需要注册算子信息，以便在后端做算子选择时进行数据类型和数据格式的检查。对于不提供算子信息的情况，则在后端做算子选择和映射的时候，将会从当前算子的输入中推导信息。
定义算子反向传播函数
如果算子要支持自动微分，需要定义其反向传播函数（bprop），然后将bprop函数传入Custom原语构造函数的bprop参数。你需要在bprop中描述利用正向输入、正向输出和输出梯度得到输入梯度的反向计算逻辑。反向计算逻辑可以使用内置算子或自定义Custom算子。
定义算子反向传播函数时需注意以下几点：
bprop函数的入参顺序约定为正向的输入、正向的输出、输出梯度。若算子为多输出算子，正向输出和输出梯度将以元组的形式提供。
bprop函数的返回值形式约定为输入梯度组成的元组，元组中元素的顺序与正向输入参数顺序一致。即使只有一个输入梯度，返回值也要求是元组的形式。
下面test_grad.py为例，展示反向传播函数的用法：
[1]:
import numpy as np
import mindspore as ms
from mindspore.nn import Cell
import mindspore.ops as ops
ms.set_context(mode=ms.GRAPH_MODE)
ms.set_device(device_target="CPU")
# 自定义算子正向实现
def square(x):
return x * x
# 自定义算子反向实现
def square_grad(x, dout):
dx = 2.0 * x * dout
return dx
# 反向传播函数
def bprop():
op = ops.Custom(square_grad, lambda x, _: x, lambda x, _: x, func_type="pyfunc")
def custom_bprop(x, out, dout):
dx = op(x, dout)
return (dx,)
return custom_bprop
class Net(Cell):
def __init__(self):
super(Net, self).__init__()
# 定义 pyfunc 类型的自定义算子，并提供反向传播函数
self.op = ops.Custom(square, lambda x: x, lambda x: x, bprop=bprop(), func_type="pyfunc")
def construct(self, x):
return self.op(x)
if __name__ == "__main__":
x = np.array([1.0, 4.0, 9.0]).astype(np.float32)
dx = ms.grad(Net())(ms.Tensor(x))
print(dx)
[ 2. 8. 18.]
其中：
反向传播函数中使用是的akg类型的自定义算子，算子定义与使用需要分开，即自定义算子在
custom_bprop函数外面定义，在custom_bprop函数内部使用。
执行用例：
python test_grad.py
执行结果：
[ 2. 8. 18.]
更多示例可参考MindSpore源码中tests/st/graph_kernel/custom下的用例。
