TVM深度学习编译器框架详解：整体架构与关键组件

TVM深度学习编译器框架详解：整体架构与关键组件

引用

CSDN

1.

https://blog.csdn.net/u012294613/article/details/145421249

TVM是一个开源的端到端深度学习编译器框架，旨在优化深度学习模型在各种硬件平台上的性能。本文将从整体架构、关键组件、代码生成流程等方面，帮助读者全面了解TVM的核心机制。

TVM的整体架构

首先，让我们来看一下TVM的完整架构图：

TVM的完整架构图

从架构图中可以看出，TVM主要由以下几个部分组成：

• 前端（Frontend）：负责将各种深度学习框架的计算图转换为Relay IR。
• 中间表示（IR）：分为两层，上层是面向前端的Relay IR，下层是面向LLVM的底层IR（TIR）。
• 优化（Transform）：包括硬件无关的Pass和硬件相关的优化。
• 目标翻译（Target Translation）：将IR转换为目标硬件可执行的代码。
• 运行时（Runtime）：负责执行生成的代码。

虽然TVM抽象了Relay IR和TIR两个层次，但实际实现上有共同之处。例如，Relay表达式树的节点都继承自ExprNode，而ExprNode又是从Object类派生出来的。这种设计使得两个层次的IR在实现上具有统一性。

接下来，简要介绍一下架构图中的关键组成部分：

• Import：这是TVM的前端组件，负责将各种深度学习框架的模型导入到TVM的IRModule中。
• Transform：可以将一个IRModule等价转换成另外一个功能相同的IRModule，达到优化的效果。
• IRModule：它是functions的集合，包含两种关键的Function集合：relay::Function和tir::PrimFunc。
• Target Translation：编译器将IRModule变换为目标硬件上可执行的格式。
• Passes：是对计算图的一些优化和转换，比如常量折叠、算子融合、死代码消除等。

这里以一个例子来直观了解一下IRModule长什么样子：

import tvm
from tvm import relay
import numpy as np
from tvm.contrib import graph_executor

# build model
n = 2
x = relay.var("x", shape=(n,), dtype='float32')
y = relay.nn.softmax(x)
net = relay.Function([x], y)

# build and lowering
module = tvm.IRModule.from_expr(net)
lib = relay.build(module, "llvm")
dev = tvm.cpu(0)
input = tvm.nd.array(np.random.uniform(size=[n]).astype('float32'), dev)
m = graph_executor.GraphModule(lib["default"](dev))

# set inputs
m.set_input("x", input)

# execute
m.run()

# get outputs
tvm_output = m.get_output(0)
print(module)

输出结果：

def @main(%x: Tensor[(2), float32]) {
  nn.softmax(%x)
}

TVM的关键组件

TVM中一些基础组件的交互图如下：

TVM的组件交互图

下面分别介绍一下这些组件：

• Support：TVM的一些通用组件，比如socketing、logging等。
• Frontends：TVM的前端，完成各种深度学习框架的计算图到Relay IR的转化。
• Relay：这是一个high-level 计算图的描述，它有自己的IR表示，用这些IR表示来描述神经网络结构。
• AutoTVM/Ansor：可选的自动调优模块。
• Topi：这是一个Tensor计算库，里边包含了很多神经网络通用的算子，比如矩阵乘法、点乘、卷积等。
• Tir：相对于Relay IR，这个层次的IR更接近底层和硬件实现。
• Te：Te表示Tensor Expression，用户可以通过调用te中的函数来构建Tir。
• Node：node是在IR的基础上增加了一些新的特性，可以允许用户对一些函数进行访问。
• Driver：基于硬件的驱动。
• Arith：这个和Tir有关，可以在Tir优化时进行一些分析。
• Runtime：runtime封装了图结构的转化、优化、代码生成，以及程序在硬件上的执行，为客户提供一个API接口完成所有的编译过程。

注意一下这个组件交互图中的箭头：

• 虚线箭头表示组件之间在实现level上有依赖，比如Tir pass在实现上必须依赖arith来分析。
• 实线箭头表示两个组件在前端level上有依赖关系，比如Relay IR是依赖于Tir的，因为Relay IR会被lower成Tir。

在前后端的交互上，TVM将所有的核心数据结构都暴露到了Python前端，这使得它具有足够的灵活性和易用性。具体体现为：

• 所有的核心对象都可以通过Python API直接构造和操作，比如IR Module。
• 支持在前端自定义组合Pass。
• 可以通过TVM的API直接操作IR，支持Python端写Pass。

TVM的代码生成流程

了解了TVM的基本组件和整体架构之后，有必要再了解一下TVM的代码生成流程。具体可以参考《TVM的CodeGen流程》。

参考资料

• TVM官方文档
• TVM入门教程
• TVM GitHub仓库