編寫可調模板並使用自動調諧器

wujianming_110117發表於2020-12-24

編寫可調模板並使用自動調諧器
這是TVM中自動除錯模組的入門說明。
自動除錯有兩個步驟。第一步是定義搜尋空間。第二步是執行搜尋演算法來探索這個空間。本文可以學習如何在TVM中執行這兩個步驟。整個工作流程由矩陣乘法示例說明。
注意,本文無法在Windows或最新版本的macOS上執行。要使其執行,需要將本文的內容包裝在一個if name == “main”:塊中。
安裝依賴
要在TVM中使用autotvm軟體包,需要安裝一些額外的依賴項。可以跳過此步驟(安裝xgboost),因為它不需要XGBoost(如果使用python2,請將“ 3”更改為“ 2”):
pip3 install --user psutil xgboost
為了使TVM在除錯中更快地執行,建議使用cython作為TVM的FFI。在TVM的根目錄中,執行(如果使用python2,請將“ 3”更改為“ 2”):
pip3 install --user cython
sudo make cython3
現在返回python程式碼,匯入包。
import logging
import sys

import numpy as np
import tvm
from tvm import te, testing

the module is called autotvm

from tvm import autotvm
步驟1:定義搜尋空間
在本節中,將確定性TVM排程程式碼重寫為可調排程模板。可以將搜尋空間定義的過程視為現有排程程式碼的引數化。
首先,這是如何在TVM中實現分塊矩陣乘法。

Matmul V0: Constant tiling factor

def matmul_v0(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)

# schedule
y, x = s[C].op.axis
k = 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]

引數化排程
在先前的排程程式碼中,使用常數“ 8”作為切片因子。但是,可能不是最好的,因為最好的切片因子取決於實際的硬體環境和輸入形態。
如果希望排程程式碼可在更廣泛的輸入形態和目標硬體中移植,則最好定義一組候選值並根據目標硬體上的測量結果選擇最佳值。
在autotvm中,可以為此類值定義可調引數或“旋鈕”。

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)

# schedule
y, x = s[C].op.axis
k = s[C].op.reduce_axis[0]

# 2. get the config object
cfg = autotvm.get_config()

# 3. define search space
cfg.define_knob("tile_y", [1, 2, 4, 8, 16])
cfg.define_knob("tile_x", [1, 2, 4, 8, 16])

# 4. schedule according to config
yo, 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]

在這裡,對先前的排程程式碼進行了四個修改,並獲得了一個可調的“模板”。可以一一解釋這些修改。
• 使用decorator將此功能標記為簡單模板。
• 獲取配置物件:可以將此cfg作為該函式的引數,但是可以通過其他方式獲得。使用此引數,此功能不再是確定性的排程程式碼。相反,可以將不同的配置傳遞給此功能並獲得不同的排程,因此此功能是一個“模板”。
為了使模板函式更緊湊,在單個函式中做了兩件事。(1)定義搜尋空間,(2)根據該空間中的實體進行排程。為了達到這個目的,將cfg其設為aConfigSpace或ConfigEntityobject。
如果是ConfigSpace,將收集此功能中的所有可tune調旋鈕並建立搜尋空間。如果是ConfigEntity,它將忽略所有空間定義API(即cfg.define_XXXXX(…))。相反,儲存所有可調旋鈕的確定性值,根據這些值進行排程。
在自動除錯期間,將首先使用一個ConfigSpace 物件呼叫此模板以構建搜尋空間。然後,ConfigEntity 在構建空間中將此模板稱為不同的模板,以獲取不同的排程。最後,將測量由不同排程生成的程式碼,並選擇最佳排程。
• 定義兩個可調旋鈕。第一個tile_y具有5個可能的值。第二個tile_x具有相同的可能值列表。這兩個旋鈕是獨立的,因此它們跨越的搜尋空間的大小為5x5 = 25
• 根據cfg中的確定性值進行排程
使用更好的空間定義API
在上一個模板中,手動列出了旋鈕的所有可能值。這是定義空間的最低階別的API。但是,還提供了另一組API,以使空間定義更輕鬆,更智慧。建議使用這套高階API。
在以下示例中,用ConfigSpace.define_split定義拆分旋鈕。它將列舉所有可能的方式來分割軸並構造空間。
還有ConfigSpace.define_reorder用於重新排序的旋鈕和 ConfigSpace.define_annotate用於展開,向量化,執行緒繫結的註釋。當高階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)

# schedule
y, x = s[C].op.axis
k = 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 config
yo, 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]

注意
有關cfg.defile_split更多說明
在此模板中,將列舉所有可以將y軸分解為長度為y的兩個軸的可能組合。例如,如果y的長度為32,並且想使用32的因數將其分成兩個軸,則(外軸的長度,內軸的長度)對有6個可能的值,即(32,1) ,(16、2),(8、4),(4、8),(2、16)或(1、32)。它們只有6個可能tile_y的值。cfg.define_split(“tile_y”, y, num_outputs=2)
在排程期間,cfg[“tile_y”]是一個SplitEntity物件。將外軸和內軸的長度儲存在cfg[‘tile_y’].size (具有兩個元素的元組)中。在此模板中,使用yo, yi = cfg[‘tile_y’].apply(s, C, y)來應用它。實際上,這等於 yo, yi = s[C].split(y, cfg[“tile_y”].size[1])或 yo, yi = s[C].split(y, nparts=cfg['tile_y"].size[0])
使用cfg.apply API的優點是,使多層拆分(當num_outputs> = 3時)更加容易。
步驟2:搜尋空間
在第1步中,通過將舊的排程程式碼擴充套件到模板中來構建搜尋空間。下一步是選擇一個tune調諧器,並在這個空間中進行探索。
TVM中的自動tune調諧器
tune調諧器的工作可以通過以下虛擬碼來描述
ct = 0
while ct < max_number_of_trials:
propose a batch of configs
measure this batch of configs on real hardware and get results
ct += batch_size
當提議下一批配置時,調諧器可以採取不同的策略。在autotvm中為四個調諧器提供了不同的策略。
• RandomTuner:以隨機順序列舉空間
• GridSearchTuner:按網格搜尋順序列舉空間
• GATuner:使用遺傳演算法搜尋空間
• XGBTuner:使用基於模型的方法。訓練XGBoost模型以預測降低的IR的速度,並根據預測選擇下一批。
可以根據空間大小,時間預算和其他因素選擇調諧器。例如,如果空間很小(小於1000),那麼使用gridsearch調諧器或隨機調諧器就足夠了。如果空間為10 ^ 9的水平(這是CUDA GPU上conv2d運算子的空間大小),則XGBoostTuner可以更有效地進行探索並找到更好的配置。
開始除錯
在這裡,繼續矩陣乘法示例。首先,應該建立一個除錯任務。還可以檢查初始化的搜尋空間。在這種情況下,對於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)
出:
ConfigSpace (len=100, space_map=
0 tile_y: Split(policy=factors, product=512, num_outputs=2) len=10
1 tile_x: Split(policy=factors, product=512, num_outputs=2) len=10
)
然後,需要定義如何測量生成的程式碼並選擇一個調諧器。由於的空間很小,所以可以使用隨機調諧器。
在本文中,僅進行10個試驗進行演示。實際上,可以根據自己的時間預算進行更多試驗。會將除錯結果記錄到日誌檔案中。此檔案可用於稍後獲得最佳配置。

logging config (for printing tuning log to the screen)

logging.getLogger(“autotvm”).setLevel(logging.DEBUG)
logging.getLogger(“autotvm”).addHandler(logging.StreamHandler(sys.stdout))

There are two steps for measuring a config: build and run.

By default, we use all CPU cores to compile program. Then measure them sequentially.

We measure 5 times and take average to reduce variance.

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”)],
)
出:
Get devices for measurement successfully!
No: 1 GFLOPS: 0.52/0.52 result: MeasureResult(costs=(0.5179643672,), error_no=0, all_cost=8.699557542800903, timestamp=1607225778.9184623) [(‘tile_y’, [-1, 64]), (‘tile_x’, [-1, 1])],None,6
No: 2 GFLOPS: 2.05/2.05 result: MeasureResult(costs=(0.1307110214,), error_no=0, all_cost=2.452157735824585, timestamp=1607225781.4836178) [(‘tile_y’, [-1, 512]), (‘tile_x’, [-1, 8])],None,39
No: 3 GFLOPS: 2.77/2.77 result: MeasureResult(costs=(0.0968108324,), error_no=0, all_cost=2.015434741973877, timestamp=1607225783.5040994) [(‘tile_y’, [-1, 2]), (‘tile_x’, [-1, 8])],None,31
No: 4 GFLOPS: 7.71/7.71 result: MeasureResult(costs=(0.0348177938,), error_no=0, all_cost=0.9887301921844482, timestamp=1607225784.5313203) [(‘tile_y’, [-1, 1]), (‘tile_x’, [-1, 32])],None,50
No: 5 GFLOPS: 13.46/13.46 result: MeasureResult(costs=(0.0199451586,), error_no=0, all_cost=0.7833263874053955, timestamp=1607225785.3334467) [(‘tile_y’, [-1, 256]), (‘tile_x’, [-1, 64])],None,68
No: 6 GFLOPS: 11.91/13.46 result: MeasureResult(costs=(0.0225446656,), error_no=0, all_cost=0.7622959613800049, timestamp=1607225786.1802726) [(‘tile_y’, [-1, 256]), (‘tile_x’, [-1, 512])],None,98
No: 7 GFLOPS: 0.92/13.46 result: MeasureResult(costs=(0.2913359364,), error_no=0, all_cost=5.074311971664429, timestamp=1607225791.3119547) [(‘tile_y’, [-1, 128]), (‘tile_x’, [-1, 2])],None,17
No: 8 GFLOPS: 2.37/13.46 result: MeasureResult(costs=(0.1133100596,), error_no=0, all_cost=2.2167930603027344, timestamp=1607225793.595454) [(‘tile_y’, [-1, 8]), (‘tile_x’, [-1, 4])],None,23
No: 9 GFLOPS: 11.52/13.46 result: MeasureResult(costs=(0.0233022846,), error_no=0, all_cost=0.7279143333435059, timestamp=1607225795.1428313) [(‘tile_y’, [-1, 256]), (‘tile_x’, [-1, 32])],None,58
No: 10 GFLOPS: 14.67/14.67 result: MeasureResult(costs=(0.0182990712,), error_no=0, all_cost=0.7626948356628418, timestamp=1607225795.9127738) [(‘tile_y’, [-1, 64]), (‘tile_x’, [-1, 128])],None,76
最後,最好從快取檔案中應用歷史記錄,並檢查其正確性。可以matmul直接在 autotvm.apply_history_best上下文中呼叫該函式。當呼叫此函式時,它將使用其引數查詢排程上下文,並使用相同的引數獲得最佳配置。

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.asnumpy(), rtol=1e-2)

相關文章