向Relay添加算子
向Relay添加算子
为了在Relay IR中使用TVM算子，需要在Relay中注册算子，以确保将其集成到Relay的类型系统中。
注册算子需要三个步骤：
- 使用RELAY_REGISTER_OPC ++中的宏注册算子的Arity和类型信息
- 定义一个C ++函数为算子生成一个调用节点，并为该函数注册一个Python API挂钩
- 将上述Python API挂钩包装在更整洁的界面中
该文件src/relay/op/tensor/binary.cc提供了前两个步骤的python/tvm/relay/op/tensor.py示例，同时提供了后两个步骤的示例。
注册算子
TVM已经具有算子注册表，如果没有其它类型信息，Relay无法正确合并TVM算子。
为了在注册算子时具有更大的灵活性，并在Relay中表达类型时，提高了表达性和粒度，使用输入和输出类型之间的关系来对算子进行类型化。这些关系表示为接受输入类型和输出类型的列表（这些类型中的任何一个都不完整），返回满足该关系的输入和输出类型的列表的函数。本质上，算子的关系除了计算输出类型外，还可以强制执行所有必要的键入规则（即，通过检查输入类型）。
例如，参阅src/relay/op/type_relations.h及其实现。例如，BroadcastRel接受两个输入类型和一个输出类型，都是具有相同基础数据类型的张量类型，最后确保输出类型的形状是输入类型的形状的广播。
type_relations.h 如果现有类型未捕获所需算子的行为，则可能有必要添加另一种类型关系。
使用RELAY_REGISTER_OPC ++中的宏，开发人员可以在Relay中指定有关算子的以下信息：
- Arity（参数个数）
- 位置参数的名称和说明
- 支持级别（1表示内部固有；较高的数字表示积分较少或不受外部支持的算子）
- 算子的类型关系
下面的示例来自binary.cc张量，并将其用于张量广播：
RELAY_REGISTER_OP("add")
.set_num_inputs(2)
.add_argument("lhs", "Tensor", "The left hand side tensor.")
.add_argument("rhs", "Tensor", "The right hand side tensor.")
.set_support_level(1)
.add_type_rel("Broadcast", BroadcastRel);
创建调用节点
此步骤仅需要简单地编写一个将参数带给算子的函数（如Relay表达式），然后将调用节点返回给算子（即，应将其放置在Relay AST中的节点，该AST是要向算子进行调用的位置）。
目前不支持调用属性和类型参数（最后两个字段），足以用于Op::Get从算子注册表中获取算子信息，并将参数传递给调用节点，如下所示。
TVM_REGISTER_GLOBAL("relay.op._make.add")
.set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) {
static const Op& op = Op::Get("add");
return Call(op, {lhs, rhs}, Attrs(), {});
});
包括Python API挂钩
在Relay中通常是约定，通过导出的函数TVM_REGISTER_GLOBAL应该包装在单独的Python函数中，而不是在Python中直接调用。对于产生对算子的调用的函数，捆绑起来可能很方便，其中python/tvm/relay/op/tensor.py，都提供了张量上的元素算子。例如，以下是上一节中的add函数在Python中的显示方式：
def add(lhs, rhs):
"""Elementwise addition.
Parameters
----------
lhs : relay.Expr
The left hand side input data
rhs : relay.Expr
The right hand side input data
Returns
-------
result : relay.Expr
The computed result.
"""
return _make.add(lhs, rhs)
这些Python库也可能是向算子提供更简单接口的好机会。例如，该 concat算子被注册为仅接受一个算子，即一个具有要连接的张量的元组，Python库将这些张量作为参数，组合成一个元组，然后生成调用节点：
def concat(*args):
"""Concatenate the input tensors along the zero axis.
Parameters
----------
args: list of Tensor
Returns
-------
tensor: The concatenated tensor.
"""
tup = Tuple(list(args))
return _make.concat(tup)
梯度算子
梯度算子对于在Relay中编写可区分的程序很重要。尽管Relay的autodiff算法可以区分一流的语言结构，算子是不透明的。由于Relay无法调查实现，必须提供明确的区分规则。
Python和C ++均可用于编写梯度算子，将示例重点放在Python上，因为更常用。
在Python中添加渐变
可以找到Python梯度算子的集合 python/tvm/relay/op/_tensor_grad.py。将通过两个有代表性的示例：sigmoid和multiply。
@register_gradient("sigmoid")
def sigmoid_grad(orig, grad):
"""Returns [grad * sigmoid(x) * (1 - sigmoid(x))]."""
return [grad * orig * (ones_like(orig) - orig)]
这里的输入是原始算子orig和grad要累加到的渐变。返回的是一个列表，其中第i个索引处的元素是算子相对于算子第i个输入的派生。通常，渐变将返回一个列表，其中包含与基本算子输入相同数量的元素。
在进一步分析这个定义之前，应该回想一下S型函数的导数： ∂σ∂x=σ(x)(1−σ(x))∂σ∂x=σ(x)(1−σ(x))。上面的定义看起来与数学定义相似，有一个重要的补充，将在下面进行描述。
该术语直接与导数匹配，因为这是S型函数，不仅对如何计算此函数的梯度感兴趣。有兴趣将此梯度与其它梯度组成，可以在整个程序中累积该梯度。这是该术语出现的地方。指定到目前为止如何用梯度组成导数。orig * (ones_like(orig) - orig)origgradgrad * orig * (ones_like(orig) - orig)grad
来看multiply一个更有趣的示例：
@register_gradient("multiply")
def multiply_grad(orig, grad):
"""Returns [grad * y, grad * x]"""
x, y = orig.args
return [collapse_sum_like(grad * y, x),
collapse_sum_like(grad * x, y)]
在此示例中，返回列表中有两个元素， multiply是一个二进制算子。如果f(x,y)=xyf(x,y)=xy，偏导数是 ∂f∂x=y∂f∂x=y 和 ∂f∂y=x∂f∂y=x。
有一个必需的步骤，因为广播具有语义，multiply不需要。可能与输入的形状不匹配，习惯于获取术语的内容，并使形状与要区分的输入的形状相匹配。sigmoidmultiplygradcollapse_sum_likegrad * <var>
在C ++中添加渐变
在C ++中添加渐变类似于在Python中添加渐变，注册的界面略有不同。
确保src/relay/pass/pattern_utils.h包含其中。提供了用于在Relay AST中创建节点的辅助功能。然后，以类似于Python示例的方式定义渐变：
tvm::Array<Expr> MultiplyGrad(const Expr& orig_call, const Expr& output_grad) {
const Call& call = orig_call.Downcast<Call>();
return { CollapseSumLike(Multiply(output_grad, call.args[1]), call.args[0]),
CollapseSumLike(Multiply(output_grad, call.args[0]), call.args[1]) };
}
在C ++中，不能使用与Python中相同的算子重载，并且需要向下转换，实现更为冗长。即使这样，仍可以轻松地验证此定义，是否与Python中的先前示例相同。
代替使用Python渲染，需要set_attr在基本算子注册的末尾添加对“ FPrimalGradient”的调用，以注册渐变。
RELAY_REGISTER_OP("multiply")
// ...
// Set other attributes
// ...
.set_attr<FPrimalGradient>("FPrimalGradient", MultiplyGrad);
总结
- TVM算子可以使用表示适当类型信息的关系在中继中注册。
- 在中继中使用算子需要一个函数来为算子生成调用节点。
- 最好有一个简单的Python库来生成调用节点。
