向Relay添加算子(operator)

为了在Relay IR中使用TVM算子，需要在Relay中注册算子，以确保将其集成到Relay的类型系统中。

注册算子**需要三个步骤：

  • 在C++中使用宏RELAY_REGISTER_OP注册算子的arity(参数个数)和类型信息
  • 定义一个C++函数为算子生成一个调用节点CallNode，并为该函数注册一个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及其实现(src/relay/op/type_relations.cc)。例如，BroadcastRel接受两个输入类型和一个输出类型，检查它们都是具有相同基础数据类型的张量类型，最后确保输出类型的shape是输入类型shape的广播。

  1. bool BroadcastRel(const Array<Type>& types,
  2.                   int num_inputs,
  3.                   const Attrs& attrs,
  4.                   const TypeReporter& reporter) {
  5.   CHECK_EQ(types.size(), 3);//2个输入类型和1个输出类型
  6.   // DLOG(INFO) << "In1:" << types[0] << ",In2:" << types[1]
  7.   //                 << ",Out:" << types[2] << std::endl;
  8.   if (auto t0 = ToTensorType(types[0])) {
  9.     if (auto t1 = ToTensorType(types[1])) {
  10.       CHECK_EQ(t0->dtype, t1->dtype);
  11.       reporter->Assign(types[2],
  12.         ConcreteBroadcast(t0, t1, t0->dtype));
  13.       return true;
  14.     }
  15.   }
  16.   return false;
  17. }

如果现有的类型关系无法捕获所需算子的行为，则可能需要添加其他类型关系到type_relations.h。
C++中的RELAY_REGISTER_OP宏允许开发人员指定以下有关Relay中算子的信息：

  • Arity（参数个数）
  • 位置参数的名称和描述
  • 支持级别（1表示内部固有；较大的数字表示不够完整或外部支持的算子）
  • 算子的类型关系

下面的示例来自binary.cc张量，并将其用于张量：

  1. RELAY_REGISTER_OP("add")
  2.     .set_num_inputs(2) //参数个数
  3.     .add_argument("lhs", "Tensor", "The left hand side tensor.")
  4.     .add_argument("rhs", "Tensor", "The right hand side tensor.")
  5.     .set_support_level(1) //支持级别
  6.     .add_type_rel("Broadcast", BroadcastRel); // 类型关系

创建一个调用节点(Call Node)

此步骤仅需要简单地编写一个将参数带给算子的函数（作为Relay表达式），并将CallNode返回给算子（即，应将其放置在打算调用算子的Relay AST中的节点）。
当前不支持调用属性和类型参数（最后两个字段），因此足以用于Op::Get从算子注册表中获取算子信息，并将参数传递给调用节点，如下所示。

  1. TVM_REGISTER_API("relay.op._make.add")
  2.     .set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) {
  3.         static const Op &op = Op::Get("add");
  4.         return CallNode::make(op, {lhs, rhs}, Attrs(), {});
  5.     });

引用一个Python API钩子

在Relay中通常习惯这么做：通过TVM_REGISTER_API导出的函数应该封装在单独的Python函数中，而不是在Python中直接调用。对于产生对算子的调用的函数，将它们绑定起来可能很方便，如python/tvm/relay/op/tensor.py中所示，其中都提供了张量上的元素算子。例如，以下是上一节的add函数在Python中的公开方式：

  1. def add(lhs, rhs):
  2.     """Elementwise addition.
  3.     Parameters
  4.     ----------
  5.     lhs : relay.Expr
  6.     The left hand side input data
  7.     rhs : relay.Expr
  8.     The right hand side input data
  9.     Returns
  10.     -------
  11.     result : relay.Expr
  12.     The computed result.
  13.     """
  14.     return _make.add(lhs, rhs)

请注意，这些Python包装器也可能是向算子提供更简单接口的好机会。例如，concat算子被注册为仅接受一个参数，即一个由多个张量组合起来的元组，Python包装器将这些张量作为参数并将它们组合成一个元组，然后生成调用节点：

  1. def concat(*args):
  2.     """Concatenate the input tensors along the zero axis.
  3.     Parameters
  4.     ----------
  5.     args: list of Tensor
  6.     Returns
  7.     -------
  8.     tensor: The concatenated tensor.
  9.     """
  10.     tup = Tuple(list(args))
  11.     return _make.concat(tup)

**

梯度算子

梯度算子对于在Relay中编写可区分的程序很重要。尽管Relay的autodiff算法可以区分一流的语言结构，但算子是不透明的。由于Relay无法调查实现，因此必须提供明确的区分规则。
Python和C++均可用于编写梯度算子，但是我们将示例重点放在Python上，因为它更常用。

在Python中添加梯度

可以在python/tvm/relay/op/tensor_grad.py中找到Python梯度算子的集合 。我们将通过两个有代表性的示例：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。返回的是一个列表list，其中索引i的元素是算子相对于第i个输入的导数。通常，梯度将返回一个列表，其中包含与基本算子输入相同数量的元素。
在进一步分析此定义之前，首先我们应该回顾一下 sigmoid 型函数的导数：
[翻译]向Relay添加算子(operator)— tvm 0.6.0文档 - 图1
∂σ∂x=σ(x)(1−σ(x))
上面的定义看起来与数学定义相似，但是有一个重要的补充，我们将在下面进行描述。
术语orig (ones_like(orig) - orig)直接与导数匹配，因为orig是sigmoid 函数，但我们不仅对如何计算此函数的梯度感兴趣。我们有兴趣将此梯度与其他梯度组合，因此我们可以在整个程序中累积该梯度。这是术语grad出现的地方。在表达式grad 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)=xy，偏导数是
[翻译]向Relay添加算子(operator)— tvm 0.6.0文档 - 图2
multiply 有一个必需的步骤，但不是 sigmoid的必要步骤。因为multiply具有广播语义，因此multiply不需要 。由于grad的形状可能与输入的形状不匹配，因此我们习惯用collapse_sum_like获取术语grad * 的内容，并使形状与我们要区分的输入的形状相匹配。

在C++中添加梯度

在C++中添加梯度类似于在Python中添加梯度，但是注册的接口略有不同。
首先，请确保引用src/relay/pass/pattern_util.h。它提供了用于在Relay AST中创建节点的辅助功能。然后，以类似于Python示例的方式定义梯度：
tvm::Array MultiplyGrad(const Expr& origcall, const Expr& output_grad) {
const Call& call = orig_call.Downcast();
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”, MultiplyGrad);

总结

  • TVM算子可以使用表示合适类型信息的关系在relay中注册。
  • 在relay中使用算子需要一个函数来为生成Call Node。
  • 最好有一个简单的Python包装器来生成Call Node。