TVM Pass优化 -- InferType 类型推导

定义（What）

InferType，类型推断，顾名思义，给表达式进行类型的推断

直接上代码

import tvm
from tvm import relay
import numpy as np

def get_demo_mod():
    a = relay.var("a", shape=(2, 3, 10), dtype="float32")
    b = relay.var("b", shape=(1, 10), dtype="float32")
    c = relay.add(a, b)
    func = relay.Function([a, b], c)
    mod = tvm.IRModule.from_expr(func)
    return mod

mod = get_demo_mod()

print("------before InferType------")
try:
    print(mod["main"].body.checked_type)
except Exception:
    print("can't get checked_type")

print("------after InferType------")

mod = relay.transform.InferType()(mod)
print(mod["main"].body.checked_type)

执行结果如下：

作用 （Why）

推断表达式的类型及输入输出尺寸

另：在 Relay 优化过程中, 每个 pass 都可以修改/添加/删除 op, 所以每个 pass 之后都需要重新 InferType

如，TVM Pass优化 -- 公共子表达式消除(Common Subexpr Elimination, CSE)对公共子表达式消除一节中FunctionPass()第四个参数就是InferType进行类型推断

怎么做（How）

这块代码主要在src/relay/transforms/type_infer.cc文件中，具体实现如下：

Pass InferType() {
  auto pass_info = PassInfo(0, "InferType", {});
  return tvm::transform::CreateModulePass(
      [=](IRModule mod, const PassContext& pass_ctx) {
	...
        AddGlobalTypes(mod);
        VLOG(1) << "AddGlobalTypes'" << PrettyPrint(mod);
        std::vector<std::pair<GlobalVar, Function>> updates;
        for (const auto& it : updated_mod->functions) {
          if (auto func = it.second.as<Function>()) {
            auto inferencer = TypeInferencer(mod, pass_ctx->diag_ctx.value());
            VLOG(1) << "it.first'" << PrettyPrint(it.first) << "it.second"<< PrettyPrint(it.second);

            auto updated_func = inferencer.Infer(it.first, func.value());
            VLOG(1) << "updated_func'" << PrettyPrint(updated_func);
      		...
            it.first->checked_type_ = updated_func->checked_type();

            if (!WellFormed(updated_func, pass_ctx->diag_ctx)) {
              LOG(FATAL) << "The type checked intermediate representation is malformed";
            }

            auto free_tvars = FreeTypeVars(updated_func, mod);
            ICHECK(free_tvars.size() == 0)
                << "Found unbound type variables in " << updated_func << ": " << free_tvars;
            EnsureCheckedType(updated_func);
            updates.push_back({it.first, Downcast<Function>(updated_func)});
          }
        }

        for (const auto& pair : updates) {
          updated_mod->Add(pair.first, pair.second, true);
        }

        return updated_mod;
      },
      0, "InferType", {});
}

TVM_REGISTER_GLOBAL("relay._transform.InferType").set_body_typed([]() { return InferType(); });

和公共子表达式消除的实现可发现，该算子调用的是CreateModulePass，因此它是一个模块级的优化，

模块级优化用于实现过程间优化和分析，模块级优化pass工作在tvm.IRModule对象上，将整个程序作为处理单元，几乎可以对程序执行任何操作。

其中，AddGlobalTypes 给mod添加全局参数，为后续的参数推断做准备，

真正进行推断的是TypeInferencer类的Infer()方法，实现如下：

Expr TypeInferencer::Infer(GlobalVar var, Function function) {
	...
  // Step 1: Populate the constraints.
  GetType(function);

  // Step 2: Solve the constraints.
  Solve();

  // Step 3: Attach resolved types to checked_type field.
  auto resolved_expr = Resolver(type_map_, &solver_).VisitExpr(function);
	...
  }
  return resolved_expr;
}

第一步，填充约束

  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);
    ICHECK(ret.defined()) << "expression:" << std::endl << PrettyPrint(expr);
    KindCheck(ret, mod_, this->diag_ctx);
    ResolvedTypeInfo& rti = type_map_[expr];
    rti.checked_type = ret;
    return ret;
  }

会先从type_map_map表中查找该Expr,第一次执行，如果type_map_中未找到该expr，便会通过VisitExpr()方法在该map表中添加，具体实现如下：

  void VisitLeaf(const Expr& expr) {
    if (!memo_.count(expr)) {
      Type ret = this->DispatchVisitExpr(expr);
      memo_[expr] = ret;
    }
  }

  bool CheckVisited(const Expr& expr) {
    if (memo_.count(expr)) {
      return true;
    } else {
      return false;
    }
  }

  Type DispatchVisitExpr(const Expr& expr) { return ExprFunctor::VisitExpr(expr); }

  Type VisitExpr(const Expr& expr) final {
    auto fcheck_visited = [this](const Expr& expr) { return this->CheckVisited(expr); };
    auto fvisit_leaf = [this](const Expr& expr) { return this->VisitLeaf(expr); };
    if (memo_.count(expr)) {
      return memo_[expr];
    } else {
      ExpandDataflow(expr, fcheck_visited, fvisit_leaf);
      return memo_[expr];
    }
  }

其中fcheck_visited()匿名函数通过调用VisitLeaf方法中的DispatchVisitExpr方法，该函数会调用到ExprFunctor类中构建的包含各种类型的虚表中，根据类型调用对应的VisitExpr_方法，如CallNode类型的参数，代码如下：

  Type VisitExpr_(const CallNode* call) final {
    Array<Type> arg_types;
    for (Expr arg : call->args) {
      arg_types.push_back(GetType(arg));
    }

    if (const OpNode* opnode = call->op.as<OpNode>()) {
      Type rtype =
          PrimitiveCall(opnode->op_type.as<FuncTypeNode>(), arg_types, call->attrs, call->span);

      if (rtype.defined()) {
        AddTypeArgs(GetRef<Call>(call), arg_types);
        return rtype;
      }
    }

其中，AddTypeArgs()会向type_map_表中插入该expr

  void AddTypeArgs(const Expr& expr, Array<Type> type_args) {
    auto type_info = type_map_.find(expr);
    if (type_info == type_map_.end()) {
      type_map_.insert({expr, ResolvedTypeInfo(Type(), type_args)});
    } else {
      ICHECK(!type_info->second.type_args.defined());
      type_info->second.type_args = type_args;
    }
  }

第二步，解决约束

bool TypeSolver::Solve() {
  while (!update_queue_.empty()) {
    RelationNode* rnode = update_queue_.front();
    const auto& rel = rnode->rel;
    update_queue_.pop();
    ICHECK(!rnode->resolved);
    // update the relation with given evidence.
    Array<Type> args;
    for (auto* tlink = rnode->type_list.head; tlink != nullptr; tlink = tlink->next) {
      args.push_back(Resolve(tlink->value->FindRoot()->resolved_type));
      ICHECK_LE(args.size(), rel->args.size());
    }

    // We need to set this in order to understand where unification
    // errors generated by the error reporting are coming from.
    reporter_->SetSpan(rnode->span);

    try {
      // Call the Type Relation's function.
      bool resolved = rel->func(args, rel->num_inputs, rel->attrs, reporter_);

      if (resolved) {
        ++num_resolved_rels_;
      }

      rnode->resolved = resolved;
    } catch (const CompileError& err) {
      this->Emit(Diagnostic::Error(rnode->span) << err.what());
      rnode->resolved = false;
    }

    // Mark inqueue as false after the function call
    // so that rnode itself won't get enqueued again.
    rnode->inqueue = false;
  }

  // This criterion is not necessarily right for all the possible cases
  // TODO(tqchen): We should also count the number of in-complete types.
  return num_resolved_rels_ == rel_nodes_.size();
}

通过调用 Solve() 方法，我们求解填充好的类型约束。解决约束的过程使用了类型约束求解器（constraint solver）来尝试找到满足约束条件的类型赋值方案。

第三步，

Resolver(const std::unordered_map<Expr, ResolvedTypeInfo, ObjectPtrHash, ObjectPtrEqual>& tmap,
           TypeSolver* solver)
      : tmap_(tmap), solver_(solver) {}

Expr MixedModeMutator::VisitExpr(const Expr& expr) {
  auto fcheck_visited = [this](const Expr& expr) { return this->CheckVisited(expr); };
  auto fvisit_leaf = [this](const Expr& expr) { return this->VisitLeaf(expr); };
  if (memo_.count(expr)) {
    return memo_[expr];
  } else {
    ExpandDataflow(expr, fcheck_visited, fvisit_leaf);
    return memo_[expr];
  }
}

使用 Resolver 类的实例来将解析后的类型信息附加到已解析的表达式的checked_type 字段上。Resolver 类是负责类型解析和处理的工具类。它通过访问表达式的结构，并使用之前求解出的类型信息来确定每个表达式的准确类型。

respect～