向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库来生成调用节点。