TVM：使用 Schedule 模板和 AutoTVM 來優化算子

在本文中，我們將介紹如何使用 TVM 張量表達式（Tensor Expression，TE）語言編寫 Schedule 模板，AutoTVM 可以搜索通過這些模板找到最佳 Schedule。這個過程稱為自動調整（Auto Tuning），它有助于自動優化張量計算的過程。

本教程需基于之前介紹的如何使用 TE 來寫一個矩陣乘法的教程。

Auto Tuning 有兩步：

• 第一步是定義一個搜索空間
• 第二步是運行相應的搜索算法來探索這個空間

本教程將展示如何在 TVM 中完成這兩步，整個流程將以矩陣乘法為例。

安裝依賴

要使用 TVM 中的 AutoTVM 包，需要安裝這些額外的依賴：

pip install --user psutil xgboost cloudpickle

為了使 TVM 在 tuning 中運行得更快，建議使用 cython 作為 TVM 的 FFI。在 TVM 的根目錄中，執行：

pip3 install --user cython
sudo make cython3

現在我們開始寫 Python 代碼，先引入包：

import logging
import sysimport numpy as np
import tvm
from tvm import te
import tvm.testingfrom tvm import autotvm

TE 實現基本的矩陣乘法

回想一下使用TE實現矩陣乘法的基本方法。我們把它寫在這里，稍作修改。我們將把乘法封裝在Python 函數定義中。簡單起見，我們將把注意力集中在分割優化（split optimization）上，使用一個固定值來定義重排（reordering）的塊的大小。

def matmul_basic(N, L, M, dtype):A = te.placeholder((N, L), name="A", dtype=dtype)B = te.placeholder((L, M), name="B", dtype=dtype)k = te.reduce_axis((0, L), name="k")C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")s = te.create_schedule(C.op)# scheduley, x = s[C].op.axisk = s[C].op.reduce_axis[0]yo, yi = s[C].split(y, 8)xo, xi = s[C].split(x, 8)s[C].reorder(yo, xo, k, yi, xi)return s, [A, B, C]

使用 AutoTVM 優化矩陣乘法

在前面的明細表代碼中，我們使用常量值 8 作為 tiling 因子。但是，它可能不是最好的，因為最佳的tiling 因子取決于實際的硬件環境和輸入的形狀。

如果希望 Schedule 代碼能夠在更大范圍的輸入形狀和目標硬件上移植，最好定義一組候選值，并根據目標硬件上的測量結果選擇最佳值。

在 autotvm 中，我們可以定義一個可調參數，或為此類參數定義一個“knob”。

一個基本的矩陣乘法模板

我們的一個示例，介紹如何為 spliting schedule 操作的塊大小創建一個可調參數集。

# Matmul V1: List candidate values
@autotvm.template("tutorial/matmul_v1")  # 1. use a decorator
def matmul_v1(N, L, M, dtype):A = te.placeholder((N, L), name="A", dtype=dtype)B = te.placeholder((L, M), name="B", dtype=dtype)k = te.reduce_axis((0, L), name="k")C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")s = te.create_schedule(C.op)# scheduley, x = s[C].op.axisk = s[C].op.reduce_axis[0]# 2. get the config objectcfg = autotvm.get_config()# 3. define search spacecfg.define_knob("tile_y", [1, 2, 4, 8, 16])cfg.define_knob("tile_x", [1, 2, 4, 8, 16])# 4. schedule according to configyo, yi = s[C].split(y, cfg["tile_y"].val)xo, xi = s[C].split(x, cfg["tile_x"].val)s[C].reorder(yo, xo, k, yi, xi)return s, [A, B, C]

這里對于之前的 schedule 代碼，我們有四處調整，從而可以得到一個可調節的 “模板”：

1. 使用裝飾器 @autotvm.template() 來將此函數標記為一個模板

2. 獲取配置項： cfg 可以視為本函數的一個參數，但是我們是通過一種不同的方式得到它的。有了這個參數，這個函數就不再是一個確定的 schedule。而是我們可以為這個函數傳入不同的配置來得到不同的 schedules。一個像這樣含有配置項的函數稱為 “模板”。

   為了使得模板函數更加 compact，我們通過完成以下兩件事情來在一個函數中定義參數搜索空間：

   • 通過一組值定義搜索空間。這是通過將 cfg 設置為一個 ConfigSpace 對象完成的。它會收集本函數中所有的可以調節的 knobs ，并從中建立一個搜索空間。
   • 根據此空間中的實體進行 schedule。這是通過將 cfg 設置為一個 ConfigEntity 對象完成的。當它是一個 ConfigEntity 時，它會忽略所有空間定義的API（即 cfg.define_XXX(...)）。相反，它會所有可調節的 knobs 保存確定的值，然后我們根據這些值進行 schedule。

   在 auto-tuning 時，我們將首先使用 ConfigSpace 對象調用此模板以構建搜索空間。然后，我們在構建空間中使用不同的 ConfigEntity 調用此模板，以獲得不同的 schedule。最后，我們將度量由不同計劃生成的代碼，并選擇最優的。

3. 定義兩個可調節的 knobs。第一個是具有 5 個可能值的 tile_y。第二個是 tile_x，具有相同的可能值列表。這兩個 knob 是獨立的，因此它們會有大小為 25=5x5 的搜索空間。

4. 配置 knobs 被傳遞到 split schedule 操作，允許我們根據之前在 cfg 中定義的 5x5 確定值進行 schedule。

具有高級參數API的矩陣乘法模板

在前面的模板中，我們手動列出了 knobs 的所有可能值。這是定義空間的最低級別API，并提供要搜索的參數空間的顯式枚舉。但是，我們還提供了另一組API，可以使搜索空間的定義更簡單、更智能。在可能的情況下，我們建議您使用此高級API。

在下面的示例中，我們使用 ConfigSpace.define_split 來定義拆分 knob。它將列舉所有分割軸和構建空間的可能方法。

我們還有 ConfigSpace.define_reorder用于 reorder knob，ConfigSpace.define_annotation 用于展開、矢量化、線程綁定等注釋。當高級API不能滿足您的需求時，您可以隨時使用低級API。

@autotvm.template("tutorial/matmul")
def matmul(N, L, M, dtype):A = te.placeholder((N, L), name="A", dtype=dtype)B = te.placeholder((L, M), name="B", dtype=dtype)k = te.reduce_axis((0, L), name="k")C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")s = te.create_schedule(C.op)# scheduley, x = s[C].op.axisk = s[C].op.reduce_axis[0]##### define space begin #####cfg = autotvm.get_config()cfg.define_split("tile_y", y, num_outputs=2)cfg.define_split("tile_x", x, num_outputs=2)##### define space end ###### schedule according to configyo, yi = cfg["tile_y"].apply(s, C, y)xo, xi = cfg["tile_x"].apply(s, C, x)s[C].reorder(yo, xo, k, yi, xi)return s, [A, B, C]

注意：More Explanation on cfg.define_split

更多關于 cfg.define_split 的解釋

在此模板中，cfg.define_split(“tile_y”，y，num_outputs=2) 將枚舉所有可能的組合，這些組合可以將y軸拆分為兩個具有y長度因子的軸。例如，如果y的長度為32，我們希望使用32的因子將其拆分為兩個軸，那么（外軸長度、內軸長度）對有6個可能的值，即（32,1）、（16,2）、（8,4）、（4,8）、（2,16）或（1,32）。這些都是 tile_y 的6個可能值。

在 schedule 期間，cfg[“tile_y”] 是一個 SplitEntity 對象。我們將外軸和內軸的長度存儲在 cfg['tile_y'].size（包含兩個元素的元組）中。在這個模板中，我們使用yo, yi=cfg['tile_y']] 來應用它。實際上，這相當于yo, yi=s[C].split（y，cfg[“tile\u y”].size[1]）或 yo, yi=s[C].split（y，npart=cfg[“tile\u y”].size[0]）

使用 cfg.apply API 的優點是，它使多級拆分（即當num_outputs>=3時）更加容易。

第二步：使用AutoTVM優化矩陣乘法

在第一步中，我們編寫了一個矩陣乘法模板，該模板允許我們對分割 schedule 中使用的塊大小進行參數化。現在我們可以在這個參數空間上進行搜索。下一步是選擇一個調諧器（tuner）來指導這個空間的探索。

TVM 中的 Auto-tuners

調諧器的工作可以通過以下偽代碼來描述：

ct = 0
while ct < max_number_of_trials:propose a batch of configsmeasure this batch of configs on real hardware and get resultsct += batch_size

提出下一批配置時，調諧器可以采取不同的策略。TVM提供的一些調諧器策略包括：

• tvm.autotvm.tuner.RandomTuner：按隨機順序枚舉空間
• tvm.autotvm.tuner.GridSearchTuner：按網格搜索順序枚舉空間
• tvm.autotvm.tuner.GATuner：利用遺傳算法進行空間搜索
• tvm.autotvm.tuner.XGBTuner：使用基于模型的方法。訓練 XGBoost 模型預測降低 IR 的速度，并根據預測選擇下一批。

我們可以根據空間大小、時間預算和其他因素選擇調諧器。例如，如果您的空間非常小（小于1000），gridsearch 調諧器或隨機調諧器就足夠了。如果您的空間級別為 10^9（這是CUDA GPU上conv2d 算子的空間大小），XGBoostTuner 可以更高效地探索并找到更好的配置。

開始 tuning

這里我們繼續我們的矩陣乘法示例。首先，我們創建一個 tuning 任務。我們還可以檢查初始化的搜索空間。在這種情況下，對于 512x512 平方矩陣乘法，空間大小為 10x10=100。請注意，任務和搜索空間與選擇的調諧器無關。

N, L, M = 512, 512, 512
task = autotvm.task.create("tutorial/matmul", args=(N, L, M, "float32"), target="llvm")
print(task.config_space)

此處輸出：

N, L, M = 512, 512, 512
task = autotvm.task.create("tutorial/matmul", args=(N, L, M, "float32"), target="llvm")
print(task.config_space)

然后我們需要定義如何測量生成的代碼并選擇調諧器。因為我們的空間很小，隨機調諧器也可以。

在本教程中，我們只做了 10 次試驗來演示。實際上，你可以根據你的時間預算做更多的試驗。我們將把調優結果記錄到日志文件中。此文件可用于選擇調諧器稍后發現的最佳配置。

# logging config (for printing tuning log to the screen)
logging.getLogger("autotvm").setLevel(logging.DEBUG)
logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

測量配置有兩個步驟：構建和運行。默認情況下，我們使用所有CPU核來編譯程序。然后我們依次測量它們。為了減少方差，我們進行了5次測量并取平均值。

measure_option = autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(number=5))# Begin tuning with RandomTuner, log records to file `matmul.log`
# You can use alternatives like XGBTuner.
tuner = autotvm.tuner.RandomTuner(task)
tuner.tune(n_trial=10,measure_option=measure_option,callbacks=[autotvm.callback.log_to_file("matmul.log")],
)

調優 tuning 完成后，我們可以從日志文件中選擇性能最好的配置，并使用相應的參數編譯 schedule。我們還快速驗證了 schedule 是否正確。我們可以直接在 autotvm.apply_history_best 上下文下調用函數 matmul。當我們調用此函數時，它將使用其參數查詢分派上下文，并使用相同的參數獲得最佳配置。

# apply history best from log file
with autotvm.apply_history_best("matmul.log"):with tvm.target.Target("llvm"):s, arg_bufs = matmul(N, L, M, "float32")func = tvm.build(s, arg_bufs)# check correctness
a_np = np.random.uniform(size=(N, L)).astype(np.float32)
b_np = np.random.uniform(size=(L, M)).astype(np.float32)
c_np = a_np.dot(b_np)c_tvm = tvm.nd.empty(c_np.shape)
func(tvm.nd.array(a_np), tvm.nd.array(b_np), c_tvm)tvm.testing.assert_allclose(c_np, c_tvm.numpy(), rtol=1e-4)

最后的注意事項與總結

在本教程中，我們展示了如何構建算子模板，以使得 TVM 能夠對參數空間進行搜索并選擇最優的的 schedule 配置。為了更深入地了解其工作原理，我們建議在此示例上進行擴展，可以根據使用張量表達式（TE）入門教程中演示的 schedule 操作向計劃添加新的搜索參數。在接下來的部分中，我們將演示 AutoScheduler，這是一種TVM優化常用算子的方法，用戶無需提供用戶定義的模板。

Ref：
https://tvm.apache.org/docs/tutorial/autotvm_matmul_x86.html