【TVM 學習資料】用 Schedule 模板和 AutoTVM 最佳化運算元

超神經HyperAI發表於2023-02-13

完整 TVM 中文文件，訪問→TVM 中文站
作者：Lianmin Zheng，Chris Hoge

本教程將展示如何用 TVM 張量表示式（TE）語言編寫 schedule 模板，並透過 AutoTVM 對模板進行搜尋，從而找到最佳 schedule。這個自動最佳化張量計算的過程被稱為 Auto-Tuning。

本教程基於前面的 TE 編寫矩陣乘法教程 設立。

auto-tuning 包括兩個步驟：

第一步：定義搜尋空間。
第二步：執行搜尋演算法來探索這個空間。
透過本教程可以瞭解如何在 TVM 中執行這兩個步驟。整個工作流程由一個矩陣乘法示例來說明。

備註
注意，本教程不會在 Windows 或最新版本的 macOS 上執行。如需執行，請將本教程的主體放在 if name == "__main__": 程式碼塊中。

安裝依賴

要在 TVM 中使用 autotvm 包，需安裝一些額外的依賴。

pip3 install --user psutil xgboost cloudpickle

為了讓 TVM 在調優過程中執行更快，建議使用 Cython 作為 TVM 的 FFI。在 TVM 的根目錄下，執行：

pip3 install --user cython
sudo make cython3

現在我們一起來看如何用 Python 程式碼實現。首先匯入所需的包：

import logging
import sys

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

# 模組名叫 `autotvm`
from tvm import autotvm

TE 的基本矩陣乘法

回想一下用 TE 進行矩陣乘法的基本實現，下面做一些改變。將矩陣乘法放在 Python 函式定義中。簡單起見，重點關注拆分的最佳化，將重新排序的塊大小設為固定值。

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)

    # 排程
    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]

用 AutoTVM 進行矩陣乘法

前面的排程程式碼用常量“8”作為迴圈切分因子，但是它可能不是最佳的。因為最佳的迴圈切分因子取決於真實的硬體環境和輸入 shape。

如果希望排程程式碼能夠在更廣泛的輸入 shape 和目標硬體上可移植，最好定義一組候選值，並根據目標硬體上的評估結果選擇最佳值。

autotvm 中可以為這種值定義一個可調引數，或者一個 "knob"。

基本矩陣乘法模板

以下示例將演示，如何為 split 排程操作的 block 大小建立一個可調的引數集。

# Matmul V1: 列出候選值
@autotvm.template("tutorial/matmul_v1")  # 1. 使用裝飾器
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)

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

    # 2. 獲取 config 物件
    cfg = autotvm.get_config()

    # 3. 定義搜尋空間
    cfg.define_knob("tile_y", [1, 2, 4, 8, 16])
    cfg.define_knob("tile_x", [1, 2, 4, 8, 16])

    # 4. 根據 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]

下面將對前面的排程程式碼作出四個修改，然後得到一個可調的“模板”。一一解釋這些修改：

  1. 使用裝飾器將此函式標記為簡單模板。
  2. 獲取 config 物件：將 cfg 視為此函式的引數，但我們以另外的方式獲取它。cfg 引數使得這個函式不再是一個確定的 schedule。將不同的配置傳遞給這個函式，可以得到不同的 schedule。這種使用配置物件的函式稱為“模板”。

為使模板函式更精煉，可在單個函式中定義引數搜尋空間：

  • 用一組值來定義搜尋空間。將 cfg 轉為 ConfigSpace 物件，收集此函式中的所有可調 knob，然後從中構建一個搜尋空間。
  • 根據空間中的實體進行排程。將 cfg 轉為 ConfigEntity 物件，當它被轉為 ConfigEntity 後，會忽略所有空間定義
    API（即 cfg.define_XXXXX(...)），但會儲存所有可調 knob 的確定值，並根據這些值進行排程。

在 auto-tuning 的過程中，首先用 ConfigSpace 物件呼叫這個模板來構建搜尋空間，然後在構建的空間中用不同的 ConfigEntity 呼叫這個模板，來得到不同的 schedule。最後，我們將評估由不同 schedule 生成的程式碼，然後選擇最佳的 schedule。

  1. 定義兩個可調 knob。第一個是 tile_y，它有 5 個可能值。第二個是 tile_x，它和前者具有相同的可能值。這兩個 knob 是獨立的，所以它們跨越大小為 25 = 5x5 的搜尋空間。
  2. 配置 knob 被傳遞給 split 排程操作，然後可以根據之前在 cfg 中定義的 5x5 確定值進行排程。

帶有高階引數 API 的矩陣乘法模板

前面的模板手動列出了 konb 的所有可能值，它是用來定義空間的最底層 API，顯示列出了要搜尋的引數空間。這裡推薦使用另一組更高階的 API，它可以更簡單、更智慧地定義搜尋空間。

下面的示例用 ConfigSpace.define_split 來定義拆分 knob。它列舉了所有可能的拆分 axis 和構造空間的方法。

同時，ConfigSpace.define_reorder 用於對 knob 重新排序，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)

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

    ##### 開始定義空間 #####
    cfg = autotvm.get_config()
    cfg.define_split("tile_y", y, num_outputs=2)
    cfg.define_split("tile_x", x, num_outputs=2)
    ##### 結束定義空間 #####

    # 根據 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.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 個可能值。

排程過程中，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 步：使用 AutoTVM 最佳化矩陣乘法

第 1 步編寫的矩陣乘法模板，可對拆分的 schedule 中的塊大小進行引數化。透過第 1 步，可以實現對這個引數空間進行搜尋。下一步是選擇一個調優器來指導如何對空間進行探索。

TVM 的自動調優器

調優器的任務可用以下虛擬碼來描述：

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

調優器可採取不同的策略來計劃下一批配置，包括：

  • tvm.autotvm.tuner.RandomTuner ：以隨機順序列舉空間
  • tvm.autotvm.tuner.GridSearchTuner ：以網格搜尋順序列舉空間
  • tvm.autotvm.tuner.GATuner ：使用遺傳演算法搜尋空間
  • tvm.autotvm.tuner.XGBTuner ：用基於模型的方法訓練一個 XGBoost 模型，來預測降級 IR
    的速度，並根據預測值選擇下一批配置。

可根據空間大小、時間預算和其他因素來選擇調優器。例如，如果你的空間非常小（小於 1000），則網格搜尋調優器或隨機調優器就夠了。如果你的空間在 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 次試驗進行演示。實際上可以根據自己的時間預算進行更多試驗。調優結果會記錄到日誌檔案中。這個檔案可用於選擇之後發現的調優器的最佳配置。

# 記錄 config（為了將 tuning 日誌列印到螢幕）
logging.getLogger("autotvm").setLevel(logging.DEBUG)
logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

評估配置有兩個步驟：構建和執行。預設用所有 CPU core 來編譯程式。然後依次進行評估。為了減少方差，對 5 次評估結果取平均值。

measure_option = autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(number=5))

# 用 RandomTuner 開始調優, 日誌記錄到 `matmul.log` 檔案中
# 可用 XGBTuner 來替代.
tuner = autotvm.tuner.RandomTuner(task)
tuner.tune(
    n_trial=10,
    measure_option=measure_option,
    callbacks=[autotvm.callback.log_to_file("matmul.log")],
)

輸出結果：

waiting for device...
device available
Get devices for measurement successfully!
No: 1   GFLOPS: 8.48/8.48       result: MeasureResult(costs=(0.0316434228,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.638512134552002, timestamp=1657225928.6342561)        [('tile_y', [-1, 1]), ('tile_x', [-1, 256])],None,80
No: 2   GFLOPS: 2.30/8.48       result: MeasureResult(costs=(0.1165478966,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.0105199813842773, timestamp=1657225930.6636436)       [('tile_y', [-1, 4]), ('tile_x', [-1, 8])],None,32
No: 3   GFLOPS: 11.82/11.82     result: MeasureResult(costs=(0.0227097348,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5589795112609863, timestamp=1657225931.7059512)       [('tile_y', [-1, 64]), ('tile_x', [-1, 32])],None,56
No: 4   GFLOPS: 1.66/11.82      result: MeasureResult(costs=(0.1616202114,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.6911513805389404, timestamp=1657225934.9635096)       [('tile_y', [-1, 1]), ('tile_x', [-1, 4])],None,20
No: 5   GFLOPS: 3.65/11.82      result: MeasureResult(costs=(0.073561817,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.3051848411560059, timestamp=1657225936.3988533)        [('tile_y', [-1, 256]), ('tile_x', [-1, 16])],None,48
No: 6   GFLOPS: 1.85/11.82      result: MeasureResult(costs=(0.1452834464,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.5179028511047363, timestamp=1657225938.961955)        [('tile_y', [-1, 512]), ('tile_x', [-1, 4])],None,29
No: 7   GFLOPS: 0.87/11.82      result: MeasureResult(costs=(0.30933780240000003,), error_no=MeasureErrorNo.NO_ERROR, all_cost=5.067087888717651, timestamp=1657225944.589149)  [('tile_y', [-1, 512]), ('tile_x', [-1, 2])],None,19
No: 8   GFLOPS: 10.53/11.82     result: MeasureResult(costs=(0.025489421,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5452830791473389, timestamp=1657225945.1592515)        [('tile_y', [-1, 4]), ('tile_x', [-1, 64])],None,62
No: 9   GFLOPS: 1.58/11.82      result: MeasureResult(costs=(0.16960762680000002,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.8109781742095947, timestamp=1657225948.0900776)        [('tile_y', [-1, 2]), ('tile_x', [-1, 2])],None,11
No: 10  GFLOPS: 2.42/11.82      result: MeasureResult(costs=(0.11083148779999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.8757600784301758, timestamp=1657225950.0266354)        [('tile_y', [-1, 4]), ('tile_x', [-1, 4])],None,22

調優完成後，可從日誌檔案中選擇具有最佳評估效能的配置，並用相應引數來編譯 schedule。快速驗證 schedule 是否產生了正確的結果，可直接在 autotvm.apply_history_best 上下文中呼叫 matmul 函式，它會用引數查詢排程上下文，然後可用相同的引數獲取最優配置。

# 從日誌檔案中應用歷史最佳
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)

# 驗證正確性
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)

輸出結果：

Finish loading 10 records

總結

本教程展示瞭如何構建運算元模板，使得 TVM 能夠搜尋引數空間，並選擇最佳化的排程配置。為了更深入地瞭解其工作原理，推薦基於 :ref: 張量表示式入門 <tensorexpr_get_started> 教程中演示的排程操作，向排程新增新的搜尋引數。接下來的章節將演示 AutoScheduler，它是TVM 中一種最佳化常用運算元的方法，同時無需使用者提供自定義的模板。

相關文章