TVM：使用 Schedule 模板和 AutoTVM 來優(yōu)化算子

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

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

Auto Tuning 有兩步：

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

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

安裝依賴

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

pip install --user psutil xgboost cloudpickle

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

pip3 install --user cython sudo make cython3

現(xiàn)在我們開始寫 Python 代碼，先引入包：

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

TE 實現(xiàn)基本的矩陣乘法

回想一下使用TE實現(xiàn)矩陣乘法的基本方法。我們把它寫在這里，稍作修改。我們將把乘法封裝在Python 函數(shù)定義中。簡單起見，我們將把注意力集中在分割優(yōu)化（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 優(yōu)化矩陣乘法

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

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

在 autotvm 中，我們可以定義一個可調(diào)參數(shù)，或為此類參數(shù)定義一個“knob”。

一個基本的矩陣乘法模板

我們的一個示例，介紹如何為 spliting schedule 操作的塊大小創(chuàng)建一個可調(diào)參數(shù)集。

# 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 代碼，我們有四處調(diào)整，從而可以得到一個可調(diào)節(jié)的 “模板”：

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

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

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

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

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

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

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

具有高級參數(shù)API的矩陣乘法模板

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

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

我們還有 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的因子將其拆分為兩個軸，那么（外軸長度、內(nèi)軸長度）對有6個可能的值，即（32,1）、（16,2）、（8,4）、（4,8）、（2,16）或（1,32）。這些都是 tile_y 的6個可能值。

在 schedule 期間，cfg[“tile_y”] 是一個 SplitEntity 對象。我們將外軸和內(nèi)軸的長度存儲在 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 的優(yōu)點是，它使多級拆分（即當num_outputs>=3時）更加容易。

第二步：使用AutoTVM優(yōu)化矩陣乘法

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

TVM 中的 Auto-tuners

調(diào)諧器的工作可以通過以下偽代碼來描述：

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

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

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

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

開始 tuning

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

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)

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

在本教程中，我們只做了 10 次試驗來演示。實際上，你可以根據(jù)你的時間預算做更多的試驗。我們將把調(diào)優(yōu)結(jié)果記錄到日志文件中。此文件可用于選擇調(diào)諧器稍后發(fā)現(xiàn)的最佳配置。

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

測量配置有兩個步驟：構(gòu)建和運行。默認情況下，我們使用所有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")], )

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

# 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)

最后的注意事項與總結(jié)

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

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

總結(jié)

以上是生活随笔為你收集整理的TVM：使用 Schedule 模板和 AutoTVM 来优化算子的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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