TVM代码生成codegen
TVM代码生成codegen
硬件后端提供程序（例如Intel，NVIDIA，ARM等），提供诸如cuBLAS或cuDNN之类的内核库以及许多常用的深度学习内核，或者提供框架例，如带有图形引擎的DNNL或TensorRT，使用户以某种方式描述模型，实现高性能。此外，新兴的深度学习加速器还具有自己的编译器，内核库或runtime框架。
当用户尝试在新的内核库或设备上工作时，必须学习新的编程接口。对统一编程接口的需求变得越来越重要，使所有用户和硬件后端提供程序都在同一页面上。
为了与广泛使用的深度学习框架共享编程接口，许多硬件设备提供商已尝试将其设备后端集成到TensorFlow。由于TensorFlow没有为新的后端提供正式的后端接口，必须破解TensorFlow进行注册，这需要对许多源文件进行更改，从而使将来的维护变得困难。
本文演示了作为硬件后端提供程序，如何轻松利用自带代码生成（BYOC）框架，将硬件设备的内核库/编译器/框架集成到TVM。利用BYOC框架的最重要优点，设备的所有相关源文件都是独立的，设备的代码源/Runtime可插入TVM代码库。这意味着
1）使用代码源的TVM代码库将在上游兼容
2）TVM用户可以根据需要选择启用代码源/runtime。
在本文的其余部分中，首先说明可能需要带有BYOC的TVM的情况，然后概述BYOC编译和runtime流程。然后，分步说明如何使用英特尔DNNL（又名MKL-DNN，OneDNN），作为运行示例，将供应商库或执行引擎与BYOC集成到TVM。
将ASIC加速器带入TVM
首先，做一个场景来说明，为什么要将加速器引入TVM，以及BYOC框架可以期待哪些功能。
想象一下，刚刚构建了一个具有ARM CPU和出色的加速器的边缘设备平台，该平台为常见的图像分类模型，提供了出色的性能。换句话说，加速器在Conv2D，ReLU，GEMM和其他广泛使用的CNN算子上表现良好。
不幸的是，目标检测模型也越来越受欢迎，并且客户需要在平台上同时运行图像分类和目标检测模型。尽管加速器能够执行目标检测模型中的几乎所有算子，但缺少一个算子（例如，非最大抑制，NMS）。
让TVM执行不受支持的算子
由于TVM具有用于不同后端的多个代码源，开源社区很容易在短时间内在CPU或GPU上实现新的算子。理想情况下，如果将加速器的编译流程与BYOC集成到TVM，TVM将执行Relay图分区，以将部分图卸载到加速器，而其它图保持在TVM上。因此，可以申明平台能够运行所有模型，而不必担心新的算子。
自定义图形级优化
ASIC加速器必须具有编译流程。可能是以下情况之一：
生成图形表示并将其提供给图形引擎：可能拥有图形引擎，该引擎能够在加速器上执行图形（或神经网络模型）。例如，英特尔DNNL和NVIDIA TensorRT都使用引擎来运行整个图形或模型，因此能够1）减少算子之间的内存事务，以及2）通过算子融合优化图形执行。
为了实现以上两个优化，需要在编译期间处理图形。例如，Conv2D和偏差加法是TVM中的两个单独的算子，可能是加速器上的一个算子（具有偏差加法功能的Conv2D）。在这种情况下，需要通过将conv2d - add图形模式替换为your_conv2d_with_bias节点来优化图形。
如果编译流程属于这种情况，建议阅读本文中的所有其余部分，但跳过将DNNL带到TVM：C源代码生成。
生成汇编代码并将其编译为可执行的二进制文件：如果没有像前面那样的平台的端到端执行框架，则可能有编译器以ISA的汇编代码编译程序。为了将汇编代码提供给编译器，将需要一个代码生成器来从Relay图生成和优化汇编代码。
如果编译流程属于这种情况，建议阅读本文中的所有其余部分，但跳过将DNNL引入TVM：JSON Codegen / Runtime。
BYOC的工作方式
简要解释BYOC框架是如何工作的。简而言之，给定图1中的Relay图，BYOC框架执行以下步骤：
图1：原始Relay图。
1.图注解
制作用户提供的Relay图，第一步是在图中注释可能卸载到加速器的节点。需要遵循“将DNNL引入TVM：注释规则”，实现受支持的算子的白名单，或定制组合算子的图形模式列表。示例注释结果如图2所示。
图2：带注解的图。
2.图变换
第二步是基于注释对图形进行转换和优化。具体来说，BYOC执行以下转换。
2.1：合并编译器区域：如图2所示，图中现在有许多“区域”可以卸载到加速器中，实际上可以合并其中的一些区域，减少数据传输和内核启动开销。因此，步骤2.1使用贪婪算法来合并尽可能多的那些区域，同时保证功能正确性。结果如图3所示。
图3：合并编译器区域后。
2.2：分区图：对于上一步中的每个区域，创建一个带有属性的Relay函数，Compiler以指示该Relay函数应该完全卸载到加速器上，如图4所示。
图4：图分区之后。
3.代码生成
现在知道应该卸载Relay图的哪一部分了。将每个Relay功能依次发送Compiler=your_accelerator到代码生成器。代码生成器应将Relay函数编译为与编译流程相匹配的形式。可以是C源代码或任何文本格式。
最后，所有已编译的函数将与其它未卸载的Relay函数一起.so由TVM export_libraryPython API序列化为单个文件。换句话说，.so运行此流程后，用户将仅获得一个文件。
4.runtime
需要实现Runtime以初始化图形引擎（如果适用）并执行已编译的函数。在推理期间，当TVMRuntime遇到图4中的相应函数调用时，TVM Runtime（即图形Runtime或VM）将利用Runtime来调用已卸载的函数。Runtime负责使用给定的输入张量启动编译后的函数。数组并将结果填充到输出张量数组中。
在本文的其余部分，以DNNL为例，演示如何使用BYOC框架实现上述工作流程。本文中所有引用的代码和行号均基于TVM存储库的master分支commit 8a0249c。
将DNNL带到TVM：注释规则
BYOC框架提供了两种描述受支持的算子和模式的方法，可以同时使用。以DNNL为例，说明如何使用。将代码源的注释规则放在下python/tvm/relay/op/contrib/your_codegen_name.py。
单一运营商规则
可以使用BYOC API直观地指定加速器支持哪些Relay算子。例如，使用以下代码段构建一条规则，该规则说DNNL代码源支持Conv2D：
@tvm.ir.register_op_attr("nn.conv2d", "target.dnnl")
def _dnnl_conv2d_wrapper(attrs, args):
return True
这target.dnnl将向Relaynn.conv2d算子注册一个新属性。通过这种方式，BYOC注释可以target.dnnl()为图中的每个算子调用以检查DNNL代码源中是否支持。
另一方面，为每个算子编写上面的代码段可能很繁琐。对于DNNL实施，实现了一个辅助函数_register_external_op_helper，更简洁：
def _register_external_op_helper(op_name, supported=True):
@tvm.ir.register_op_attr(op_name, "target.dnnl")
def _func_wrapper(attrs, args):
return supported
return _func_wrapper
_register_external_op_helper("nn.batch_norm")
_register_external_op_helper("nn.conv2d")
_register_external_op_helper("nn.dense")
_register_external_op_helper("nn.relu")
_register_external_op_helper("add")
_register_external_op_helper("subtract")
_register_external_op_helper("multiply")
在上面的示例中，指定了DNNL代码源可以支持的算子列表。
图形模式规则
加速器或编译器可能已将某些模式（例如Conv2D + add + ReLU）优化为单个指令或API。在这种情况下，可以指定从图形模式到指令/ API的映射。对于DNNL，Conv2D API已经包含了偏差加法，并且允许连接下一个ReLU，可以将DNNL称为以下代码片段：
DNNLConv2d(const bool has_bias = false, const bool has_relu = false) {
// ... skip ...
auto conv_desc = dnnl::convolution_forward::desc(
dnnl::prop_kind::forward_inference,
dnnl::algorithm::convolution_direct,
conv_src_md, conv_weights_md, conv_bias_md, conv_dst_md,
strides_dims, padding_dims_l, padding_dims_r);
// Attach ReLU
dnnl::primitive_attr attr;
if (has_relu) {
dnnl::post_ops ops;
ops.append_eltwise(1.f, dnnl::algorithm::eltwise_relu, 0.f, 0.f);
attr.set_post_ops(ops);
}
auto conv2d_prim_desc = dnnl::convolution_forward::primitive_desc(
conv_desc, attr, engine_);
// ... skip ...
在这种情况下，除了用于单个conv2d，想映射图模式conv2d+relu到DNNLConv2d(false, true)，并映射conv2d+add+relu到DNNLConv2d(true, true)。可以使用以下代码片段实现此目的：
def make_pattern(with_bias=True):
data = wildcard()
weight = wildcard()
bias = wildcard()
conv = is_op('nn.conv2d')(data, weight)
if with_bias:
conv_out = is_op('add')(conv, bias)
else:
conv_out = conv
return is_op('nn.relu')(conv_out)
@register_pattern_table("dnnl")
def pattern_table():
conv2d_bias_relu_pat = ("dnnl.conv2d_bias_relu", make_pattern(with_bias=True))
conv2d_relu_pat = ("dnnl.conv2d_relu", make_pattern(with_bias=False))
dnnl_patterns = [conv2d_bias_relu_pat, conv2d_relu_pat]
return dnnl_patterns
在DNNL示例中，实现了两个具有不同名称的模式，以便可以轻松地在代码生成中识别。注意，这些模式以Relay模式语言实现。
使用模式表，然后可以使用从Relay传递来执行
%1 = nn.conv2d(%data, %weight, ...)
%2 = add(%1, %bias)
%3 = nn.relu(%2)
至
%1 = fn(%input1, %input2, %input3,
Composite="dnnl.conv2d_bias_relu",
PartitionedFromPattern="nn.conv2d_add_nn.relu_") {
%1 = nn.conv2d(%input1, %input2, ...)
%2 = add(%1, %input3)
nn.relu(%2)
}
%2 = %1(%data, %weight, %bias)
因此，DNNL代码生成器可以获取模式名称conv2d_bias_relu并映射%1到DNNLConv2d(true, true)。
复合函数中还有一个名为“ PartitionedFromPattern”的属性。如果模式包含wildcard算子，这可能会有所帮助。例如，可能有一个模式表("conv2d_with_something", conv2d -> *)：
def make_pattern(with_bias=True):
data = wildcard()
weight = wildcard()
conv = is_op('nn.conv2d')(data, weight)
return wildcard()(conv)
在这种情况下，将获得带有的复合函数Composite=conv2d_with_something，但是不知道实际匹配的图形。那就是PartitionedFromPattern起作用的地方。通过查看匹配图是否为conv2d -> add或conv2d -> relu，可以知道是否PartitionedFromPattern为nn.conv2d_add_或nn.conv2d_nn.relu_。
将DNNL引入TVM：Relay图转换
利用上一步中的注释规则，现在可以应用BYOCRelay传递列表，以将Relay图从图1转换为图4：
mod = create_relay_module_from_model() # Output: Figure 1
mod = transform.MergeComposite(pattern_table)(mod)
mod = transform.AnnotateTarget(["dnnl"])(mod) # Output: Figure 2
mod = transform.MergeCompilerRegions()(mod) # Output: Figure 3
mod = transform.PartitionGraph()(mod) # Output: Figure 4
可以看出，每个Relay传递都可以映射到在BYOC工作原理中引入的步骤。
将DNNL引入TVM：JSON代码生成/Runtime
现在，实现将Relay图序列化为JSON表示的DNNL代码源，然后实现DNNL JSONRuntime以反序列化并执行该图。如果尝试实现一个代码生成器来生成C兼容程序，则可能需要直接进入下一部分。
为了使DNNL JSON的代码生成/运行在TVM就这个例子中工作，确保DNNL可以在机器上，并与建立TVMset(USE_DNNL_CODEGEN ON)中config.cmake。
DNNL代码生成在中实现src/relay/backend/contrib/dnnl/codegen.cc。在此文件中以两种形式实现了DNNLUSE_JSON_RUNTIME代码生成，在跟踪代码时，可以专注于宏所覆盖的部分。
首先使用TVM注册API（L510），注册代码源。该注册使TVM编译引擎使用Compiler=<your codegen> 来分发Relay功能relay.ext.<your codegen>。然后，实现DNNL编译器（L490）的入口函数。阅读代码段中嵌入的注释以获取详细信息：
runtime::Module DNNLCompiler(const ObjectRef& ref) {
// "ref" should be the paritioned Relay function with kCompiler=dnnl.
CHECK(ref->IsInstance<FunctionNode>());
auto func = Downcast<Function>(ref);
// Get the function name as the symbol to match in runtime.
auto func_name = GetExtSymbol(func);
// Serialize the function to a JSON string (introduce later).
DNNLJSONSerializer serializer(func_name, func);
serializer.serialize();
std::string graph_json = serializer.GetJSON();
// The constant tensor names that have been bound to the module.
// All constant tensors will be serialzied along with the JSON graph
// when export_library is invoked.
auto params = serializer.GetParams();
// The function to create DNNL JSON runtime (introduce later).
const auto* pf = runtime::Registry::Get("runtime.DNNLJSONRuntimeCreate");
CHECK(pf != nullptr) << "Cannot find JSON runtime module to create";
// Create a DNNL runtime module that can run the serialized function.
auto mod = (*pf)(func_name, graph_json, params);
return mod;
}
TVM_REGISTER_GLOBAL("relay.ext.dnnl").set_body_typed(DNNLCompiler);
注意，每个Runtime模块仅负责一个Relay功能，这意味着可能在单个.so文件中包含多个DNNLRuntime模块。
DNNL JSON序列化
接下来，实现DNNL JSON序列化器（L429）。从BYOC JSON代码生成器（src / relay / backend / contrib / codegen_json / codegen_json.h）派生了它。DNNL JSON序列化程序中的特殊过程尝试，将组合函数调用序列化为DNNL JSON Runtime，可以解释的JSON节点。假设有一个与pattern匹配的复合函数dnnl.conv2d_relu，那么BYOC JSON代码生成器将生成以下JSON节点：
{
op: "kernel",
name: "dnnl.conv2d_relu",
inputs: [[0, 0, 0], [1, 0, 0]],
attrs: {
PartitionedFromPattern: ["nn.conv2d_nn.relu_"],
shape: [1, 32, 14, 14]
}
}
问题在于，在Runtime仍然需要Conv2D属性，例如padding和stride，但是BYOC JSON序列化器仅附加复合函数的属性，而不附加主体算子。另一方面，定制的DNNL JSON序列化程序将第一个也是唯一的Conv2D的属性附加到复合函数中，以生成以下JSON节点：
{
op: "kernel",
name: "dnnl.conv2d_relu",
inputs: [[0, 0, 0], [1, 0, 0]],
attrs: {
shape: [1, 32, 14, 14],
data_layout: ["NCHW"],
kernel_layout: ["OIHW"],
strides: [1, 1],
padding: [1, 1, 1, 1]
}
}
从DNNL JSON序列化器可以看出，可以自定义序列化器以生成JSON中的任何形式，只要JSON Runtime可以解释它们即可。
DNNL JSON Runtime
然后，实现DNNL JSON Runtime以解释和执行序列化的JSON图。放在下面src/runtime/contrib/dnnl/dnnl_json_runtime.cc。
同样，首先注册两个API来创建Runtime，以便可以在任何地方使用。在runtime.DNNLJSONRuntimeCreate被序列化后的上一部分中使用，并且runtime.module.loadbinary_dnnl_json装载时也可以使用.so了。
// Create a DNNL JSON runtime to interpret and execute the given JSON graph.
runtime::Module DNNLJSONRuntimeCreate(String symbol_name, String graph_json,
const Array<String>& const_names) {
auto n = make_object<DNNLJSONRuntime>(symbol_name, graph_json, const_names);
return runtime::Module(n);
}
TVM_REGISTER_GLOBAL("runtime.DNNLJSONRuntimeCreate")
.set_body_typed(DNNLJSONRuntimeCreate);
TVM_REGISTER_GLOBAL("runtime.module.loadbinary_dnnl_json")
.set_body_typed(JSONRuntimeBase::LoadFromBinary<DNNLJSONRuntime>);
现在，解释DNNL JSON Runtime实现。基本的类结构为：
class DNNLJSONRuntime : public JSONRuntimeBase {
const char* type_key() const { return "dnnl_json"; }
void Init(const Array<NDArray>& consts) override {
// Initialize the DNNL graph engine.
BuildEngine();
// Setup constants entries for weights.
CHECK_EQ(consts.size(), const_idx_.size())
<< "The number of input constants must match the number of required.";
SetupConstants(consts);
}
void Run() override {
// 1. Fill in the input buffers.
// 2. Invoke the engine through intepreting the stream.
// 3. Read and fill output buffers.
}
}
该Init功能是负责通过解释JSON图形字符串，建设DNNL引擎（见L93的BuildEngine），并填补了固定的权重，以相应的数据输入缓冲区（SetupConstant在JSON运行基类来实现，需要调用它在Init）。注意，即使运行了多次推理，该函数也只会被调用一次。
接下来，Run函数（L64）首先将输入张量（可能来自用户输入或恒定权重）写入在构建DNNL引擎时初始化的相应DNNL存储缓冲区。然后启动DNNL引擎以执行JSON图。最后，将DNNL输出存储缓冲区写回到相应的输出张量。
由于DNNL JSONRuntime中的其余实现都是DNNL特有的，不再细说。想强调一点，尽管DNNL JSONRuntime是一个很好的开始，但JSON Runtime可以完全自定义以满足要求。
将DNNL带到TVM：C源代码生成
现在，让实现DNNL代码生成器，该代码生成器生成C源代码，该源代码调用DNNL API来执行Relay图。注意，如果尝试实现一个代码生成器，生成其它图形表示形式（如JSON格式），则可能需要阅读将DNNL带到TVM：JSON代码生成器/Runtime，并跳过本节。
为了能够在TVM CODEGEN对这个例子的工作DNNL C源代码，确保DNNL可以在机器上，并与建立TVMset(USE_DNNL_CODEGEN C_SRC)中config.cmake。
DNNL代码生成在中实现src/relay/backend/contrib/dnnl/codegen.cc。由于在这个文件用于说明目的实现的代码，生成DNNL两种形式，可以专注于部分不被覆盖USE_JSON_RUNTIME宏跟踪代码时。
首先，使用TVM注册API（L510）注册代码源。该注册使TVM编译引擎使用Compiler=<your codegen> 来分发Relay功能relay.ext.<your codegen>。然后，实现DNNL编译器的入口函数（L490）：
runtime::Module DNNLCompiler(const ObjectRef& ref) {
DNNLModuleCodegen dnnl;
return dnnl.CreateCSourceModule(ref);
}
TVM_REGISTER_GLOBAL("relay.ext.dnnl").set_body_typed(DNNLCompiler);
注意，每个Runtime模块仅负责一个Relay功能，这意味着可能在单个.so文件中包含多个DNNL Runtime模块。
然后，在L362中派生CSourceModuleCodegenBase实施。而负责其它模块级过程，如序列化的，只需要实现在所述DNNL代码生成函数（L389）：DNNLModuleCodegenCSourceModuleCodegenBaseCreateCSourceModule
runtime::Module CreateCSourceModule(const ObjectRef& ref) override {
// Include headers
// ...skip...
code_stream_ << "#include <dnnl/dnnl_kernel.h>\n";
// ...skip...
// "ref" should be the paritioned Relay function with kCompiler=dnnl.
CHECK(ref->IsInstance<FunctionNode>());
auto res = GenDNNLFunc(Downcast<Function>(ref));
// "code" is the generated C code with DNNL APIs.
std::string code = code_stream_.str();
// "res" is a tuple of constant weights (symbols, values).
// All constant tensors will be serialzied along with the generated C code
// when export_library is invoked.
String sym = std::get<0>(res);
Array<String> variables = std::get<1>(res);
// Create a CSource module with all above artifacts.
const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate");
CHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
return (*pf)(code, "c", sym, variables);
}
接下来，实现GenDNNLFunc（L365）来使用DNNL API生成可编译的C代码，如下所示。参阅嵌入的注释，以获取与TVM C源Runtime模块兼容的功能接口的说明。
// The example Relay graph: conv2d -> add -> relu.
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/container.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include <dnnl/dnnl_kernel.h>
using namespace tvm::runtime;
using namespace tvm::runtime::contrib;
// Execute the conv2d->add->relu graph with DNNL.
extern "C" void dnnl_0_(float* dnnl_0_i0, float* dnnl_0_i1,
float* dnnl_0_i2, float* out0) {
// Allocate intermediate buffers.
float* buf_0 = (float*)std::malloc(4 * 4608);
float* buf_1 = (float*)std::malloc(4 * 4608);
float* buf_2 = (float*)std::malloc(4 * 4608);
// Pre-implemented op-based DNNL functions.
dnnl_conv2d(dnnl_0_i0, dnnl_0_i1, buf_0, 1, 32, 14, 14, 32, 1, 0, 0, 3, 3, 1, 1);
dnnl_add(buf_0, dnnl_0_i2, buf_1, 1, 32, 12, 12);
dnnl_relu(buf_1, buf_2, 1, 32, 12, 12);
// Copy the final output to the corresponding buffer.
std::memcpy(out0, buf_2, 4 * 4608);
std::free(buf_0);
std::free(buf_1);
std::free(buf_2);
}
// The wrapper function with all arguments in DLTensor type.
extern "C" int dnnl_0_wrapper_(DLTensor* arg0,
DLTensor* arg1,
DLTensor* arg2,
DLTensor* out0) {
// Cast all DLTensor to primitive type buffers and invoke the above
// execution function.
dnnl_0_(static_cast<float*>(arg0->data),
static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data),
static_cast<float*>(out0->data));
return 0;
}
// The TVM macro to generate TVM runtime compatible function "dnnl_0"
// from our generated "dnnl_0_wrapper_".
TVM_DLL_EXPORT_TYPED_FUNC(dnnl_0, dnnl_0_wrapper_);
注意，预先实现的基于op的DNNL函数位于src / runtime / contrib / dnnl / dnnl.cc中。
由于本文中的其余实现src/relay/backend/contrib/dnnl/codegen.cc都过于DNNL而无法在本文中进行详细介绍。主要思想是实现一个Relay图访问者（L138），访问给定的Relay函数并生成上面的C代码。只要代码生成器能够生成与TVM Runtime兼容的C代码，就可以完全自定义代码生成器以符合要求。
C源代码编译
可能已经注意到，输出的DNNLCompiler是一个带有生成的C代码的文本格式的模块，该模块尚未被编译gcc为可执行二进制文件。实际上，生成的C代码将在用户调用时进行编译export_libray(mod)，如以下代码片段所示：
def update_lib(lib):
# Include the path of src/runtime/contrib/dnnl/dnnl.cc
test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
source_dir = os.path.join(test_dir, "..", "..", "..")
contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")
# Setup the gcc flag to compile DNNL code.
kwargs = {}
kwargs["options"] = ["-O2", "-std=c++14", "-I" + contrib_path]
tmp_path = util.tempdir()
lib_name = 'lib.so'
lib_path = tmp_path.relpath(lib_name)
# The generated C code with DNNL APIs is compiled to a binary lib.so.
lib.export_library(lib_path, fcompile=False, **kwargs)
# Load the lib.so back to a runtime module.
lib = runtime.load_module(lib_path)
return lib
with tvm.transform.PassContext(opt_level=3):
json, lib, param = relay.build(mod, target=target, params=params)
lib = update_lib(lib)
rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx)
将DNNL引入TVM：使用DNNL Codegen / Runtime构建TVM
最后，在构建TVM时创建cmake / modules / contrib / DNNL.cmake，包含DNNL代码源。出于演示目的，DNNL代码生成器在同一cmake文件中具有两个实现。只能根据需要专注于其中之一。
在准备好cmake文件之后，现在用户可以set(USE_DNNL_CODEGEN ON)在其中指定build/config.cmake启用DNNL代码生成。
