[Relay] Mutual Recursion Support

include/tvm/relay/transform.h

@@ -210,6 +210,16 @@ TVM_DLL Pass FastMath();
  */
 TVM_DLL Pass InferType();
 
+/*!
+ * \brief Infer the type of all functions in a module.
+ *
+ * This pass should be used when typechecking modules
+ * with mutually recursive functions.
+ *
+ * \return The pass.
+ */
+TVM_DLL Pass InferTypeAll();
+
 /*!
  * \brief Search and eliminate common subexpression. For example, if there are
  * two expressions evaluated to an identical value, a single variable is created
@@ -493,6 +503,15 @@ TVM_DLL Function UnCPS(const Function& f);
  */
 TVM_DLL Expr DeDup(const Expr& e);
 
+/*!
+ * \brief Deduplicate the bound type variables in the type.
+ *
+ * \param e the type (does not have to be typechecked).
+ *
+ * \return the deduplicated type.
+ */
+TVM_DLL Type DeDupType(const Type& e);
+
 }  // namespace relay
 }  // namespace tvm
 

python/tvm/ir/module.py

@@ -87,6 +87,10 @@ def _add(self, var, val, update=False):
                 var = _ty.GlobalTypeVar(var)
             _ffi_api.Module_AddDef(self, var, val, update)
 
+    def add_unchecked(self, var, val):
+        assert isinstance(val, _expr.RelayExpr)
+        _ffi_api.Module_AddUnchecked(self, var, val)
+
     def __getitem__(self, var):
         """Lookup a global definition by name or by variable.
 

python/tvm/relay/transform/transform.py

@@ -98,6 +98,8 @@ def InferType():
     """
     return _ffi_api.InferType()
 
+def InferTypeAll():
+    return _ffi_api.InferTypeAll()
 
 def FoldScaleAxis():
     """Fold the scaling of axis into weights of conv2d/dense. This pass will

src/ir/module.cc

@@ -403,6 +403,9 @@ TVM_REGISTER_GLOBAL("ir.Module_Add").set_body([](TVMArgs args, TVMRetValue* ret)
 
 TVM_REGISTER_GLOBAL("ir.Module_AddDef").set_body_method<IRModule>(&IRModuleNode::AddTypeDef);
 
+TVM_REGISTER_GLOBAL("ir.Module_AddUnchecked")
+    .set_body_method<IRModule>(&IRModuleNode::AddUnchecked);
+
 TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar")
     .set_body_method<IRModule>(&IRModuleNode::GetGlobalVar);
 

src/relay/analysis/type_solver.h

@@ -65,6 +65,9 @@ class TypeSolver {
  public:
   TypeSolver(const GlobalVar& current_func, const IRModule& _mod, ErrorReporter* err_reporter);
   ~TypeSolver();
+
+  void SetCurrentFunc(const GlobalVar& current_func) { this->current_func = current_func; }
+
   /*!
    * \brief Add a type constraint to the solver.
    * \param constraint The constraint to be added.

src/relay/transforms/de_duplicate.cc

@@ -99,6 +99,33 @@ Expr DeDup(const Expr& e) {
   return ret;
 }
 
+// dedup bound type variables in type
+// - types do not have to be already typechecked
+Type DeDupType(const Type& e) {
+  class DeDupTypeMutator : public TypeMutator, public ExprMutator, public PatternMutator {
+   public:
+    TypeVar Fresh(const TypeVar& tv) {
+      TypeVar ret = TypeVar(tv->name_hint, tv->kind);
+      type_rename_[tv] = ret;
+      return ret;
+    }
+
+    Type VisitType(const Type& t) final { return t.defined() ? TypeMutator::VisitType(t) : t; }
+
+    Pattern VisitPattern(const Pattern& p) final { return PatternFunctor::VisitPattern(p); }
+
+    Type VisitType_(const TypeVarNode* op) final {
+      TypeVar v = GetRef<TypeVar>(op);
+      return type_rename_.count(v) != 0 ? type_rename_.at(v) : Fresh(v);
+    }
+
+   private:
+    std::unordered_map<TypeVar, TypeVar, ObjectPtrHash, ObjectPtrEqual> type_rename_;
+  };
+
+  Type ret = DeDupTypeMutator().VisitType(e);
+  return ret;
+}
 TVM_REGISTER_GLOBAL("relay._transform.dedup").set_body_typed(DeDup);
 
 }  // namespace relay

src/relay/transforms/type_infer.cc

@@ -109,6 +109,44 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)>,
   // inference the type of expr.
   Expr Infer(Expr expr);
 
+  void SetCurrentFunc(GlobalVar current_func) {
+    this->current_func_ = current_func;
+    this->solver_.SetCurrentFunc(current_func);
+  }
+
+  void Solve();
+  Expr ResolveType(Expr expr);
+
+  // Lazily get type for expr
+  // expression, we will populate it now, and return the result.
+  Type GetType(const Expr& expr) {
+    auto it = type_map_.find(expr);
+    if (it != type_map_.end() && it->second.checked_type.defined()) {
+      if (expr.as<GlobalVarNode>() != nullptr) {
+        // if we don't dedup GlobalVarNode, two functions that use the same GlobalVar
+        // may resolve to the same type incorrectly
+        return DeDupType(it->second.checked_type);
+      }
+      return it->second.checked_type;
+    }
+    Type ret = this->VisitExpr(expr);
+    CHECK(ret.defined());
+    KindCheck(ret, mod_);
+    ResolvedTypeInfo& rti = type_map_[expr];
+    rti.checked_type = ret;
+    return ret;
+  }
+
+  // Lazily get type for GlobalVar
+  // expression, we will populate it now, and return the result.
+  Type GetTypeGlobalVar(const GlobalVar& expr) {
+    // we have to visit functiion
+    // or else it may not be type-checked
+    auto f = Downcast<Function>(mod_->Lookup(expr));
+    Type ret = GetType(f);
+    return GetType(expr);
+  }
+
  private:
   // type resolver that maps back to type
   class Resolver;
@@ -143,21 +181,6 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)>,
     }
   }
 
-  // Lazily get type for expr
-  // expression, we will populate it now, and return the result.
-  Type GetType(const Expr& expr) {
-    auto it = type_map_.find(expr);
-    if (it != type_map_.end() && it->second.checked_type.defined()) {
-      return it->second.checked_type;
-    }
-    Type ret = this->VisitExpr(expr);
-    CHECK(ret.defined());
-    KindCheck(ret, mod_);
-    ResolvedTypeInfo& rti = type_map_[expr];
-    rti.checked_type = ret;
-    return ret;
-  }
-
   void ReportFatalError(const ObjectRef& expr, const Error& err) {
     CHECK(this->current_func_.defined());
     this->err_reporter.ReportAt(this->current_func_, expr, err);
@@ -303,7 +326,6 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)>,
         this->ReportFatalError(match, ss);
       }
     }
-
     return rtype;
   }
 
@@ -543,14 +565,6 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)>,
     }
     return FuncType(c->inputs, TypeCall(c->belong_to, types), td->type_vars, {});
   }
-
-  void Solve() {
-    solver_.Solve();
-
-    if (err_reporter.AnyErrors()) {
-      err_reporter.RenderErrors(mod_);
-    }
-  }
 };
 
 class TypeInferencer::Resolver : public ExprMutator, PatternMutator {
@@ -698,6 +712,20 @@ Expr TypeInferencer::Infer(Expr expr) {
   return resolved_expr;
 }
 
+void TypeInferencer::Solve() {
+  solver_.Solve();
+
+  if (err_reporter.AnyErrors()) {
+    err_reporter.RenderErrors(mod_);
+  }
+}
+
+Expr TypeInferencer::ResolveType(Expr expr) {
+  auto resolved_expr = Resolver(type_map_, &solver_).VisitExpr(expr);
+  CHECK(WellFormed(resolved_expr));
+  return resolved_expr;
+}
+
 struct AllCheckTypePopulated : ExprVisitor {
   void VisitExpr(const Expr& e) {
     if (e.as<OpNode>()) {
@@ -742,6 +770,43 @@ Function InferType(const Function& func, const IRModule& mod, const GlobalVar& v
   return Downcast<Function>(func_ret);
 }
 
+IRModule InferTypeAll(const IRModule& mod) {
+  CHECK(mod.defined()) << "internal error: module must be set for type inference";
+  const auto& globalvars = mod->GetGlobalVars();
+
+  // first pass: fill in type for all functions
+  for (const auto& var : globalvars) {
+    relay::Function func = Downcast<relay::Function>(mod->Lookup(var));
+    func->checked_type_ = func->func_type_annotation();
+    mod->AddUnchecked(var, func);
+  }
+
+  // use dummy var, will be updated as we fill constraints/
+  // solve for each GlobalVar
+  TypeInferencer ti = TypeInferencer(mod, GlobalVar("dummy"));
+
+  // second pass, fill in constraints
+  for (const auto& var : globalvars) {
+    ti.SetCurrentFunc(var);
+    ti.GetTypeGlobalVar(var);
+  }
+
+  // solve constraints
+  ti.Solve();
+
+  // third pass, resolve types
+  for (const auto& var : globalvars) {
+    ti.SetCurrentFunc(var);
+
+    relay::Function func = Downcast<relay::Function>(mod->Lookup(var));
+    Expr func_ret = ti.ResolveType(func);
+    // add function back to module
+    mod->AddUnchecked(var, Downcast<relay::Function>(func_ret));
+  }
+
+  return mod;
+}
+
 namespace transform {
 
 Pass InferType() {
@@ -750,8 +815,18 @@ Pass InferType() {
   return CreateFunctionPass(pass_func, 0, "InferType", {});
 }
 
+Pass InferTypeAll() {
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
+      [=](IRModule m, PassContext pc) { return relay::InferTypeAll(m); };
+  return CreateModulePass(pass_func, 0, "InferTypeAll", {});
+}
+
 TVM_REGISTER_GLOBAL("relay._transform.InferType").set_body_typed([]() { return InferType(); });
 
+TVM_REGISTER_GLOBAL("relay._transform.InferTypeAll").set_body_typed([]() {
+  return InferTypeAll();
+});
+
 }  // namespace transform
 
 }  // namespace relay