将代码生成器带入TVM
将代码生成器带入TVM
为了使数据科学家不必担心开发新模型时的性能,硬件后端提供程序(例如Intel,NVIDIA,ARM等)可以提供诸如cuBLAS或cuDNN之类的内核库以及许多常用的深度学习内核,或者提供诸如此类的框架。例如带有图形引擎的DNNL或TensorRT,使用户以某种方式描述其模型以实现高性能。此外,新兴的深度学习加速器还具有自己的编译器,内核库或运行时runtime框架。
当用户尝试在新的内核库或设备上工作时,必须学习新的编程接口。结果,对统一编程接口的需求变得越来越重要,使所有用户和硬件后端提供程序都站在同一页面上。
为了与广泛使用的深度学习框架共享编程接口,许多硬件设备提供商,尝试将其设备后端集成到TensorFlow。由于TensorFlow没有为新的后端提供正式的后端接口,必须破解TensorFlow进行注册,这需要对许多源文件进行更改,使将来的维护变得困难。
演示了作为硬件后端提供的程序,如何轻松利用自带代码生成(BYOC)框架,将硬件设备的内核库/编译器/框架集成到TVM。利用BYOC框架的最重要优点,设备的所有相关源文件都是独立的,设备的代码源/运行时可插入TVM代码库。这意味着1)使用代码源的TVM代码库将在上游兼容,以及2)TVM用户可以根据需要选择启用代码源/运行时。
的其余部分,首先说明可能需要带有BYOC的TVM的情况,然后概述BYOC编译和运行时流程。然后,分步说明如何使用英特尔DNNL(又名MKL-DNN,OneDNN)作为运行示例,将供应商库或执行引擎与BYOC集成到TVM。
将ASIC加速器带入TVM
首先,让做一个场景,说明为什么要将加速器引入TVM,以及BYOC框架可以包括哪些功能。
刚刚构建了一个具有ARM CPU和出色的加速器的边缘设备平台,该平台为常见的图像分类模型提供了出色的性能。加速器在Conv2D,ReLU,GEMM和其它广泛使用的CNN算子上表现良好。
不幸的是,对象检测模型也越来越受欢迎,并且客户需要在平台上同时运行图像分类和对象检测模型。尽管加速器能够执行对象检测模型中的几乎所有算子,但缺少一个算子(例如,非最大抑制,NMS)。
让TVM执行不受支持的算子
由于TVM具有用于不同后端的多个代码源,开源社区很容易在短时间内在CPU或GPU上实现新的算子。理想情况下,如果将加速器的编译流程与BYOC集成到TVM,则TVM将执行中继图分区,以将部分图卸载到加速器,同时将其它图保持在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中的中继图,BYOC框架执行以下步骤:
图1:原始中继图。
1.图注解
制作用户提供的中继图,第一步是在图中注释可能卸载到加速器的节点。将需要遵循“将DNNL引入TVM:注释规则”以实现受支持的算子的白名单,或定制组合算子的图形模式列表。示例注释结果如图2所示。
图2:带注解的图。
2.图变换
第二步是基于注释对图形进行转换和优化。具体来说,BYOC执行以下转换。
2.1:合并编译器区域:如图2所示,图中现在有许多“区域”可以卸载到加速器中,实际上其中一些区域可以合并,以减少数据传输和内核启动开销。因此,步骤2.1使用贪婪算法来合并尽可能多的那些区域,同时保证功能正确性。结果如图3所示。
图3:合并编译器区域后。
2.2:分区图:对于上一步中的每个区域,创建一个带有属性的Relay函数,Compiler以指示该Relay函数应该完全卸载到加速器上,如图4所示。
图4:图分区之后。
3.代码生成
现在知道应该卸载中继图的哪一部分了。在此步骤中,将每个中继功能依次发送Compiler=your_accelerator到代码源。代码生成器应将Relay函数编译为,与自己的编译流程相匹配的形式。可以是C源代码或任何文本格式。
最后,所有已编译的函数将与其它未卸载的Relay函数一起.so由TVM export_libraryPython API序列化为单个文件。换句话说,.so运行此流程后,用户将仅获得一个文件。
4.运行时runtime
可能还需要实现运行时,初始化图形引擎(如果适用)并执行已编译的函数。在推理期间,当TVM运行时遇到图4中的相应函数调用时,TVM运行时(即图形运行时或VM)将利用运行时来调用已卸载的函数。运行时负责使用给定的输入张量,启动编译后的函数。将数组结果填充到输出张量数组中。
以DNNL为例,演示如何使用BYOC框架实现上述工作流程。所有引用的代码和行号均基于TVM存储库的master分支commit 8a0249c。
将DNNL带到TVM:注释规则
BYOC框架为提供了两种描述受支持的算子和模式的方法。可以同时使用。以DNNL为例来说明如何使用。将代码源的注释规则放在下python/tvm/relay/op/contrib/your_codegen_name.py。
单一运营商规则
可使用BYOC API直观地指定加速器支持哪些中继算子。例如,使用以下代码段,构建一条规则,该规则表明DNNL代码源支持Conv2D:
@tvm.ir.register_op_attr("nn.conv2d", "target.dnnl")
def _dnnl_conv2d_wrapper(attrs, args):
return True
target.dnnl将向中继nn.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传递来执行
%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:中继图转换
利用上一步中的注释规则,现在可以应用BYOC中继传递列表,将中继图从图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
每个中继传递都可以映射到在BYOC工作原理中引入的步骤。
将DNNL引入TVM:JSON代码生成/运行时
让实现将中继图序列化为JSON表示的DNNL代码源,然后实现DNNL JSON运行时以反序列化并执行该图。如果尝试实现一个代码生成器来生成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);
每个运行时模块仅负责一个中继功能,这意味着可能在单个.so文件中包含多个DNNL运行时模块。
DNNL JSON序列化
接下来,实现DNNL JSON序列化器(L429)。从BYOC JSON代码生成器(src / relay / backend / contrib / codegen_json / codegen_json.h)派生。DNNL JSON序列化程序中的特殊过程,将组合函数调用序列化为DNNL JSON运行时可以解释的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]
}
}
问题在于,在运行时仍然需要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运行时可以解释即可。
DNNL JSON运行时
然后,实现DNNL JSON运行时以解释和执行序列化的JSON图。放在下面src/runtime/contrib/dnnl/dnnl_json_runtime.cc。
同样,首先注册两个API来创建运行时,以便可以在任何地方使用。在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运行时实现。基本的类结构为:
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 JSON运行时中的其余实现都是DNNL特有的,不做详细介绍。尽管DNNL JSON运行时是一个很好的开始,但JSON运行时可以完全自定义以满足要求。
将DNNL带到TVM:C源代码生成
实现DNNL代码生成器,该代码生成器生成C源代码,该源代码调用DNNL API来执行中继图。如果尝试实现一个代码生成器,生成其它图形表示形式(如JSON格式)。
为了能够在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);
每个运行时模块仅负责一个中继功能,这意味着可能在单个.so文件中包含多个DNNL运行时模块。
然后,在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源运行时模块兼容的功能接口的说明。
// 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,无法进行详细介绍。主要思想是实现一个中继图访问者(L138)以访问给定的Relay函数,生成上面的C代码。只要代码生成器能够生成与TVM运行时兼容的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代码生成。