如何在TVM上集成Codegen（上）
許多常用的深度學(xué)習(xí)內(nèi)核，或者提供DNNL或TensorRT等框架和圖形引擎，讓用戶以某種方式描述模型，從而獲得高性能。此外，新興的深度學(xué)習(xí)加速器也有自己的編譯器、內(nèi)核庫(kù)或runtime框架。
當(dāng)用戶試圖在一個(gè)新的內(nèi)核庫(kù)或設(shè)備上工作時(shí)，必須學(xué)習(xí)一個(gè)新的編程接口。對(duì)于統(tǒng)一編程接口的需求變得越來(lái)越重要，讓所有用戶和硬件后端提供商站在同一個(gè)頁(yè)面上。
廣泛使用的深度學(xué)習(xí)框架共享編程接口，許多硬件設(shè)備提供商，嘗試將設(shè)備后端集成到TensorFlow。由于TensorFlow沒(méi)有為新的后端提供正式的后端接口，必須對(duì)TensorFlow進(jìn)行注冊(cè)，這涉及到許多源文件的更改，維護(hù)變得困難。
本節(jié)將展示作為一個(gè)硬件后端提供商，如何輕松地利用自帶的Codegen（BYOC）框架，將硬件設(shè)備的內(nèi)核庫(kù)/編譯器/框架，集成到TVM。利用BYOC框架最重要的優(yōu)勢(shì)，設(shè)備的所有相關(guān)源文件都是自包含的，設(shè)備的codegen/runtime，可以嵌入到TVM代碼庫(kù)。
1）帶有codegen的TVM代碼庫(kù)，將與上游兼容
2）TVM用戶可以根據(jù)需要，選擇啟用codegen/runtime。
首先說(shuō)明一個(gè)場(chǎng)景，可能需要使用BYOC實(shí)現(xiàn)TVM，然后概述BYOC編譯和runtime流。最后，以Intel DNNL（又稱MKL-DNN，OneDNN）為例，逐步說(shuō)明如何將供應(yīng)商庫(kù)或執(zhí)行引擎，集成到TVM與BYOC。
Bring an ASIC Accelerator to TVM
先做一個(gè)場(chǎng)景，說(shuō)明為什么要將加速器引入TVM，可以從BYOC框架中，獲得哪些特性。如果不確定案例是否適合BYOC，歡迎在tvm.ai里討論。
假如制作了一個(gè)邊緣設(shè)備平臺(tái)，有一個(gè)ARM CPU和一個(gè)很棒的加速器，在常見(jiàn)的圖像分類模型中，取得了驚人的性能。加速器在Conv2D、ReLU、GEMM和其它廣泛使用的CNN算子上表現(xiàn)良好。
但是，目標(biāo)檢測(cè)模型也越來(lái)越流行，客戶需要在平臺(tái)上，同時(shí)運(yùn)行圖像分類和目標(biāo)檢測(cè)模型。雖然加速器能夠執(zhí)行目標(biāo)檢測(cè)模型中的幾乎所有算子，但缺少一個(gè)算子（例如，非最大抑制，NMS）。
Let TVM execute unsupported operators
由于TVM為不同的后端，提供了多個(gè)代碼源，開(kāi)源社區(qū)很容易在短時(shí)間內(nèi)，在CPU或GPU上實(shí)現(xiàn)新的操作程序。如果將加速器的編譯流與BYOC集成到TVM，TVM將執(zhí)行Relay圖分區(qū)，將圖的一部分 load 到加速器上，將其它部分保留在TVM上。因此，可以宣稱平臺(tái)能夠運(yùn)行所有模型，而不必?fù)?dān)心新的算子。
Customize graph-level optimization
ASIC加速器必須有自己的編譯流。通常，可能是以下情況之一：
生成一個(gè)圖形表示，將其輸入圖形引擎：
可能有自己的圖形引擎，能夠在加速器上執(zhí)行圖形（或神經(jīng)網(wǎng)絡(luò)模型）。例如，Intel DNNL和NVIDIA TensorRT，都使用引擎運(yùn)行整個(gè)圖形或模型，因此能夠
1）減少算子之間的內(nèi)存事務(wù)；
2）使用算子融合優(yōu)化圖形執(zhí)行。
為了實(shí)現(xiàn)上述兩個(gè)優(yōu)化，可能需要在編譯期間處理該圖。例如，Conv2D和bias addition，在TVM中是兩個(gè)獨(dú)立的算子，但可能是加速器上的一個(gè)算子（具有bias addition功能的Conv2D）。可能希望通過(guò)將conv2d-add graph模式，替換為帶有“bias”節(jié)點(diǎn)的“uconv2d”，優(yōu)化圖形。
如果編譯流程屬于這種情況，建議閱讀本文的其余部分，但跳過(guò)將DNNL帶到TVM:C源代碼生成。
生成匯編代碼，編譯為可執(zhí)行的二進(jìn)制文件：
如果平臺(tái)不像前面的例子，有一個(gè)端到端的執(zhí)行框架，可能有一個(gè)編譯器，用ISA的匯編代碼編譯程序。為了向編譯器提供匯編代碼，需要一個(gè)codegen，從Relay圖生成和優(yōu)化匯編代碼。
如果編譯流程屬于這種情況，建議閱讀本文的所有其余部分，但跳過(guò)將DNNL到TVM:JSON Codegen/Runtime。
How BYOC Works
簡(jiǎn)單地解釋一下BYOC框架，如何工作的。有關(guān)底層框架組件及實(shí)現(xiàn)的詳細(xì)說(shuō)明，參閱開(kāi)發(fā)人員文檔。給定圖1中的Relay圖，BYOC框架執(zhí)行以下步驟：

Figure 1: The Original Relay Graph.

1. Graph Annotation
   以用戶提供的Relay圖為例，第一步是在圖中注釋，可能 load 到加速器的節(jié)點(diǎn)。需要遵循Bring DNNL to TVM:
   來(lái)實(shí)現(xiàn)受支持算子的白名單，或者自定義復(fù)合算子的圖形模式列表。圖2顯示了一個(gè)示例注釋結(jié)果。
   

Figure 2: The Graph with Annotations.

1. Graph Transformation
   第二步是基于注釋，對(duì)圖形進(jìn)行變換和優(yōu)化。具體來(lái)說(shuō)，BYOC執(zhí)行以下轉(zhuǎn)換。
   2.1：合并編譯器區(qū)域：
   如圖2所示，圖中現(xiàn)在有許多“區(qū)域”，可以 load 到加速器上，可以合并其中一些區(qū)域，減少數(shù)據(jù)傳輸和內(nèi)核啟動(dòng)開(kāi)銷。步驟2.1使用貪婪算法，合并盡可能多的這些區(qū)域，保證功能的正確性。結(jié)果如圖3所示。
   

Figure 3: After Merging Compiler Regions.
2.2: Partition Graph:
對(duì)于上一步中的每個(gè)區(qū)域，創(chuàng)建一個(gè)帶有屬性編譯器的Relay函數(shù)，指示該Relay函數(shù)，應(yīng)該完全 load 到加速器上，如圖4所示。

Figure 4: After Graph Partitioning.
3. Code Generation
現(xiàn)在我們知道應(yīng)該 load Relay圖的哪個(gè)部分。在這一步中，按順序?qū)⒚總€(gè)帶有Compiler=your_accelerator加速器的Relay函數(shù)，發(fā)送到codegen。
codegen應(yīng)該將Relay函數(shù)，編譯成與編譯流相匹配的形式。可以是C源代碼，或任何文本格式。
最后，所有編譯的函數(shù)，將與其它未 load 的Relay函數(shù)一起，通過(guò)TVM export_library Python API，序列化到一個(gè)single .so文件中。換句話說(shuō)，用戶在運(yùn)行此flow后，將只獲得一個(gè)one .so文件。
4. Runtime
可能還需要實(shí)現(xiàn)一個(gè)runtime，初始化圖形引擎（如果適用），執(zhí)行編譯后的函數(shù)。在推理過(guò)程中，當(dāng)TVM runtime遇到圖4中相應(yīng)的函數(shù)調(diào)用時(shí)，TVM runtime（即圖runtime或VM）將利用runtime，調(diào)用 load 的函數(shù)。runtime負(fù)責(zé)使用給定的輸入張量數(shù)組，啟動(dòng)編譯函數(shù)，將結(jié)果填充到輸出張量數(shù)組中。
以DNNL為例，演示如何使用BYOC框架，實(shí)現(xiàn)上述工作流。本文引用的所有代碼和行號(hào)，都基于TVM存儲(chǔ)庫(kù)的主分支提交8a0249c。
Bring DNNL to TVM: Annotation Rules
BYOC框架提供了兩種方法，描述支持的算子和模式，以DNNL為例來(lái)說(shuō)明如何使用。這里提供了完整的實(shí)現(xiàn)。將codegen的注釋規(guī)則放在
python/tvm/relay/op/contrib/your_codegen_name.py.
Rules for single operators
可以使用BYOC API，直觀地指定加速器，支持哪些Relay算子。例如，使用下面的代碼片段，構(gòu)建一個(gè)規(guī)則，說(shuō)明DNNL codegen支持Conv2D：
@tvm.ir.register_op_attr(“nn.conv2d”, “target.dnnl”)
def _dnnl_conv2d_wrapper(attrs, args):
return True
這將注冊(cè)一個(gè)新屬性target.dnnl接力nn.conv2d算子。通過(guò)這種方式，BYOC注釋可以調(diào)用target.dnnl（），檢查DNNL codegen中是否支持。
另一方面，每個(gè)算子編寫(xiě)上面的代碼片段，可能很乏味。對(duì)于DNNL實(shí)現(xiàn)，實(shí)現(xiàn)了一個(gè)helper函數(shù)，即_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 codegen支持的算子列表。
Rules for graph patterns
加速器或編譯器，可能已將某些模式（例如Conv2D+add+ReLU），優(yōu)化為單個(gè)指令或API。可以指定從圖形模式，到指令/API的映射。對(duì)于DNNL來(lái)說(shuō)，Conv2D API已經(jīng)包含了bias addition，允許附加下一個(gè)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 …
在本例中，除了單個(gè)conv2d，希望將圖形模式conv2d+relu映射到DNNLConv2d（false，true），將conv2d+add+relu映射到DNNLConv2d（true，true）。用下面的代碼片段實(shí)現(xiàn)：
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(‘a(chǎn)dd’)(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示例中，實(shí)現(xiàn)了兩個(gè)具有不同名稱的模式，可以在codegen中輕松地識(shí)別。這些模式是用Relay模式語(yǔ)言實(shí)現(xiàn)的。可以學(xué)習(xí)如何編寫(xiě)模式。
通過(guò)模式表，可以使用一個(gè)Relay pass執(zhí)行
%1 = nn.conv2d(%data, %weight, …)
%2 = add(%1, %bias)
%3 = nn.relu(%2)
to
%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)
hus，DNNL codegen可以獲得模式名conv2d_bias_relu，將%1映射到DNNLConv2d（true，true）。
在復(fù)合函數(shù)中，還有一個(gè)名為“PartitionedFromPattern”的屬性。如果模式包含通配符算子，這可能會(huì)很有幫助。例如，可能有一個(gè)模式表(“conv2d_with_something”, conv2d -> *):
def make_pattern(with_bias=True):
data = wildcard()
weight = wildcard()
conv = is_op(‘nn.conv2d’)(data, weight)
return wildcard()(conv)
In this case, you will get a composite function with Composite=conv2d_with_something, but you have no idea about what graph it actually matched. That’s where PartitionedFromPattern comes into play. You can know that if the matched graph is conv2d -> add or conv2d -> relu by looking at PartitionedFromPattern to see if it is nn.conv2d_add_ or nn.conv2d_nn.relu_.
Bring DNNL to TVM: Relay Graph Transformation
使用上一步中的注釋規(guī)則，可以應(yīng)用BYOC Relay pass列表，將Relay圖從圖1轉(zhuǎn)換為圖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
As can be seen, each Relay pass can be mapped to a step we have introduced in How BYOC Works.

總結(jié)

以上是生活随笔為你收集整理的如何在TVM上集成Codegen（上）的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。