跳到主要内容
版本:0.13.0

向 TVM 中添加 Codegen

随着深度学习工作负载所针对的硬件设备数量不断增加,用户在各种设备上实现高性能所需的知识也在不断增加。为了让数据科学家在开发新模型时不必担心性能问题,硬件厂商或是基于一些常见的深度学习算子,提供 MKLDNN 或 cuDNN 等库,或是提供 TensorRT 等框架,让用户按照某种方式描述模型,从而提高模型性能。

然而,用户在尝试使用新的库或设备时,必须学习新的编程接口。因此,一个统一的编程接口变得越来越重要:1)让所有用户及硬件厂商信息同步,2)提供一个可行的解决方案,让特定硬件或库只支持具有极高性能的、广泛使用的算子,不受支持的算子则回退到 CPU/GPU 等通用设备。

本开发手册演示了硬件厂商如何轻松实现自己的 Codegen,并将其注册为 Relay 后端编译器,从而支持自己的硬件设备/库。本手册涵盖了两种基于不同计算图的 codegen:

1. 希望生成 C 代码。

如果你的硬件已经具备了一个高度优化的 C/C++ 库,如对于 CPU 而言的 Intel CBLAS/MKL 库,或针对 GPU 而言的 NVIDIA CUBLAS 库,那么本节内容非常适合你。幸运的是,C 源代码模块与 TVM runtime 模块完全兼容,这意味着生成的代码可以由任何具有适当编译标志的 C/C++ 编译器编译,因此用户只需实现一个能为子图生成 C 代码的 codegen,并将 C 源代码模块集成到 TVM runtime 模块中。下一节内容讲详细演示如何为硬件实现 C codegen。

2. 希望生成任意计算图。

有时候,硬件可能需要其他形式的计算图如 JSON。这种情况下,用户不仅要实现一个 codegen,还要实现一个自定义 TVM runtime 模块,从而使得 TVM runtime 知道如何执行这个计算图。如果你的硬件已经拥有完整的计算图执行引擎(graph execution engine),如适用于 GPU 的 TensorRT,那么该解决方案对你而言非常具有参考价值。

完成 codegen 和 runtime 后,可以让客户借助你的自定义标签,对模型进行注释并加以利用。终端用户如何注释和启动特定 codegen 的教程,将在后续进行补充。

实现 C Codegen

在这一部分中,我们将演示如何借助预实现的算子函数,生成 C 代码的 codegen。简单起见,本示例 codegen 不依赖于第三方库。相反,我们在 C 中手动实现了两个宏:

#define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)         \
extern "C" void p_ID_(float* a, float* b, float* out) { \
for (int64_t i = 0; i < p_DIM1_; ++i) { \
out[i] = a[i] p_OP_ b[i]; \
} \
}

#define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_) \
extern "C" void p_ID_(float* a, float* b, float* out) { \
for (int64_t i = 0; i < p_DIM1_; ++i) { \
for (int64_t j = 0; j < p_DIM2_; ++j) { \
int64_t k = i * p_DIM2_ + j; \
out[k] = a[k] p_OP_ b[k]; \
} \
} \
}

使用这两个宏,可以为一维和二维张量生成二元算子(binary operator)。例如,给定如下所示的子图,假设所有输入都是 shape 为(10, 10)的二维张量:

c_compiler_input0
|
add <-- c_compiler_input1
|
subtract <-- c_compiler_input2
|
multiply <-- c_compiler_input3
|
out

我们的目标是生成以下可编译代码来执行子图:

#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include <cstdint>
#include <cstring>
#include <iostream>

#define GCC_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_) \
extern "C" void p_ID_(float* a, float* b, float* out) { \
for (int64_t i = 0; i < p_DIM1_; ++i) { \
out[i] = a[i] p_OP_ b[i]; \
} \
}

#define GCC_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_) \
extern "C" void p_ID_(float* a, float* b, float* out) { \
for (int64_t i = 0; i < p_DIM1_; ++i) { \
for (int64_t j = 0; j < p_DIM2_; ++j) { \
int64_t k = i * p_DIM2_ + j; \
out[k] = a[k] p_OP_ b[k]; \
} \
} \
}

// 注 1
GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);
GCC_BINARY_OP_2D(gcc_0_1, -, 10, 10);
GCC_BINARY_OP_2D(gcc_0_2, +, 10, 10);

// 注 2
extern "C" void gcc_0_(float* gcc_input0, float* gcc_input1,
float* gcc_input2, float* gcc_input3, float* out) {
float* buf_0 = (float*)malloc(4 * 100);
float* buf_1 = (float*)malloc(4 * 100);
gcc_0_2(gcc_input0, gcc_input1, buf_0);
gcc_0_1(buf_0, gcc_input2, buf_1);
gcc_0_0(buf_1, gcc_input3, out);
free(buf_0);
free(buf_1);
}

// 注 3
extern "C" int gcc_0_wrapper(DLTensor* arg0, DLTensor* arg1, DLTensor* arg2,
DLTensor* arg3, DLTensor* out) {
gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
}
TVM_DLL_EXPORT_TYPED_FUNC(gcc_0, gcc_0_wrapper);

这里详细介绍一下上面代码里的注释:

  • 注1:子图中三个节点的函数实现。
  • 注2:通过分配中间数组(intermediate buffer)并调用相应函数来执行子图的函数。
  • 注3:TVM runtime 兼容的包装函数。它接收一个输入张量列表和一个输出张量(最后一个参数),并将其转换为正确的数据类型,调用注2 中描述的子图函数。此外,TVM_DLL_EXPORT_TYPED_FUNC 是一个 TVM 宏,它通过将所有张量打包到 TVMArgs 来生成另一个函数 gcc_0,该函数具有统一的函数参数。因此,TVM runtime 可以直接调用 gcc_0 来执行子图,无需其他操作。生成上述代码后,TVM 能够将其与计算图的其余部分一起编译并导出单个库以进行部署。

在本节的其余部分,我们将逐步创建一个 codegen,来实现上述代码。你的 codegen 必须位于 src/relay/backend/contrib/<your-codegen-name>/。在这个例子中,我们将 codegen 命名为 "codegen_c",并将其放在 /src/relay/backend/contrib/codegen_c/ 目录下。你可以随时查看这个文件,了解完整的实现过程。

具体来说,我们将在这个文件中实现两个类,两个类的关系如下:

            subgraph                                subgraph
TVM backend -----------------------------> CSourceCodegen -------------> CodegenC
^ | ^ |
| | | |
---------------------------------------- ------------------------
generated C source runtime module generated C code

当 TVM 后端发现 Relay 计算图中的函数(子图),用注册的编译器标签(本例中为 ccompiler)进行了注释时,TVM 后端就会调用 CSourceCodegen 并传递子图。 CSourceCodegen 的成员函数 CreateCSourceModule 将:

1)为子图生成 C 代码;

2)将生成的 C 代码包装到 C source runtime 模块中,以便 TVM 后端进行编译和部署。

特别是,C codegen 对 CodegenC 类是透明的,因为它提供了许多有用的实用程序来简化 codegen 实现。下面的章节将自下而上实现这两个类。

实现 CodegenC

src/relay/backend/contrib/codegen_c/codegen.cc 中,首先在 tvm.relay.contrib 的命名空间下创建一个 codegen 类骨架:

#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
#include <tvm/relay/type.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/object.h>

#include <fstream>
#include <sstream>

#include "codegen_c.h"

namespace tvm {
namespace relay {
namespace contrib {

class CodegenC : public ExprVisitor, public CodegenCBase {
public:
explicit CodegenC(const std::string& id) { this->ext_func_id_ = id; }

void VisitExpr_(const VarNode* node) { ; }
void VisitExpr_(const CallNode* call) final { ; }
std::string JIT() { ; }

private:
/*! \brief The function id that represents a C source function. */
std::string ext_func_id_ = "";
/*! \brief The index of a wrapped C function. */
int func_idx = 0;
/*! \brief The index of allocated buffers. */
int buf_idx_ = 0;
/*! \brief The arguments of a C compiler compatible function. */