将您的自定义代码生成器引入 TVM#
随着深度学习应用所针对的硬件设备数量持续增长，用户在各种设备上实现高性能所需的知识也在不断增加。为了使数据科学家在开发新模型时不必担心性能问题，硬件后端供应商要么提供包含许多常用深度学习算子的库，如 DNNL（Intel OneDNN）或 cuDNN，要么提供像 TensorRT 这样的框架，让用户以特定方式描述他们的模型以实现高性能。然而，当用户尝试在新的库或设备上工作时，他们必须学习新的编程接口。因此，对统一编程接口的需求变得越来越重要，原因有二：1）让所有用户和硬件后端供应商站在同一立场上，2）提供可行的解决方案，允许专用硬件或库仅支持广泛使用的高性能算子，同时将不支持的算子回退到 CPU/GPU 等通用设备上。
在本开发者指南中，将向您展示，作为硬件后端提供商，如何轻松实现自己的代码生成器，并将其注册为 Relay 后端编译器，以支持您的硬件设备或库。本指南涵盖基于不同图表示需求的两种代码生成类型：
1. 您希望生成 C 代码。
如果您的硬件已经拥有高度优化的 C/C++ 库，例如针对 CPU 的 Intel CBLAS/MKL 和针对 GPU 的 NVIDIA CUBLAS，那么这正是您所需要的。幸运的是，C 源代码模块与 TVM 运行时模块完全兼容，这意味着生成的代码可以通过任何具有适当编译标志的 C/C++ 编译器进行编译，因此您唯一需要做的就是实现为子图生成 C 代码的代码生成器，以及集成到 TVM 运行时模块中的 C 源代码模块。将在接下来的部分中演示如何为您的硬件实现 C 代码生成器。
2. 您希望生成其他图表示形式。
您的硬件可能需要其他形式的图表示，例如 JSON。在这种情况下，您不仅需要实现代码生成器，还需要实现定制的 TVM 运行时模块，以便让 TVM 运行时知道应如何执行此图表示。如果您已经为您的硬件拥有完整的图执行引擎，例如针对 GPU 的 TensorRT，那么这是您可以考虑的解决方案。
在完成代码生成器和运行时的实现后，您可以让您的客户使用您的自定义标签来标注他们的模型，以便利用这些功能。关于终端用户如何标注并启动特定代码生成器的教程，请参见 here (TBA)。
实现 C 代码生成器#
在这一部分，将演示如何实现代码生成器，该生成器利用预先实现的操作符函数生成 C 代码。为了简化，示例代码生成器不依赖于第三方库。相反，手动在 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]; \
} \
} \
}
通过这两个宏，可以为一维和二维张量生成二元算子。例如，给定如下子图。假设所有输入都是形状为 (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]; \
} \
} \
}
// Note 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);
// Note 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);
}
// Note 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);
重点标注上述代码中的注释：
Note 1 是子图中三个节点的函数实现。
Note 2 是通过分配中间缓冲区并调用相应函数来执行子图的函数。
Note 3 是与 TVM 运行时兼容的包装函数。它接受输入张量列表和输出张量（最后一个参数），将它们转换为正确的数据类型，并调用 Note 2 中描述的子图函数。此外，
TVM_DLL_EXPORT_TYPED_FUNC是 TVM 宏，它通过将所有张量打包到TVMArgs中，生成另一个具有统一函数参数的函数gcc_0。因此，TVM 运行时可以直接调用gcc_0来执行子图，而无需额外的工作。通过生成上述代码，TVM 能够将其与图的其余部分一起编译，并导出单独的库以供部署。
在本节的剩余部分，将逐步实现代码生成器来生成上述代码。您自己的代码生成器必须位于 src/relay/backend/contrib/<您的代码生成器名称>/ 目录下。在示例中，将代码生成器命名为“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 源运行时模块，供 TVM 后端编译和部署。特别地，C 代码的生成对 CodegenC 类是透明的，因为它提供了许多有用的工具来简化代码生成的实现。接下来的部分将按照自底向上的顺序实现这两个类。
实现 CodegenC#
在 src/relay/backend/contrib/codegen_c/codegen.cc 文件中，首先在 tvm.relay.contrib 命名空间下创建代码生成类的框架：
#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. */
std::vector<std::string> ext_func_args_;
/*! \brief The statements of a C compiler compatible function. */
std::vector<std::string> ext_func_body;
/*! \brief The declaration statements of a C compiler compatible function. */
std::vector<std::string> func_decl_;
/*! \brief The declaration statements of buffers. */
std::vector<std::string> buf_decl_;
/*! \brief The name and index pairs for output. */
std::vector<std::pair<std::string, int>> out_;
}
CodegenC 类继承了两个类：ExprVisitor 提供了遍历子图并收集所需信息的能力，以及生成子图函数（例如 gcc_0_）的功能；CodegenCBase 则提供了生成包装函数（如上述示例中的 gcc_0）的能力和实用工具。由此可见，只需在该代码生成类中实现三个函数即可使其正常工作。
算子代码生成#
首先实现 VisitExpr_(const CallNode* call) 函数。该函数在遍历子图时会访问所有的调用节点。每个调用节点包含希望卸载到硬件上执行的算子。因此，需要按照拓扑顺序生成带有正确算子的对应 C 代码。将逐步实现该函数，具体步骤如下。
1. 生成函数声明
示例结果：GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);
为了生成如上所示的函数声明，需要以下信息：1) 函数名称（例如 gcc_0_0），2) 算子类型（例如 *），以及3) 输入张量的形状（例如 (10, 10)）。幸运的是，这些信息可以轻松地从 CallNode 中获取：
std::ostringstream macro_stream;
std::ostringstream decl_stream;
std::ostringstream buf_stream;
// Generate a unique function name you like.
std::string func_name = ext_func_id_ + "_" + std::to_string(func_idx++);
// Make function declaration string.
macro_stream << "CSOURCE_BINARY_OP_" << call->args.size() << "D(" << func_name << ", ";
// Check the operator type.
if (IsOp(call, "add")) {
macro_stream << "+";
} else if (IsOp(call, "subtract")) {
macro_stream << "-";
} else if (IsOp(call, "multiply")) {
macro_stream << "*";
} else {
LOG(FATAL) << "Unrecognized op";
}
// Extract the input tensor shape.
auto in_shape = GetShape(call->args[0]->checked_type());
for (size_t i = 0; i < in_shape.size(); ++i) {
macro_stream << ", " << in_shape[i];
}
macro_stream << ");";
func_decl_.push_back(macro_stream.str());
可以看到，将生成的代码推送到类的成员变量 func_decl_ 中。这意味着在遍历完整个子图后，已经收集了所有所需的函数声明，唯一需要做的就是让 GCC 编译它们。VisitExpr_(const CallNode* call) 的其余实现也遵循这一概念。
2. 生成函数调用
示例结果：gcc_0_0(buf_1, gcc_input3, out);
在生成函数声明后，需要生成具有适当输入和输出的函数调用。为了知道在调用此函数时应放置哪些输入或缓冲区，必须访问其参数：
bool first = true;
decl_stream << func_name << "(";
for (size_t i = 0; i < call->args.size(); ++i) {
VisitExpr(call->args[i]); // Note 1
for (auto out : out_) {
if (!first) {
decl_stream << ", ";
}
first = false;
decl_stream << out.first;
}
}
// Note 2
再次，希望强调上述代码中的注意事项：
Note 1：VisitExpr(call->args[i]) 是递归调用，用于访问当前函数的参数。参数可能是另一个节点的输出或输入张量。在示例实现中，确保每个节点在离开访问者之前更新类变量 out_。以下是说明：
arg_node arg_node <- Visit arg (Note 1) arg_node
| | |
curr_node <- Process curr_node curr_node <- Put "buf_0" as an input buffer
(a) out_ = {} (b) out_ = {} (c) out_ = {("buf_0", 20)}
从上图可以看出，在访问参数节点之前，类变量 out_ 为空，并且它被填充了输出缓冲区名称和 arg_node 的大小。因此，当完成对参数节点的访问时，通过查看 out_，就知道应该放置的正确输入缓冲区。在本节的最后部分，您将发现我们如何更新 out_。
Note 2：您可能会注意到，在这一步中没有关闭函数调用字符串。当前的函数调用字符串看起来像这样：gcc_0_0(buf_1, gcc_input3。这是因为还没有将最后一个参数（即输出）放入此调用中。函数调用的输出可以是分配的临时缓冲区或子图输出张量。为了简化，在此示例中，为每个调用节点分配输出缓冲区（下一步），并将最后缓冲区中的结果复制到输出张量中。
3. 生成输出缓冲区
示例结果：float* buf_0 = (float*)malloc(4 * 100);
正如前一步提到的，除了子图输入和输出张量外，可能还需要缓冲区来保存中间结果。为了生成缓冲区，提取形状信息以确定缓冲区的类型和大小：
// This example only supports single output.
auto type_node = call->checked_type().as<TensorTypeNode>();
ICHECK(type_node != nullptr && runtime::TypeMatch(type_node->dtype, kDLFloat, 32))
<< "Only support single output tensor with float type";
// Generate a unique buffer name.
std::string out = "buf_" + std::to_string(buf_idx_++);
// Extract the shape to be the buffer size.
auto out_shape = GetShape(call->checked_type());
int out_size = 1;
for (size_t i = 0; i < out_shape.size(); ++i) {
out_size *= out_shape[i];
}
// Make the buffer allocation and push to the buffer declarations.
buf_stream << "float* " << out << " = (float*)std::malloc(4 * " << out_size << ");";
buf_decl_.push_back(buf_stream.str());
在分配了输出缓冲区后，现在可以关闭函数调用字符串并将生成的函数调用推送到类变量 ext_func_body 中。
decl_stream << ", " << out << ");";
ext_func_body.push_back(decl_stream.str());
4. 更新输出缓冲区
为了让下一个节点（接受当前调用节点的输出作为其输入）知道它应该使用哪个缓冲区，需要在离开此访问函数之前更新类变量 out_：
out_.clear();
out_.push_back({out, out_size});
恭喜！已经完成了这个类中最困难的功能。在接下来的两节中，只需要补充此函数中一些次要的缺失部分。
输入变量的代码生成#
回想一下，通过访问调用节点的参数（上一节的第2步）收集了输入缓冲区信息，并处理了当参数是另一个调用节点的情况（第4步）。在本节中，以 VarNode 为例，演示如何处理其他节点。
VarNode 表示模型中的输入张量。它唯一但重要的信息是名称提示（例如 data，weight 等）。当访问 VarNode 时，只需更新类变量 out_ 以传递名称提示，以便后代可以生成正确的函数调用。
void VisitExpr_(const VarNode* node) {
ext_func_args_.push_back(node->name_hint());
out_.clear();
out_.push_back({node->name_hint(), 0});
}
请注意，在此示例中，假设卸载的子图仅包含调用节点和变量节点。如果您的子图包含其他类型的节点，例如 TupleNode，那么您还需要访问它们并绕过输出缓冲区信息。
代码 Emitting#
这个代码生成类的最后一部分是 JIT 函数，它为子图生成 C 函数，并使用刚刚生成的 C 代码作为函数体。请记住，除了在前几节中生成的子图函数外，还需要具有统一参数的包装函数，以便 TVM 运行时调用并传递数据。幸运的是，继承的基类已经提供了实现 JitImpl 来生成该函数。例如，可以如下调用 JitImpl：
JitImpl("gcc_0" /* Subgraph symbol (ID) */,
{"gcc_input0", "gcc_input1", "gcc_input2", "gcc_input3"} /* Input arguments */,
{"float *buf_0 = (float*)malloc(4 * 20)", ...} /* Buffer allocations */,
{"gcc_0_2(gcc_input0, gcc_input1, buf_0);"} /* Function body */,
{"out"} /* Output */);
上述调用将生成三个函数（其中一个来自 TVM 包装宏）：
子图函数
gcc_0_（函数名称末尾多一个下划线）包含生成的所有 C 代码，用于执行子图。包装函数
gcc_0__wrapper_带有DLTensor参数列表，将数据转换为正确的类型并调用gcc_0_。TVM 运行时兼容函数
gcc_0具有 TVM 统一函数参数，解包 TVM 打包的张量并调用gcc_0__wrapper_。
因此，在 JIT 实现中，唯一需要做的就是将所有生成的子图函数代码传递给 JitImpl：
std::string JIT() {
// Write function macros
for (auto decl : func_decl_) {
code_stream_ << decl << "\n";
}
return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out_);
}
传递的所有变量（ext_func_id 等）都是类变量，并在遍历子图时填充。
实现 CSourceCodegen#
再次，创建类框架并实现所需的函数。请注意，它继承了 CSourceModuleCodegenBase：
class CSourceCodegen : public CSourceModuleCodegenBase {
public:
// Pass a subgraph function, and generate the C code.
void GenCFunc(const Function& func) { ; }
// Use GenCFunc to generate the C code and wrap it as a C source module.
runtime::Module CreateCSourceModule(const NodeRef& ref) override { ; }
private:
std::ostringstream code_stream_;
};
实现 GenCFunc#
GenCFunc 简单地使用刚刚实现的 CodegenC 来遍历 Relay 函数（子图）并获取生成的 C 代码。内置函数 GetExtSymbol 检索 Relay 函数中的唯一符号名称（例如 gcc_0），并且我们必须将其用作 C 函数名称，因为此符号将用于 DSO 运行时查找。
void GenCFunc(const Function& func) {
ICHECK(func.defined()) << "Input error: expect a Relay function.";
// Record the external symbol for runtime lookup.
auto sid = GetExtSymbol(func);
CodeGenC builder(sid);
builder.VisitExpr(func->body);
code_stream_ << builder.JIT();
}
实现 CreateCSourceModule#
此函数为外部库创建运行时模块。在此示例中，创建 CSourceModule，它可以直接编译并与 TVM 生成的 DSOModule 链接。在您实现 CodegenC 之后，实现此函数相对简单：
runtime::Module CreateCSourceModule(const NodeRef& ref) override {
// Create headers
code_stream_ << "#include <cstdint>\n";
code_stream_ << "#include <iostream>\n";
code_stream_ << "#include <cstdlib>\n";
code_stream_ << "#include <stdio.h>\n";
code_stream_ << "#include <cstring>\n";
code_stream_ << "#include <tvm/runtime/c_runtime_api.h>\n";
code_stream_ << "#include <dlpack/dlpack.h>\n";
// Append some common macro for operator definition.
const char* operator_macro = R"op_macro(
#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]; \
} \
} \
}
)op_macro";
code_stream_ << operator_macro << "\n\n";
// Generate C code for the subgraph.
if (ref->IsInstance<FunctionNode>()) {
GenCFunc(Downcast<Function>(ref));
} else if (ref->IsInstance<relay::ModuleNode>()) {
relay::Module mod = Downcast<relay::Module>(ref);
for (const auto& it : mod->functions) {
GenCFunc(Downcast<Function>(it.second));
}
} else {
LOG(FATAL) << "The input ref is expected to be a Relay function or module"
<< "\n";
}
// Create a CSourceModule
const auto* pf = runtime::Registry::Get("module.csource_module_create");
ICHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
return (*pf)(code_stream_.str(), "cc");
}
注册 Codegen#
最后一步是将您的代码生成器注册到 TVM 后端。首先实现简单的函数来调用代码生成器并生成运行时模块。
runtime::Module CCompiler(const NodeRef& ref) {
CSourceCodegen csource;
return csource.CreateCSourceModule(ref);
}
最后，将此函数注册到 TVM 后端：
TVM_REGISTER_GLOBAL("relay.ext.ccompiler").set_body_typed(CCompiler);
其中 ccompiler 是自定义标签，用于让 TVM 知道这是当子图被标记为 ccompiler 时，它应该用来生成和卸载子图的代码生成器。
最后，好的做法是设置 CMake 配置标志，以便仅为您的客户包含您的编译器。首先创建 cmake 文件：cmake/modules/contrib/CODEGENC.cmake：
if(USE_CODEGENC)
file(GLOB CSOURCE_RELAY_CONTRIB_SRC src/relay/backend/contrib/codegen_c/codegen.cc)
list(APPEND COMPILER_SRCS ${CSOURCE_RELAY_CONTRIB_SRC})
endif(USE_CODEGENC)
这样用户可以在使用 config.cmake 配置 TVM 时选择是否包含您的编译器：
set(USE_CODEGENC ON)
为您的表示实现代码生成器#
尽管已经演示了如何实现 C 代码生成器，但您的硬件可能需要其他形式的图表示，例如 JSON。在这种情况下，您可以修改实现的 CodegenC 类以生成您自己的图表示，并实现自定义的运行时模块，以让 TVM 运行时知道应如何执行此图表示。
为了简化，在本指南中定义了名为“ExampleJSON”的图表示。ExampleJSON 并不是指真正的 JSON，而只是没有控制流的图的简单表示。例如，假设有名为 subgraph_0 的子图：
input0
|
add <-- input1
|
subtract <-- input2
|
multiply <-- input3
|
out
那么此子图的 ExampleJSON 如下所示：
subgraph_0
input 0 10 10
input 1 10 10
input 2 10 10
input 3 10 10
add 4 inputs: 0 1 shape: 10 10
sub 5 inputs: 4 2 shape: 10 10
mul 6 inputs: 5 3 shape: 10 10
input 关键字声明了输入张量及其 ID 和形状；而其他语句以 <op> <output ID> inputs: [input ID] shape: [shape] 语法描述计算。
在本节中，目标是实现以下自定义 TVM 运行时模块来执行 ExampleJSON 图。
runtime::Module ExampleJsonCompiler(const NodeRef& ref) {
ExampleJsonCodeGen codegen(ref);
std::string code = codegen.gen(); // Note 1
const auto* pf = runtime::Registry::Get("module.examplejson_module_create"); // Note 2
ICHECK(pf != nullptr) << "Cannot find ExampleJson module to create the external runtime module";
return (*pf)(code);
}
TVM_REGISTER_GLOBAL("relay.ext.examplejsoncompiler").set_body_typed(ExampleJsonCompiler);
Note 1：稍后将实现自定义代码生成器，通过获取子图来生成 ExampleJSON 代码字符串。
Note 2：此行获取指向用于创建自定义运行时模块的函数的指针。您可以看到，它接受刚刚生成的 ExampleJSON 格式的子图代码并初始化运行时模块。
在接下来的部分中，将介绍1）如何实现 ExampleJsonCodeGen，以及2）如何实现并注册 examplejson_module_create。
实现 ExampleJsonCodeGen#
与 C 代码生成器类似，也将 ExampleJsonCodeGen 从 ExprVisitor 派生，以利用访问者模式进行子图遍历。另一方面，不需要继承 CodegenCBase，因为不需要 TVM C++ 包装器。代码生成类的实现如下：
#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>
namespace tvm {
namespace relay {
namespace contrib {
class ExampleJsonCodeGen : public ExprVisitor {
public:
explicit ExampleJsonCodeGen();
// Note 1
void VisitExpr_(const VarNode* node) { /* Skip in this example. */ }
void VisitExpr_(const CallNode* call) final { /* Skip in this example. */ }
// Note 2
std::string gen(NodeRef& ref) {
this->code = "";
if (ref->IsInstance<FunctionNode>()) {
this->visit(Downcast<Function>(ref));
} else if (ref->IsInstance<relay::ModuleNode>()) {
relay::Module mod = Downcast<relay::Module>(ref);
for (const auto& it : mod->functions) {
this->visit(Downcast<Function>(it.second));
}
} else {
LOG(FATAL) << "The input ref is expected to be a Relay function or module";
}
return this->code;
}
private:
/*! \brief The function id that represents a C source function. */
std::string code;
}
Note 1：再次实现了相应的访问者函数来生成 ExampleJSON 代码，并将其存储到类变量 code 中（在此示例中跳过了访问者函数的实现，因为它们的概念与 C 代码生成基本相同）。完成图的访问后，应该在 code 中得到 ExampleJSON 图。
Note 2：定义了内部 API gen，用于接收子图并生成 ExampleJSON 代码。此 API 的名称可以根据您的喜好任意命名。
接下来的步骤是实现定制化的运行时，以利用 ExampleJsonCodeGen 的输出。
实现定制化的运行时#
在本节中，将逐步实现定制的 TVM 运行时，并将其注册到 TVM 运行时模块中。定制的运行时应位于 src/runtime/contrib/<your-runtime-name>/ 目录下。在示例中，将运行时命名为“example_ext_runtime”。
同样，首先定义定制化的运行时类，如下所示。该类必须继承自 TVM 的 ModuleNode，以便与其他 TVM 运行时模块兼容。
#include <dmlc/logging.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/memory.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
#include <fstream>
#include <cmath>
#include <map>
#include <sstream>
#include <string>
#include <vector>
namespace tvm {
namespace runtime {
class ExampleJsonModule : public ModuleNode {
public:
explicit ExampleJsonModule(std::string graph_json);
PackedFunc GetFunction(const std::string& name,
const ObjectPtr<Object>& sptr_to_self) final;
const char* type_key() const { return "examplejson"; }
void SaveToBinary(dmlc::Stream* stream) final;
static Module LoadFromBinary(void* strm);
static Module Create(const std::string& path);
std::string GetSource(const std::string& format = "");
void Run(int id, const std::vector<int>& inputs, int output);
void ParseJson(const std::string& json);
private:
/* \brief The json string that represents a computational graph. */
std::string graph_json_;
/* \brief The subgraph that being processed. */
std::string curr_subgraph_;
/*! \brief A simple graph from subgraph id to node entries. */
std::map<std::string, std::vector<NodeEntry>> graph_;
/* \brief A simple pool to contain the tensor for each node in the graph. */
std::vector<NDArray> data_entry_;
/* \brief A mapping from node id to op name. */
std::vector<std::string> op_id_;
};
具体来说，有一些从 ModuleNode 继承的函数，必须在 ExampleJsonModule 中实现：
构造函数：该类的构造函数应接受子图（以您的表示形式），处理并以您喜欢的任何格式存储它。保存的子图可以被以下两个函数使用。
GetFunction：这是该类中最重要的函数。当 TVM 运行时希望使用您的编译器标签执行子图时，TVM 运行时会从您的定制运行时模块中调用此函数。它提供了函数名称以及运行时参数，而GetFunction应返回打包的函数实现，供 TVM 运行时执行。SaveToBinary和LoadFromBinary：SaveToBinary将运行时模块序列化为二进制格式，以便后续部署。当用户使用export_libraryAPI 时，TVM 会调用此函数。另一方面，由于现在使用自己的图表示形式，必须确保LoadFromBinary能够通过读取SaveToBinary生成的序列化二进制文件，构建相同的运行时模块。GetSource（可选）：如果您希望查看生成的 ExampleJSON 代码，可以实现此函数以将其导出；否则，您可以跳过此实现。
其他函数和类变量将随着上述必需函数的实现一起引入。
实现构造函数#
explicit ExampleJsonModule(std::string graph_json) {
this->graph_json_ = graph_json;
ParseJson(this->graph_json_);
}
接下来，实现 ParseJson 来解析 ExampleJSON 格式的子图，并在内存中构建图以供后续使用。由于在此示例中不支持带有分支的子图，因此简单地使用数组按顺序存储子图中的每个节点。
void ParseJson(const std::string& json) {
std::string line;
std::string curr_subgraph;
std::stringstream ss(json);
while (std::getline(ss, line, '\n')) {
std::stringstream ss2(line);
std::string token;
int id = 0;
ss2 >> token;
if (token.find("subgraph_") != std::string::npos) {
curr_subgraph = token;
continue;
}
ss2 >> id;
if (op_id_.size() <= static_cast<size_t>(id)) {
op_id_.resize(id + 1);
data_entry_.resize(id + 1);
}
int64_t total_elements = 1;
std::vector<int64_t> shape;
if (token == "input") {
int64_t size = 0;
while (ss2 >> size) {
total_elements *= size;
shape.push_back(size);
}
} else {
op_id_[id] = token; // Note 1
bool shape_data = false;
NodeEntry entry;
while (ss2 >> token) {
if (token == "shape:") {
shape_data = true;
} else if (shape_data) {
total_elements *= std::stoll(token);
shape.push_back(std::stoll(token));
} else if (token != "inputs:") {
entry.inputs.push_back(std::stoi(token));
}
}
entry.id = id;
entry.output = id;
graph_[curr_subgraph].push_back(entry); // Note 2
}
DLDevice dev;
dev.device_type = static_cast<DLDeviceType>(1);
dev.device_id = 0;
data_entry_[id] = NDArray::Empty(shape, DLDataType{kDLFloat, 32, 1}, dev); // Note 3
}
}
Note 1：使用类变量 op_id_ 来将子图节点 ID 映射到算子名称（例如，add），以便可以在运行时调用相应的算子函数。
Note 2：使用类变量 graph_ 来将子图名称映射到节点数组。GetFunction 将在运行时通过子图 ID 查询图节点。
Note 3：使用类变量 data_entry_ 来将子图节点 ID 映射到张量数据占位符。将在运行时将输入和输出放入相应的数据条目中。
实现 GetFunction#
在构建完成后，应该已经准备好上述类变量。接着，实现 GetFunction，以向 TVM 运行时提供可执行的子图函数：
PackedFunc GetFunction(const std::string& name,
const ObjectPtr<Object>& sptr_to_self) final {
if (this->graph_.find(name) != this->graph_.end()) {
this->curr_subgraph_ = name;
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
// Copy input tensors to corresponding data entries.
for (auto i = 0; i < args.size(); ++i) {
ICHECK(args[i].type_code() == kNDArrayContainer || args[i].type_code() == kArrayHandle)
<< "Expect NDArray or DLTensor as inputs\n";
if (args[i].type_code() == kArrayHandle) {
DLTensor* arg = args[i];
this->data_entry_[i].CopyFrom(arg);
} else {
NDArray arg = args[i];
this->data_entry_[i].CopyFrom(arg);
}
}
// Execute the subgraph.
for (const auto& it : this->graph_[this->curr_subgraph_]) {
this->Run(it.id, it.inputs, it.output);
}
ICHECK_GT(graph_.count(this->curr_subgraph_), 0U);
// Copy the output from a data entry back to TVM runtime argument.
auto out_idx = graph_[this->curr_subgraph_].back().output;
if (args[args.size() - 1].type_code() == kArrayHandle) {
DLTensor* arg = args[args.size() - 1];
this->data_entry_[out_idx].CopyTo(arg);
} else {
NDArray arg = args[args.size() - 1];
this->data_entry_[out_idx].CopyTo(arg);
}
*rv = data_entry_.back();
});
} else {
LOG(FATAL) << "Unknown subgraph: " << name << "\n";
return PackedFunc();
}
}
可以看出，GetFunction 由三个主要部分组成。第一部分将数据从 TVM 运行时参数复制到在构造函数中分配的相应数据条目中。第二部分使用 Run 函数（稍后将实现）执行子图，并将结果保存到另一个数据条目中。第三部分将结果从输出数据条目复制回相应的 TVM 运行时参数以进行输出。
实现 Run#
现在实现 Run 函数。该函数接受1）子图 ID，2）输入数据条目索引列表，以及3）输出数据条目索引。
void Run(int id, const std::vector<int>& inputs, int output) {
// Make a list data entry indexs.
std::vector<int> args(inputs.begin(), inputs.end());
args.push_back(output);
// Initialize data holders.
std::vector<TVMValue> values(args.size());
std::vector<int> type_codes(args.size());
// Initialize a TVM arg setter with TVMValue and its type code.
TVMArgsSetter setter(values.data(), type_codes.data());
// Set each argument to its corresponding data entry.
if (op_id_[id] == "add" || op_id_[id] == "sub" || op_id_[id] == "mul") {
for (size_t i = 0; i < args.size(); i++) {
setter(i, data_entry_[args[i]]);
}
}
// Invoke the corresponding operator function.
if (op_id_[id] == "add") {
Add(values.data(), type_codes.data(), args.size());
} else if (op_id_[id] == "sub") {
Sub(values.data(), type_codes.data(), args.size());
} else if (op_id_[id] == "mul") {
Mul(values.data(), type_codes.data(), args.size());
} else {
LOG(FATAL) << "Unknown op: " << op_id_[id] << "\n";
}
}
Run 函数主要有两个部分。第一部分分配 TVMValue 列表，并映射相应的数据条目块。这将作为算子函数的参数。第二部分随后调用算子函数。尽管使用与之前相同的 C 函数，您可以用自己的引擎替换 Add、Sub 和 Mul。您只需要确保您的引擎将结果存储到最后一个参数中，以便它们可以传输回 TVM 运行时。
通过实现上述函数，定制代码生成器和运行时现在可以执行子图。最后一步是注册API（examplejson_module_create）来创建此模块：
TVM_REGISTER_GLOBAL("module.examplejson_module_create")
.set_body_typed([](std::string code){
auto n = make_object<ExampleJsonModule>(code);
return runtime::Module(n);
});
实现 SaveToBinary 和 LoadFromBinary#
到目前为止，已经实现了定制运行时的主要功能，以便它可以像其他 TVM 运行时一样使用。然而，当用户希望将构建的运行时保存到磁盘以进行部署时，TVM 并不知道如何保存它。这就是为什么要实现 SaveToBinary 和 LoadFromBinary，它们告诉 TVM 应该如何持久化和恢复这个定制的运行时。
首先实现 SaveToBinary 函数，以允许用户将此模块保存到磁盘中。
void SaveToBinary(dmlc::Stream* stream) final {
stream->Write(this->graph_json_);
}
可以发现这个函数非常简单。回想一下，在构造函数中接受的唯一参数是子图表示，这意味着只需要子图表示来构建/恢复这个定制的运行时模块。因此，SaveToBinary 简单地将子图写入输出 DMLC 流。也就是说，当用户使用 export_library API导出模块时，定制的模块将是子图的 ExampleJSON 流。
类似地，LoadFromBinary 读取子图流并重新构建定制的运行时模块：
static Module LoadFromBinary(void* strm) {
dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
std::string graph_json;
stream->Read(&graph_json);
auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);
return Module(n);
}
还需要注册此函数以启用相应的 Python API：
TVM_REGISTER_GLOBAL("module.loadbinary_examplejson")
.set_body_typed(ExampleJsonModule::LoadFromBinary);
上述注册意味着当用户调用 tvm.runtime.load_module(lib_path) API 并且导出的库包含 ExampleJSON 流时，LoadFromBinary 将被调用来创建相同的定制运行时模块。
此外，如果您希望支持直接从 ExampleJSON 文件创建模块，您还可以实现简单的函数并注册 Python API，如下所示：
static Module Create(const std::string& path) {
std::ifstream filep;
filep.open(path, std::ios::in);
std::string graph_json;
std::string line;
while (std::getline(filep, line)) {
graph_json += line;
graph_json += "\n";
}
filep.close();
auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);
return Module(n);
}
TVM_REGISTER_GLOBAL("module.loadfile_examplejson")
.set_body([](TVMArgs args, TVMRetValue* rv) {
*rv = ExampleJsonModule::Create(args[0]);
});
这意味着用户可以手动编写/修改 ExampleJSON 文件，并使用 Python API tvm.runtime.load_module("mysubgraph.examplejson", "examplejson") 来构建定制模块。
小结#
总结一下，以下是供您参考的清单：
从
ExprVisitor和CodegenCBase（仅适用于 C 代码生成）派生的代码生成类，包含以下函数。VisitExpr_(const CallNode* call)用于收集调用节点信息。您需要收集子图信息的其他访问者函数。
JIT用于生成子图代码。注册代码生成器。
用于创建
CSourceModule的函数（适用于 C 代码生成）。从
ModuleNode派生的运行时模块类，包含以下函数（适用于您的图表示）。Constructor.
GetFunction用于生成与 TVM 运行时兼容的PackedFunc。Run用于执行子图。注册运行时创建 API。
SaveToBinary和LoadFromBinary用于序列化/反序列化定制的运行时模块。注册
LoadFromBinaryAPI 以支持tvm.runtime.load_module(your_module_lib_path)。（可选）
Create用于支持从子图文件中构建定制的运行时模块。
annotator，用于注解用户的 Relay 程序以利用您的编译器和运行时（待添加）。