深入浅出 tvm – (8) Relay IR 之 Let binding

Let Binding 这个概念也来自函数式编程,为了理解引入 Let Binding 的原因,我们先看下面这个例子
0
这是一个经典的计算图表达方式,唯一与传统不太一致的,就是使用 Function 来代替 Graph 的概念。在这个例子中,需要注意的是 add 的两个参数都是 1%,在 DAG 中出现了一个共享节点,我们假设 add 表达式的代码如下
class Add : public Expr {
  Expr lhs;
  Expr rhs;
};
在上面的例子中,也就意味着 lhs 成员和 rhs 成员指向了同一个表达式。在传统的计算图模式中,我们会按照拓扑序求值,1% 只会计算一次。但是从做编译的角度来说,实现的 codegen 大概率是这样
auto Add::Codegen() {
  auto lhs_val = lhs->Codegen();
  auto rhs_val = rhs->Codegen();
  return add(lhs_val, rhs_val);
}
也就意味着 log(%x)+log(%x) ,所以这里就出现了歧义,而为了能够保持和传统计算图的兼容,TVM 在打印 IR 到文本格式时,会为每个节点打印一行,并分配一个临时 id (也就是上面的%1, %2),所以从文本格式看可能发现不了什么。为了解决歧义,TVM 引入了 Let Binding, 看下面这个例子
0
Let 表达式为 Let (var, value, body) , 其将value求值赋给var,然后返回body的求值结果,其对应的 IR/AST 称为 A-normal form,如上图,上面那个 Let 将 log(%x) 作为 value 求值赋给其 var(也就是 %v1),被 add 引用。通过这种方式,解决了歧义问题。
更简单地说,Let Binding 将表达式 Expr 绑定到局部的不可变 var 中。从代码实现角度,如果将传统计算图的方式称为 Graph Binding,与 Let Binding 对立的话。在 Graph Binding中,add 的 lhs 和 rhs 都直接指向了 log(%x) 这个计算表达式 Expr,而在 Let Binding 中,add 的 lhs 和 rhs 都指向了一个局部变量表达式,而这个局部变量表达式的值,由 Let 表达式负责赋值。我们可以看 Let 表达式的代码
class LetNode : public ExprNode {
 protected:
  // LetNode uses own deleter to indirectly call non-recursive destructor
  Object::FDeleter saved_deleter_;
  static void Deleter_(Object* ptr);

 public:
  /*! \brief The variable we bind to */
  Var var;
  /*! \brief The value we bind var to */
  Expr value;
  /*! \brief The body of the let binding */
  Expr body;
事实上,Let Binding 的一个核心用途就是指定计算的Scope,我们看下面这个没有使用 Let Binding 的例子
0
如果看文本格式,我们会觉得 %1 应该在 if 表达式外计算,但看 AST/IR, 我们又迷惑了,在下面这个例子中,这个迷惑就更明显了
fn (%x) {
  %1 = log(%x)
  %2 = fn(%y) {
    add(%y, %1)
  }
  %2
}
我们应该在闭包的外面还是里面做 %1 的计算呢?而 Let Binding 解决了这个问题,我们只要加上 Let 就能解决这个歧义。