TVM代码流程分析

TVM代码流程分析

TVM - 代码生成流程

本节主要介绍TVM的代码生成流程，即调用relay.build或tvm.build之后发生了什么，将深入到TVM的源代码进行剖析。（这里采用的依然是TVM v0.6）

首先区分两个build的区别：tvm.build主要针对单一算子（参照Tensor Expression一文），而relay.build是针对整个模型进行编译（参照GCN优化一文），而Relay最后也会调用到tvm::build做代码生成。

relay.build

通常的模型编译由以下两条语句完成。

# Build with Relay
with relay.build_config(opt_level=0):
    graph, lib, params = relay.build(func, target, params=params)

 

 跟踪细节

 

那么对relay.build进行跟踪，跳转进来是python/tvm/relay/build_module.py（这里是因为在relay/__init__.py中将build函数直接import到relay的命名空间，因此跳过了build_module这一层），其中的build函数是build_module内的全局函数(helper)。

def build(mod, target=None, target_host=None, params=None):
    # do somthing
 
    if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
        tophub_context = autotvm.tophub.context(list(target.values()))
    else:
        tophub_context = autotvm.util.EmptyContext()
 
    with tophub_context:
        bld_mod = BuildModule()
        graph_json, mod, params = bld_mod.build(func, target, target_host, params)
    return graph_json, mod, params

首先是寻找AutoTVM是否有预先tune好的参数记录，然后构造tophub_context，在其内部构建了BuildModule之后，才跳转到BuildModule.build，然后返回BuildModule.__init__中的内容。

class BuildModule(object):
    """Build a Relay function to run on TVM graph runtime. This class is used
    to expose the `RelayBuildModule` APIs implemented in C++.
    """
    def __init__(self):
        self.mod = _build_module._BuildModule()
        self._get_graph_json = self.mod["get_graph_json"]
        self._get_module = self.mod["get_module"]
        self._build = self.mod["build"]
        self._optimize = self.mod["optimize"]
        self._set_params_func = self.mod["set_params"]
        self._get_params_func = self.mod["get_params"]
 
    def build(self, func, target=None, target_host=None, params=None):
        target = _update_target(target)
 
        # Setup the params.
        if params:
            self._set_params(params)
        # Build the function
        self._build(func, target, target_host)
        # Get artifacts
        graph_json = self.get_json()
        mod = self.get_module()
        params = self.get_params()
 
        return graph_json, mod, params

而_build_module._BuildModule()又通过FFI在python/tvm/relay/_build_module.py中与C++函数建立联系（tvm._ffi._cytpes.function.Function.__call__）。

from tvm._ffi.function import _init_api
_init_api("relay.build_module", __name__)

对应的C++函数在src/relay/backend/build_module.cc

runtime::Module RelayBuildCreate() {
  auto exec = make_object<RelayBuildModule>();
  return runtime::Module(exec);
}
 
TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")
.set_body([](TVMArgs args, TVMRetValue* rv) {
  *rv = RelayBuildCreate();
});

也就是注册了一个RelayBuildModule供调用，由于主要用的是build函数，因此到RelayBuildModule中找对应的函数。这里TVM又用PackedFunc做了一层封装，见下。

PackedFunc GetFunction(const std::string& name,
                         const ObjectPtr<Object>& sptr_to_self) final {
      // ...
      if (name == "build") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        CHECK_EQ(args.num_args, 3);
        this->Build(args[0], args[1], args[2]);
      });
      // ...
}

也就是调用的是this->Build，再跳转过去会指向BuildRelay。

  void BuildRelay(
      Function func,
      const std::unordered_map<std::string, tvm::runtime::NDArray>& params) {
    // Optimize input Relay Function and returns Relay Module
    relay::Module relay_module = Optimize(func, targets_, params);
    // Get the updated function.
    func = relay_module->Lookup("main");
 
    // Generate code for the updated function.
    graph_codegen_ = std::unique_ptr<GraphCodegen>(new GraphCodegen());
    graph_codegen_->Init(nullptr, targets_);
    graph_codegen_->Codegen(func);
 
    ret_.graph_json = graph_codegen_->GetJSON();
    ret_.params = graph_codegen_->GetParams();
 
    auto lowered_funcs = graph_codegen_->GetLoweredFunc();
    if (lowered_funcs.size() == 0) {
      LOG(WARNING) << "no lowered funcs exist in the compiled module";
    } else {
      ret_.mod = tvm::build(
        lowered_funcs,
        target_host_,
        BuildConfig::Current());
    }
  }

经过多番跳转，终于到达build的核心模块，再来看TVM逐步做的工作。

  1. 优化
  2. 计算图生成
  3. 后端代码生成

优化

先是优化Optimize，可以看到这里的优化主要是设备无关的优化，是graph-level的针对tensor运算的优化。（这里的优化pass都已经在C++中实现，先前版本的NNVM似乎还是在Python中调用）

  relay::Module Optimize(
      Function func,
      const TargetsMap& targets,
      const std::unordered_map<std::string, runtime::NDArray>& params) {
    // BindParamsByName(func, params)
 
    // Perform Module->Module optimizations.
    relay::Module relay_module = relay::ModuleNode::FromExpr(func);
 
    Array<Pass> pass_seqs;
    // Run all dialect legalization passes.
    // ...
    pass_seqs.push_back(transform::SimplifyInference());
    //
    // ...fskip
    //
    pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
    pass_seqs.push_back(transform::CombineParallelConv2D(3));
    pass_seqs.push_back(transform::CombineParallelDense(3));
    pass_seqs.push_back(transform::FoldConstant());
    pass_seqs.push_back(transform::FoldScaleAxis());
    pass_seqs.push_back(transform::CanonicalizeCast());
    pass_seqs.push_back(transform::CanonicalizeOps());
    // ...AlterOpLayout
    pass_seqs.push_back(transform::FoldConstant());
 
    // Create a sequential pass and perform optimizations.
    transform::Pass seq = transform::Sequential(pass_seqs);
    // ... judge & do
    relay_module = seq(relay_module);
 
    // Handle heterogeneous compilation.
    transform::PassContext pass_ctx = PassContext::Current();
    if (targets_.size() > 1) {
      relay_module =
          RunDeviceAnnotationPass(relay_module, pass_ctx->fallback_device);
    }
 
    // Fuse the operations if it is needed.
    relay_module = transform::FuseOps()(relay_module);
    relay_module = transform::InferType()(relay_module);
    CHECK(relay_module.defined());
 
    return relay_module;
  }

计算图生成

对应GraphCodegen类，以同样的方式调用src/relay/backend/build_module.cc中的relay.build_module._GraphRuntimeCodegen（一样是FFI），然后跳转至src/relay/backend/graph_runtime_codegen.cc，其中已经用TVM_REGISTER_GLOBAL注册了对应函数，即用GraphRuntimeCodegenModule生成对应Object。

因此实际graph_codegen_->Codegen的函数是一个PackedFunc，定义在GraphRuntimeCodegen.Codegen，用来将relay::Function func进行遍历，然后生成计算图。

后端代码生成

Relay得到lower后的函数，最后一步则是交给tvm::build做代码生成，跳转到src/codegen/build_module.cc中的build函数（注意这里重载了几个版本），然后跳转到核心build，注意这里的build函数支持异构编译，只要再inputs划分好不同硬件设施即可。

// Build for heterogeneous execution.
runtime::Module build(const Map<Target, Array<LoweredFunc>>& inputs,
                      const Target& target_host,
                      const BuildConfig& config) {
  Array<LoweredFunc> fhost_all;
  std::vector<runtime::Module> device_modules;
 
  Target target_host_val = target_host;
  if (!target_host.defined()) {
    for (const auto& it : inputs) {
      if (it.first->device_type == kDLCPU) {
        target_host_val = it.first;
        break;
      }
    }
  }
 
  if (!target_host_val.defined()) {
    target_host_val = DefaultTargetHost(target_host_val);
  }
 
  for (const auto& it : inputs) {
    auto host_dev_funcs =
        split_dev_host_funcs(it.second, it.first, target_host_val, config);
    auto& fhost = host_dev_funcs[0];
    auto& fdevice = host_dev_funcs[1];
    // Get the module for a certain target.
    runtime::Module mdev = DeviceBuild(fdevice, it.first);
    for (const auto& it : fhost) {
      fhost_all.push_back(it);
    }
    device_modules.push_back(mdev);
  }
 
  runtime::Module mhost = codegen::Build(fhost_all, target_host_val->str());
  // Import all modules
  for (const auto& it : device_modules) {
    if (it.operator->()) {
      mhost.Import(it);
    }
  }
  return mhost;
}

当中最最核心的则是mhost = codegen::Build，最后跳转过去就开始调用代码生成模块了（src/codegen/codegen.cc）。

runtime::Module Build(const Array<LoweredFunc>& funcs,
                      const std::string& target) {
  // do something
 
  std::string build_f_name = "codegen.build_" + mode;
  // the build function.
  const PackedFunc* bf = runtime::Registry::Get(build_f_name);
  runtime::Module m = transformed_funcs.empty() ?
                      (*bf)(funcs, target) :
                      (*bf)(transformed_funcs, target);
  return m;
}

以生成LLVM IR为例，codegen.build_llvm会在src/codegen/llvm/llvm_module.cc注册，然后调用同个文件中的LLVMModuleNode->Init。这时会跳转到src/codegen/llvm/codegen_llvm.cc中的CodeGenLLVM类进行代码生成。

tvm.build

用tvm.build对算子进行编译则是按照以下方式进行调用，例子来自Tensor Expression。

s = tvm.create_schedule(C.op)
tgt = "llvm" # "cuda"
fadd = tvm.build(s,[A,B,C],target=tgt,name="myadd")

调用tvm.build后首先跳转到python/tvm/build_module.py，其中的build函数主要做两个步骤：

  1. lower高层次代码
  2. 后端代码生成

代码变换

lower高层次代码对应的是

flist = lower(inputs,args,name=name,binds=binds)

而lower函数同样在python/tvm/build_module.py中，类似于relay.build中的Optimize，但这里执行的是operator-level的优化，主要针对循环变换。

def lower(sch,
          args,
          name="default_function",
          binds=None,
          simple_mode=False):
 
    # initialization
 
    # Phase 0
    if isinstance(sch, schedule.Schedule):
        stmt = form_body(sch)
 
    for f in lower_phase0:
        stmt = f(stmt)
 
    compact = ir_pass.VerifyCompactBuffer(stmt)
    binds, arg_list = get_binds(args, compact, binds)
 
    # Phase 1
    stmt = ir_pass.RewriteForTensorCore(stmt, sch, binds)
    stmt = ir_pass.StorageFlatten(stmt, binds, 64, cfg.instrument_bound_checkers)
    stmt = ir_pass.CanonicalSimplify(stmt)
    for f in lower_phase1:
        stmt = f(stmt)
 
    # Phase 2
    if not simple_mode:
        stmt = ir_pass.LoopPartition(stmt, cfg.partition_const_loop)
    if cfg.disable_vectorize:
        stmt = ir_pass.SkipVectorize(stmt)
    else:
        stmt = ir_pass.VectorizeLoop(stmt)
    stmt = ir_pass.InjectVirtualThread(stmt)
    stmt = ir_pass.InjectDoubleBuffer(stmt, cfg.double_buffer_split_loop)
    stmt = ir_pass.StorageRewrite(stmt)
    stmt = ir_pass.UnrollLoop(
        stmt,
        cfg.auto_unroll_max_step,
        cfg.auto_unroll_max_depth,
        cfg.auto_unroll_max_extent,
        cfg.unroll_explicit)
    for f in lower_phase2:
        stmt = f(stmt)
 
    # Phase 3
    stmt = ir_pass.Simplify(stmt)
    stmt = ir_pass.RemoveNoOp(stmt)
    if not cfg.disable_select_rewriting:
        stmt = ir_pass.RewriteUnsafeSelect(stmt)
    for f in lower_phase3:
        stmt = f(stmt)
    # Instrument BoundCheckers
    if cfg.instrument_bound_checkers:
        stmt = ir_pass.InstrumentBoundCheckers(stmt)
    if simple_mode:
        return stmt
 
    return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)

优化Pass的主体实施都在src/api/api_pass.cc中，以tvm.ir_pass进行注册（注意由于C++函数中已经在tvm的命名空间里，故搜索时直接搜ir_pass才会出来对应的API）。

代码生成

lower完之后就进入到后端代码生成，对应build函数中的

mhost = codegen.build_module(fhost_all, str(target_host))

同样的原理，跳转至tvm/codegen.py，初始化tvm.codegen的API codegen._Build，调用FFI，跳转至src/api/api_codegen.cc，最后跳转至src/codegen/codegen.cc中的tvm::Build，之后的后端代码生成则与relay.build相同。

TVM - Tensor Expression

本节以向量加法为例，记录TVM最最基本的Tensor Expression的使用，以及简单的编译运行流程。

下面的代码为简单的向量加法，参考自Tensor Expression官方教程，在TVM v0.6下执行（注意与v0.7dev的模块有区别)。

import tvm
import numpy as np
 
# Tensor Expression
# args: (shape, label)
A = tvm.placeholder((10,), name='A')
B = tvm.placeholder((10,), name='B')
# args: (shape, function, label)
# function represented in lambda expression (element-wise)
#     lambda axis1, axis2, ... : f(axis1, axis2, ...)
C = tvm.compute((10,), lambda i: A[i] + B[i], name="C")
 
# generate schedule
s = tvm.create_schedule(C.op)
# print low level codes
print(tvm.lower(s,[A,B,C],simple_mode=True))

其中placeholder代表特定维度的张量，最后生成的代码会要求用户输入两个tensor，如果是C++代码，则要求用户输入两个float*。注意，会发现这个过程实际上是没有计算发生的，而只是定义了计算如何进行。

输出的low-level代码如下所示，还是相当好理解的，即i从0到10循环，循环内每次计算C[i]的值。

produce C {
  for (i, 0, 10) {
    C[i] = (A[i] + B[i])
  }
}

一些常用的循环优化API可以在这里找到。这里使用循环分割split作为尝试。

split(parent[, factor, nparts])

Split the stage either by factor providing outer scope, or both. Return outer, innervaiable of iteration.

bx, tx = s[C].split(C.op.axis[0],factor=2)
print(tvm.lower(s,[A,B,C],simple_mode=True))

由于对schedule的操作是原地变换，因此可以直接输出lower后的代码，发现确实已经改变了，原来的循环体变成5*2的循环。

produce C {
  for (i.outer, 0, 5) {
    for (i.inner, 0, 2) {
      C[((i.outer*2) + i.inner)] = (A[((i.outer*2) + i.inner)] + B[((i.outer*2) + i.inner)])
    }
  }
}

当然这一个schedule变换并没有带来任何好处，只是为了说明Tensor Expression应该怎么用。

之后就可以调用build生成目标代码了，可以设置target和target_host。

tgt = "c" # "llvm", "cuda"
fadd = tvm.build(s,[A,B,C],target=tgt,name="myadd")

然后可以创造运行时环境，进行运行测试。

n = 10
ctx = tvm.context(tgt,0)
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n,dtype=C.dtype), ctx)
fadd(a,b,c) # run
 
# test
tvm.testing.assert_allclose(c.asnumpy(),a.asnumpy() + b.asnumpy())
print(fadd.get_source())

生成的C代码如下

for (int32_t i_outer = 0; i_outer < 5; ++i_outer) {
  for (int32_t i_inner = 0; i_inner < 2; ++i_inner) {
    C[((i_outer * 2) + i_inner)] = (A[((i_outer * 2) + i_inner)] + B[((i_outer * 2) + i_inner)]);
  }
}

 

 生成的myadd.c完整代码如下

 

最后通过fadd.save("myadd.c")保存文件。

TVM - Relay IR Pass

本节介绍Relay IR Pass的构造。

Relay IR Pass核心依然是在C++中实现，但提供了Python接口，方便上层直接调用并对计算流图进行变换优化。

Pass管理器在include/tvm/relay/transform.h中，里面包含所有Pass的声明，希望做到

  • 管理调度不同的优化pass
  • 收集需要的分析信息，并且保持是最新的
  • 减少程序员实现新pass的麻烦

Python的接口函数声明在python/tvm/relay/transform.py中，在python/tvm/relay/_transform.py中通过FFI对C++函数进行调用，命名空间为relay._transform。

具体C++的实现则分为两个部分：

  • 高层IR图变换，源码在src/relay/pass中，集中变换则是在src/relay/backend/build_module.cc中的relay::Module Optimize
  • 后端代码的图变换，源码在src/relay/backend/vm中，集中变换在python/tvm/build_module.py中的lower函数

Pass的构造

  • PassInfo
·         class PassInfoNode : public RelayNode {
·         std::string name;
·         int opt_level;
·         std::vector<std::string> required;
·         };
  • PassContext
·         class PassContextNode : public RelayNode {
·         public:
·         ErrorReporter err_reporter;
·         int opt_level{2};
·         int fallback_device{static_cast<int>(kDLCPU)};
·         tvm::Array<tvm::Expr> required_pass;
·         tvm::Array<tvm::Expr> disabled_pass;
·         };
·        
·         class PassContext : public NodeRef {
·         public:
·         TVM_DLL static PassContext Create();
·         TVM_DLL static PassContext Current();
·         /* Other fields are omitted. */
·        
·         private:
·         // The entry of a pass context scope.
·         TVM_DLL void EnterWithScope();
·         // The exit of a pass context scope.
·         TVM_DLL void ExitWithScope();
·        
·         // Classes to get the Python `with` like syntax.
·         friend class tvm::With<PassContext>;
·         };
·        
·         struct RelayPassContextThreadLocalEntry {
·         /*! \brief The default pass context. */
·         PassContext default_context;
·         /*! \brief The current pass context. */
·         std::stack<PassContext> context_stack;
·         RelayPassContextThreadLocalEntry() {
·             default_context = PassContext(make_node<PassContextNode>());
·         }
·         };
·        
·         /*! \brief The thread-local store to hold the pass context. */
·         typedef dmlc::ThreadLocalStore<RelayPassContextThreadLocalEntry>
·             RelayPassContextThreadLocalStore;
  • Pass Constructs：提供基类
·         class PassNode : RelayNode {
·         virtual PassInfo Info() const = 0;
·         virtual Module operator()(const IRModule& mod
·                                     const PassContext& pass_ctx) const = 0;
·         };

也就是说，一个Pass一定是作用在特定context下的IRModule，所有Pass都设计成Module到Module的映射，完整Pass的定义在src/relay/ir/transform.cc和src/ir/transform.cc中。

Module-Level

class ModulePassNode : PassNode {
  PassInfo pass_info;
  runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func;
  Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
  // Other members/methods are omitted
};

Function-Level

class FunctionPassNode : PassNode {
  PassInfo pass_info;
  runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func;
  Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
  bool SkipFunction(const Function& func) const;
  // Other members/methods are omitted...
};

Sequential

类似于PyTorch中的nn.Sequential，顺序执行多个Pass

class SequentialPassNode : PassNode {
  PassInfo pass_info;
  // Passes need to be executed.
  Array<Pass> passes;
  bool PassEnabled(const PassInfo& info) const;
  Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
};

 

References

 

 

posted @ 2021-11-13 18:00  吴建明wujianming  阅读(772)  评论(1)    收藏  举报