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代码生成。