Let Binding 这个概念也来自函数式编程,为了理解引入 Let Binding 的原因,我们先看下面这个例子
这是一个经典的计算图表达方式,唯一与传统不太一致的,就是使用 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, 看下面这个例子
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 的例子
如果看文本格式,我们会觉得 %1 应该在 if 表达式外计算,但看 AST/IR, 我们又迷惑了,在下面这个例子中,这个迷惑就更明显了
fn (%x) { %1 = log(%x) %2 = fn(%y) { add(%y, %1) } %2 }
我们应该在闭包的外面还是里面做 %1 的计算呢?而 Let Binding 解决了这个问题,我们只要加上 Let 就能解决这个歧义。