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

Figure 1: The Original Relay Graph.

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

Figure 2: The Graph with Annotations.

1. Graph Transformation
   第二步是基于注釋，對圖形進行變換和優化。具體來說，BYOC執行以下轉換。
   2.1：合并編譯器區域：
   如圖2所示，圖中現在有許多“區域”，可以 load 到加速器上，可以合并其中一些區域，減少數據傳輸和內核啟動開銷。步驟2.1使用貪婪算法，合并盡可能多的這些區域，保證功能的正確性。結果如圖3所示。
   

Figure 3: After Merging Compiler Regions.
2.2: Partition Graph:
對于上一步中的每個區域，創建一個帶有屬性編譯器的Relay函數，指示該Relay函數，應該完全 load 到加速器上，如圖4所示。

Figure 4: After Graph Partitioning.
3. Code Generation
現在我們知道應該 load Relay圖的哪個部分。在這一步中，按順序將每個帶有Compiler=your_accelerator加速器的Relay函數，發送到codegen。
codegen應該將Relay函數，編譯成與編譯流相匹配的形式。可以是C源代碼，或任何文本格式。
最后，所有編譯的函數，將與其它未 load 的Relay函數一起，通過TVM export_library Python API，序列化到一個single .so文件中。換句話說，用戶在運行此flow后，將只獲得一個one .so文件。
4. Runtime
可能還需要實現一個runtime，初始化圖形引擎（如果適用），執行編譯后的函數。在推理過程中，當TVM runtime遇到圖4中相應的函數調用時，TVM runtime（即圖runtime或VM）將利用runtime，調用 load 的函數。runtime負責使用給定的輸入張量數組，啟動編譯函數，將結果填充到輸出張量數組中。
以DNNL為例，演示如何使用BYOC框架，實現上述工作流。本文引用的所有代碼和行號，都基于TVM存儲庫的主分支提交8a0249c。
Bring DNNL to TVM: Annotation Rules
BYOC框架提供了兩種方法，描述支持的算子和模式，以DNNL為例來說明如何使用。這里提供了完整的實現。將codegen的注釋規則放在
python/tvm/relay/op/contrib/your_codegen_name.py.
Rules for single operators
可以使用BYOC API，直觀地指定加速器，支持哪些Relay算子。例如，使用下面的代碼片段，構建一個規則，說明DNNL codegen支持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 codegen中是否支持。
另一方面，每個算子編寫上面的代碼片段，可能很乏味。對于DNNL實現，實現了一個helper函數，即_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），優化為單個指令或API。可以指定從圖形模式，到指令/API的映射。對于DNNL來說，Conv2D API已經包含了bias addition，允許附加下一個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示例中，實現了兩個具有不同名稱的模式，可以在codegen中輕松地識別。這些模式是用Relay模式語言實現的。可以學習如何編寫模式。
通過模式表，可以使用一個Relay pass執行
%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）。
在復合函數中，還有一個名為“PartitionedFromPattern”的屬性。如果模式包含通配符算子，這可能會很有幫助。例如，可能有一個模式表(“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
使用上一步中的注釋規則，可以應用BYOC Relay pass列表，將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
As can be seen, each Relay pass can be mapped to a step we have introduced in How BYOC Works.

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。