TVM编译流程 详细介绍

TVM编译流程

TVM是一个深度学习描述框架，通过Python代码描述算子（输入、输出、运算方法等）形成抽象语法树（Abstract Syntax Tree，AST），然后在TVM内部转换为中间表示（Intermediate Representation，IR），最终转换成目标平台的机器代码，以作为算子用于构成更复杂的神经网络。

如下图所示，TVM主要流程可以分为：

1. TVM主要的流程就是从深度学习的框架，如PyTorch、TensorFlow、MxNet等，将其转换为Relay IR。
2. 在Relay这一层会进行计算图级别优化，如常量折叠、死代码消除、算子融合等。
3. Lower：将高层IR表示转化成低阶TIR表示，并进行 算子级别的优化。
4. 最后进行内存分配和硬件可执行程序的生成

下文，将详细讲解这几个部分，并补充TVM的相关组件。

模型加载及转换

TVM 作为人工智能领域的编译器，主要处理来自 PyTorch、TensorFlow、MxNet 和 ONNX 等深度学习框架的模型。这些模型被转换成 TVM 专用的中间表示形式——Relay IR。我们关注模型的结构、参数和输入形状等关键信息。在进行模型推理前，必须对数据进行预处理，以适应模型的输入需求，这包括数据的张量转换和类型转换，例如转换为 float16 或 int8 等低精度格式。数据预处理完成后，即可对模型进行优化，以提高推理效率。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

mod, params = relay.frontend.from_pytorch(model, shape_list)

通过 relay.frontend.from_pytorch可以将pytorch代码转换成relay IR格式的代码，转换效果如下：

def @main(%x0: Tensor[(32, 32), float32] /* span=aten::linear_0.x0:0:0 */, %x1: Tensor[(32, 32), float32] /* span=aten::linear_5.x1:0:0 */, %aten::linear_0.weight: Tensor[(512, 32), float32] /* span=aten::linear_0.weight:0:0 */, %aten::linear_0.bias: Tensor[(512), float32] /* span=aten::linear_0.bias:0:0 */, %aten::linear_1.weight: Tensor[(1024, 512), float32] /* span=aten::linear_1.weight:0:0 */, %aten::linear_1.bias: Tensor[(1024), float32] /* span=aten::linear_1.bias:0:0 */, %aten::linear_2.weight: Tensor[(2048, 1024), float32] /* span=aten::linear_2.weight:0:0 */, %aten::linear_2.bias: Tensor[(2048), float32] /* span=aten::linear_2.bias:0:0 */, %aten::linear_3.weight: Tensor[(4096, 2048), float32] /* span=aten::linear_3.weight:0:0 */, %aten::linear_3.bias: Tensor[(4096), float32] /* span=aten::linear_3.bias:0:0 */, %aten::linear_4.weight: Tensor[(64, 4096), float32] /* span=aten::linear_4.weight:0:0 */, %aten::linear_4.bias: Tensor[(64), float32] /* span=aten::linear_4.bias:0:0 */, %aten::linear_5.weight: Tensor[(512, 32), float32] /* span=aten::linear_5.weight:0:0 */, %aten::linear_5.bias: Tensor[(512), float32] /* span=aten::linear_5.bias:0:0 */, %aten::linear_6.weight: Tensor[(1024, 512), float32] /* span=aten::linear_6.weight:0:0 */, %aten::linear_6.bias: Tensor[(1024), float32] /* span=aten::linear_6.bias:0:0 */, %aten::linear_7.weight: Tensor[(2048, 1024), float32] /* span=aten::linear_7.weight:0:0 */, %aten::linear_7.bias: Tensor[(2048), float32] /* span=aten::linear_7.bias:0:0 */, %aten::linear_8.weight: Tensor[(4096, 2048), float32] /* span=aten::linear_8.weight:0:0 */, %aten::linear_8.bias: Tensor[(4096), float32] /* span=aten::linear_8.bias:0:0 */, %aten::linear_9.weight: Tensor[(64, 4096), float32] /* span=aten::linear_9.weight:0:0 */, %aten::linear_9.bias: Tensor[(64), float32] /* span=aten::linear_9.bias:0:0 */, %aten::linear_10.weight: Tensor[(256, 128), float32] /* span=aten::linear_10.weight:0:0 */, %aten::linear_10.bias: Tensor[(256), float32] /* span=aten::linear_10.bias:0:0 */, %aten::linear_11.weight: Tensor[(128, 256), float32] /* span=aten::linear_11.weight:0:0 */, %aten::linear_11.bias: Tensor[(128), float32] /* span=aten::linear_11.bias:0:0 */, %aten::linear_12.weight: Tensor[(2, 128), float32] /* span=aten::linear_12.weight:0:0 */, %aten::linear_12.bias: Tensor[(2), float32] /* span=aten::linear_12.bias:0:0 */) {
  %0 = nn.dense(%x0, %aten::linear_0.weight, units=None) /* span=aten::linear_0:0:0 */;
  %1 = nn.bias_add(%0, %aten::linear_0.bias, axis=-1) /* span=aten::linear_0:0:0 */;
  %2 = nn.relu(%1) /* span=aten::relu_0:0:0 */;
  %3 = nn.dense(%2, %aten::linear_1.weight, units=None) /* span=aten::linear_1:0:0 */;
  %4 = nn.bias_add(%3, %aten::linear_1.bias, axis=-1) /* span=aten::linear_1:0:0 */;
  %5 = nn.relu(%4) /* span=aten::relu_1:0:0 */;
  %6 = nn.dense(%5, %aten::linear_2.weight, units=None) /* span=aten::linear_2:0:0 */;
  %7 = nn.bias_add(%6, %aten::linear_2.bias, axis=-1) /* span=aten::linear_2:0:0 */;
  %8 = nn.relu(%7) /* span=aten::relu_2:0:0 */;
  %9 = nn.dense(%8, %aten::linear_3.weight, units=None) /* span=aten::linear_3:0:0 */;
  %10 = nn.bias_add(%9, %aten::linear_3.bias, axis=-1) /* span=aten::linear_3:0:0 */;
  %11 = nn.relu(%10) /* span=aten::relu_3:0:0 */;
  %12 = nn.dense(%11, %aten::linear_4.weight, units=None) /* span=aten::linear_4:0:0 */;
  %13 = nn.bias_add(%12, %aten::linear_4.bias, axis=-1) /* span=aten::linear_4:0:0 */;
  %14 = nn.dense(%x1, %aten::linear_5.weight, units=None) /* span=aten::linear_5:0:0 */;
  %15 = nn.bias_add(%14, %aten::linear_5.bias, axis=-1) /* span=aten::linear_5:0:0 */;
  %16 = nn.relu(%15) /* span=aten::relu_5:0:0 */;
  %17 = nn.dense(%16, %aten::linear_6.weight, units=None) /* span=aten::linear_6:0:0 */;
  %18 = nn.bias_add(%17, %aten::linear_6.bias, axis=-1) /* span=aten::linear_6:0:0 */;
  %19 = nn.relu(%18) /* span=aten::relu_6:0:0 */;
  %20 = nn.dense(%19, %aten::linear_7.weight, units=None) /* span=aten::linear_7:0:0 */;
  %21 = nn.bias_add(%20, %aten::linear_7.bias, axis=-1) /* span=aten::linear_7:0:0 */;
  %22 = nn.relu(%21) /* span=aten::relu_7:0:0 */;
  %23 = nn.dense(%22, %aten::linear_8.weight, units=None) /* span=aten::linear_8:0:0 */;
  %24 = nn.bias_add(%23, %aten::linear_8.bias, axis=-1) /* span=aten::linear_8:0:0 */;
  %25 = nn.relu(%24) /* span=aten::relu_8:0:0 */;
  %26 = nn.dense(%25, %aten::linear_9.weight, units=None) /* span=aten::linear_9:0:0 */;
  %27 = nn.bias_add(%26, %aten::linear_9.bias, axis=-1) /* span=aten::linear_9:0:0 */;
  %28 = nn.relu(%13) /* span=aten::relu_4:0:0 */;
  %29 = nn.relu(%27) /* span=aten::relu_9:0:0 */;
  %30 = (%28, %29) /* span=aten::cat_0:0:0 */;
  %31 = concatenate(%30, axis=1) /* span=aten::cat_0:0:0 */;
  %32 = nn.dense(%31, %aten::linear_10.weight, units=None) /* span=aten::linear_10:0:0 */;
  %33 = nn.bias_add(%32, %aten::linear_10.bias, axis=-1) /* span=aten::linear_10:0:0 */;
  %34 = nn.relu(%33) /* span=aten::relu_10:0:0 */;
  %35 = nn.dense(%34, %aten::linear_11.weight, units=None) /* span=aten::linear_11:0:0 */;
  %36 = nn.bias_add(%35, %aten::linear_11.bias, axis=-1) /* span=aten::linear_11:0:0 */;
  %37 = nn.relu(%36) /* span=aten::relu_11:0:0 */;
  %38 = nn.dense(%37, %aten::linear_12.weight, units=None) /* span=aten::linear_12:0:0 */;
  %39 = nn.bias_add(%38, %aten::linear_12.bias, axis=-1) /* span=aten::linear_12:0:0 */;
  nn.relu(%39) /* span=aten::relu_12:0:0 */
}

需要 注意的是，在Relay表示中的算子并不是真正的算子实现，而是一种声明，就像声明式语言一样关注算子是什么类型，而不关注算子具体怎么执行计算。此外，Relay IR作为TVM编译中第一个对接输入模型的IR对象，具有较高的抽象层次，如果想使用TVM优化模型，首先需要Realy中间表示支持模型中的所有算子，对于不支持的算子就需要编写对应的转换函数手动拓展TVM。

tvm的 IR设计

TVM中的IR分为两层，上层是面向前端的Relay IR，下层是面向LLVM的底层IR（也可以叫Tir，将在后序部分进行讲解）。

根据TVM官网描述的这张图可以看到，Relay中间表示由IRModule组成，其中包含了输入的整个神经网络，是后续TVM编译优化的对象，IRModule又由Relay中的Function组成。无论是Relay还是后续的要介绍到的TensorIR在TVM中都共用一套IR基础设施，IR中的元素都是主要由Type和Expr两个主要的基类派生而来

再来具体看看代码中对于IRModule的定义，

class IRModuleNode : public Object {
 public:
  /*! \brief A map from ids to all global functions. */
  Map<GlobalVar, BaseFunc> functions;
  /*! \brief A map from global type vars to ADT type data. */
  Map<GlobalTypeVar, TypeData> type_definitions;
  /*! \brief The source map for the module. */
  SourceMap source_map;
  /* \brief Additional attributes storing meta-data about the module. */
  DictAttrs attrs;
  /*! \brief Globally static object that are referred by the IR itself */
  Map<String, Array<GlobalInfo>> global_infos;
  ........
  }

可以 看到，IRModule是functions的集合，结合图3可以看到它包含两种最关键的Function集合，即relay::Function和tir::PrimFuc 。

• 上层relay::Function继承自BaseFunction，relay::Function对应一个end2end的模型，可以理解为一个支持控制流，递归，以及复杂数据结构的计算图。
• 下层tir::PrimFunc也继承自BaseFunction，tir::PrimFunc包含了一些底层threading，vector/tensor的指令。通常为模型中的一个OP执行单元。
• 在编译阶段，一个relay::Function可能会被lower成多个tir::PrimFunc

还有type类型也是tvm IR中 比较 重要的 组成成分，包含bool、int8，float32等基础数据类型，以及张量Tensor和元组Touple等类型。在TVM Type类中可以看到TVM各个层级IR需要的Type。同时Relay中还提供了描述Relay函数的输入和输出类型之间关系的类型关系特性，允许用户扩展类型推断，方便算子的shape推理。

计算图优化

将模型初步转换成Relay IR的 表达格式之后，就可以正式对RelayIR的 计算图进行优化。Relay的Transform是的硬件无关的Pass，例如常规的常量折叠，算符融合，死代码消除等等，以及张量计算相关的一些特殊Pass如transformation，scaling factor folding。

Passes：pass是对计算图的一些优化和转换。

在Relay优化pipline的后期，会运行一个算子融合的pass来进行计算图的拆分 ，算子融合指在设备上执行过程中，并不是逐层当做独立操作执行，而是将多个算子融合成一个操作，避免频繁的访存操作。TVM规定了算子的四种类型，包括injectvie，1对1映射如add；reduction，多对少映射如sum；complex-out-fusable，逐元素复用映射到输出如conv2d；opaque，无法融合的类型如sort。算子融合具体的解释可以阅读深入解析算子融合。

IR 的 lower

前文 讲过 ，realy层面的IR的是没有具体的计算定义的，所以在进行了计算图优化之后，tvm会将relay IR中的计算节点转换成te表达式，并且生成具体的调度策略并将模型转换成tir的格式，并使用 tir级别的pass优化模型，也可以称之为算子级别的优化，这部分的主要功能是lower，不过也有一些optimization。比如将访问多维数据扁平化为一维指针访问、针对特定的后端进行intrinsics扩展、或者根据运行时调用约定装饰函数（方便后续call）；注意这个阶段保留了一些底层优化没有做，而是交给了下游的LLVM或者CUDA C编译器来进行，比如寄存器分配等等。

Te：Te表示Tensor Expression，用户可以通过调用te中的函数来构建Tir。（这意味着TVM允许用户直接通过Te来写神经网络）

调度策略：具体如何执行计算，如数据如何加载存储，使用何种优化手段如循环分块、循环展开、多线程等。

Tir：相对于Relay IR，这个层次的IR更接近底层和硬件实现。

调度策略的生成

TVM的设计目标之一就是希望支持生成不同平台的高性能张量化代码，所以TVM提供了两种调度策略的生成方法，一种 是基于模版的AutoTVM，一种是基于探索的Ansor。

TVM有许多典型的scheduler，感兴趣可以移步tvm schedule详细举例

AutoTVM

TVM采用了一种在scheduler形成的搜索空间中，使用机器学习算法找到最优化选择的方案。基于这些小粒度的scheduler，TVM定义了一个程序变换可执行操作的集合，叫做调度原语(scheduling primitives)，比如循环变化（split, unrool…），内联，向量化。这个阶段同时引入了target的概念，将一些target的特性考虑在内来进行搜索，比如target的寄存器可以同时处理多少个数据就决定了这部分代码向量化的行为。

接下来就是不断的跑程序，记录性能，调整调度选择，在最后执行程序时，对于某一段子图程序来说，将从log文件中选择最优的性能对应的调度方案来进行执行。

Ansor

AutoTVM需要事先编写模板来组成调度的搜索空间，最佳性能的上限取决于模板的设计，这对模板的编写带来了很高的要求。所以提出了第二代调度策略调优策略-Ansor(Auto Scheduler)，这种 方法完全取消了基于模版的调度 生成机制，而是采用自动的，无干预的优化，开发者无需手动指定优化手段。

Ansor自动生成一个覆盖全面的优化的大搜索空间，并为空间中的每个张量程序提供被选择的机会。

1. 首先，它自动构建一个大的搜索空间，以覆盖给定计算定义的尽可能多的张量程序。
2. 其次，在大搜索空间中高效搜索，该搜索空间可能比现有模板可以覆盖的范围大几个数量级。
3. 最后，在优化具有许多子图的整个 DNN 时，识别对端到端性能至关重要的子图并对其进行优先级排序，因为资源是有限的，应该将调优时间和算力资源分配给对性能有更大影响的子图。

可以看到Ansor 是在基于子图拆分(算子融合)的基础上，经过程序采样，性能微调和任务调度器组成的一个自动调优工具。

1. 程序采样器：为了在无模板的前提下自动生成搜索空间，递归地应用一组推导规则来扩展搜索空间；为了避免在搜索空间中陷入局部最优，使用随机抽取完整的程序给每个采样点相同概率。将搜索空间定为两级，高层结构称为草图（sketch），低级结构（分块大小、并行等）称为注解（annotation）。
2. 性能微调：使用遗传算法 和 成本学习 模型来微调 生成的策略
3. 任务调度器：将生成的策略下放到具体的硬件上进行性能评估，并将评估 结果反馈给上述的部分

代码生成

在代码生成阶段，就是将tir 转换成具体可以在硬件上执行的代码(指令)。

代码详解

首先区分两个build的区别：

tvm.build主要针对单一算子

relay.build是针对整个模型进行编译，而Relay最后也会调用到tvm::build做代码生成。

# 前端模型导入
mod, params = relay.frontend.from_pytorch(model, shape_list)  # 第二个add在这里就被优化了
print(mod.astext(show_meta_data=False))
# 指定优化等级
opt_level = 1
target = "cuda"
# 进行编译以及代码生成
with tvm.transform.PassContext(opt_level=opt_level):
    lib = relay.build(mod, target, params=params)
dev = tvm.cuda(0)
# 生成执行器
module = graph_executor.GraphModule(lib["default"](dev))
# 具体填充参数以及输入 数据
module.set_input(**params)
module.set_input('x0',x0)
module.set_input('x1',x1)
# 开始执行模型
module.run()
# 获取tvm输出
out = module.get_output(0, tvm.nd.empty(pytorch_result.shape)).numpy()

对 relay.build进行追踪 ，其实现主要是在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();
});

在这里 TVM又用PackedFunc 做了一层封装，

PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) final {
    if (name == "get_graph_json") {
      return PackedFunc(
          [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetGraphJSON(); });
    } else if (name == "get_module") {
      return PackedFunc(
          [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetModule(); });
    } else if (name == "build") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        ICHECK_EQ(args.num_args, 8);
        this->Build(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7]);
      });
    }
    .......
    }

可以看到，我们主要采用的build的模型，所以可以去查看this->Build(…)

  void BuildRelay(IRModule relay_module, const String& mod_name) {
    // Relay IRModule -> IRModule optimizations.
    IRModule module = WithAttrs(
        relay_module, {{tvm::attr::kExecutor, executor_}, {tvm::attr::kRuntime, runtime_}});
    relay_module = OptimizeImpl(std::move(module));

    // Get the updated function and new IRModule to build.
    // Instead of recreating the IRModule, we should look at the differences between this and the
    // incoming IRModule to see if we can just pass (IRModule, Function) to the code generator.
    Function func = Downcast<Function>(relay_module->Lookup("main"));
    IRModule func_module = WithAttrs(IRModule::FromExpr(func),
                                     {{tvm::attr::kExecutor, executor_},
                                      {tvm::attr::kRuntime, runtime_},
                                      {tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_},
                                      {tvm::attr::kConstantMemoryPools, constant_memory_pools_}});

    // Generate code for the updated function.
    executor_codegen_ = MakeExecutorCodegen(executor_->name);
    executor_codegen_->Init(nullptr, config_->primitive_targets);
    executor_codegen_->Codegen(func_module, func, mod_name);
    executor_codegen_->UpdateOutput(&ret_);
    ret_.params = executor_codegen_->GetParams();

    auto lowered_funcs = executor_codegen_->GetIRModule();

    // No need to build for external functions.
    Target ext_dev("ext_dev");
    if (lowered_funcs.find(ext_dev) != lowered_funcs.end()) {
      lowered_funcs.Set(ext_dev, IRModule());
    }

    const Target& host_target = config_->host_virtual_device->target;
    const runtime::PackedFunc* pf = runtime::Registry::Get("codegen.LLVMModuleCreate");
    // When there is no lowered_funcs due to reasons such as optimization.
    if (lowered_funcs.size() == 0) {
      if (host_target->kind->name == "llvm") {
        CHECK(pf != nullptr) << "Unable to create empty module for llvm without llvm codegen.";
        // If we can decide the target is LLVM, we then create an empty LLVM module.
        ret_.mod = (*pf)(host_target->str(), "empty_module");
      } else {
        // If we cannot decide the target is LLVM, we create an empty CSourceModule.
        // The code content is initialized with ";" to prevent complaining
        // from CSourceModuleNode::SaveToFile.
        ret_.mod = tvm::codegen::CSourceModuleCreate(";", "", Array<String>{});
      }
    } else {
      ret_.mod = tvm::TIRToRuntime(lowered_funcs, host_target);
    }

    auto ext_mods = executor_codegen_->GetExternalModules();
    ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods, host_target,
                                                  runtime_, executor_,
                                                  executor_codegen_->GetExecutorCodegenMetadata());
    // Remove external params which were stored in metadata module.
    for (tvm::runtime::Module mod : ext_mods) {
      auto pf_var = mod.GetFunction("get_const_vars");
      if (pf_var != nullptr) {
        Array<String> variables = pf_var();
        for (size_t i = 0; i < variables.size(); i++) {
          auto it = ret_.params.find(variables[i].operator std::string());
          if (it != ret_.params.end()) {
            VLOG(1) << "constant '" << variables[i] << "' has been captured in external module";
            ret_.params.erase(it);
          }
        }
      }
    }
  }

可以看到这段build代码主要做了三个工作：

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

有兴趣的 朋友 ，可以自行探索具体的代码流程。