TVM：使用 Auto-scheduling 來優化算子

在本教程中，我們將展示 TVM 的 Auto-scheduling 功能如何在無需編寫自定義模板的情況下找到最佳 schedule。

與基于模板的 AutoTVM 依賴手動模板定義搜索空間不同，auto-scheduler 不需要任何模板。 用戶只需編寫計算聲明，無需任何調度命令或模板。 auto-scheduler 可以自動生成一個大的搜索空間，并在該空間中找到一個好的 schedule。

我們在本教程中同樣使用矩陣乘法作為示例。

import osimport numpy as np import tvm from tvm import te, auto_scheduler

定義矩陣乘法

首先，我們定義一個帶有偏置的矩陣乘法。 請注意，這使用了 TVM 張量表達式語言中可用的標準操作。 主要區別在于在函數定義的開始使用了 auto_sceduler 裝飾器。 該函數應返回輸入/輸出張量列表。 從這些張量中，自動調度器可以獲得整個計算圖。

@auto_scheduler.register_workload # Note the auto_scheduler decorator def matmul_add(N, L, M, dtype):A = te.placeholder((N, L), name="A", dtype=dtype)B = te.placeholder((L, M), name="B", dtype=dtype)C = te.placeholder((N, M), name="C", dtype=dtype)k = te.reduce_axis((0, L), name="k")matmul = te.compute((N, M),lambda i, j: te.sum(A[i, k] * B[k, j], axis=k),name="matmul",attrs={"layout_free_placeholders": [B]}, # enable automatic layout transform for tensor B)out = te.compute((N, M), lambda i, j: matmul[i, j] + C[i, j], name="out")return [A, B, C, out]

創建搜索任務

定義函數后，我們現在可以創建供 auto_scheduler 搜索的任務。 我們指定此矩陣乘法的特定參數，在本例中為 1024x1024 大小的方陣的乘法。 然后我們創建一個搜索任務，其中 N=L=M=1024 ，數據類型為 ”float32”。

target = tvm.target.Target("llvm") N = L = M = 1024 task = tvm.auto_scheduler.SearchTask(func=matmul_add, args=(N, L, M, "float32"), target=target)# Inspect the computational graph print("Computational DAG:") print(task.compute_dag)

注意：自定義 target 可以提高性能

為了讓 TVM 充分利用特定硬件平臺，您需要手動指定 CPU 功能。 例如： - 將下面的“llvm”替換為“llvm -mcpu=core-avx2”以啟用 AVX2 - 將下面的“llvm”替換為“llvm -mcpu=skylake-avx512”以啟用 AVX-512

此處輸出：

Computational DAG: A = PLACEHOLDER [1024, 1024] B = PLACEHOLDER [1024, 1024] matmul(i, j) += (A[i, k]*B[k, j]) C = PLACEHOLDER [1024, 1024] out(i, j) = (matmul[i, j] + C[i, j])

為 Auto-Scheduler 設置參數

接下來，我們為自動調度程序設置參數。

• num_measure_trials 是我們在搜索過程中可以使用的測量試驗次數。 為了快速演示，我們在本教程中僅進行了 10 次試驗。 在實踐中，1000 是一個很好的搜索收斂值。 您可以根據您的時間預算進行更多試驗。

• 此外，我們使用 RecordToFile 將測量記錄記錄到文件 matmul.json 中。 測量記錄可用于最佳查詢歷史記錄、恢復搜索以及稍后進行更多分析。

• 有關更多參數，請參閱 auto_scheduler.TuningOptions

log_file = "matmul.json" tune_option = auto_scheduler.TuningOptions(num_measure_trials=10,measure_callbacks=[auto_scheduler.RecordToFile(log_file)],verbose=2, )

運行搜索

現在我們準備好所有輸入。 很簡單，不是嗎？ 我們可以開始搜索并讓自動調度程序發揮它的魔力。 經過一些測量試驗后，我們可以從日志文件中加載最佳計劃并應用它。

# Run auto-tuning (search) task.tune(tune_option) # Apply the best schedule sch, args = task.apply_best(log_file)

檢查優化過的 Schedule

我們可以在 auto-scheduling 后降低（lower）schedule 以查看 IR。 auto-schduling 程序正確執行優化，包括多級平鋪、布局轉換、并行化、矢量化、展開和算子融合。

print("Lowered TIR:") print(tvm.lower(sch, args, simple_mode=True))

此處輸出：

Lowered TIR: primfn(A_1: handle, B_1: handle, C_1: handle, out_1: handle) -> ()attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}buffers = {out: Buffer(out_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C, out_1: out} {allocate(auto_scheduler_layout_transform: Pointer(global float32), float32, [1048576]), storage_scope = global {for (ax0.ax1.fused.ax2.fused: int32, 0, 128) "parallel" {for (ax4: int32, 0, 256) {for (ax6: int32, 0, 4) {for (ax7: int32, 0, 8) {auto_scheduler_layout_transform[((((ax0.ax1.fused.ax2.fused*8192) + (ax4*32)) + (ax6*8)) + ax7)] = (float32*)B_2[((((ax4*4096) + (ax6*1024)) + (ax0.ax1.fused.ax2.fused*8)) + ax7)]}}}}for (i.outer.outer.j.outer.outer.fused: int32, 0, 16384) "parallel" {allocate(matmul: Pointer(global float32x8), float32x8, [4]), storage_scope = global;for (i.outer.inner: int32, 0, 2) {matmul[ramp(0, 1, 8)] = broadcast(0f32, 8)matmul[ramp(8, 1, 8)] = broadcast(0f32, 8)matmul[ramp(16, 1, 8)] = broadcast(0f32, 8)matmul[ramp(24, 1, 8)] = broadcast(0f32, 8)for (k.outer: int32, 0, 256) {for (k.inner: int32, 0, 4) {matmul[ramp(0, 1, 8)] = ((float32x8*)matmul[ramp(0, 1, 8)] + (broadcast((float32*)A_2[((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))matmul[ramp(8, 1, 8)] = ((float32x8*)matmul[ramp(8, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 1024)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))matmul[ramp(16, 1, 8)] = ((float32x8*)matmul[ramp(16, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 2048)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))matmul[ramp(24, 1, 8)] = ((float32x8*)matmul[ramp(24, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 3072)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))}}for (i.inner: int32, 0, 4) {out_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (i.inner*1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)*8)), 1, 8)] = ((float32x8*)matmul[ramp((i.inner*8), 1, 8)] + (float32x8*)C_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (i.inner*1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)*8)), 1, 8)])}}}} }

檢查正確性并評估性能

我們構建二進制文件并檢查其正確性和性能。

func = tvm.build(sch, args, target) a_np = np.random.uniform(size=(N, L)).astype(np.float32) b_np = np.random.uniform(size=(L, M)).astype(np.float32) c_np = np.random.uniform(size=(N, M)).astype(np.float32) out_np = a_np.dot(b_np) + c_npdev = tvm.cpu() a_tvm = tvm.nd.array(a_np, device=dev) b_tvm = tvm.nd.array(b_np, device=dev) c_tvm = tvm.nd.array(c_np, device=dev) out_tvm = tvm.nd.empty(out_np.shape, device=dev) func(a_tvm, b_tvm, c_tvm, out_tvm)# Check results np.testing.assert_allclose(out_np, out_tvm.numpy(), rtol=1e-3)# Evaluate execution time. evaluator = func.time_evaluator(func.entry_name, dev, min_repeat_ms=500) print("Execution time of this operator: %.3f ms"% (np.median(evaluator(a_tvm, b_tvm, c_tvm, out_tvm).results) * 1000) )

此處輸出：

Execution time of this operator: 45.418 ms

使用記錄文件

在搜索過程中，所有的測量記錄都被記錄到記錄文件“matmul.json”中。 測量記錄可用于重新應用搜索結果、恢復搜索和執行其他分析。

這是一個示例，我們從文件加載最佳 schedule，并打印等效的 Python schedule API。 這可用于調試和學習 auto-scheduling 程序的行為。

print("Equivalent python schedule:") print(task.print_best(log_file))

此處輸出：

Equivalent python schedule: matmul_i, matmul_j, matmul_k = tuple(matmul.op.axis) + tuple(matmul.op.reduce_axis) out_i, out_j = tuple(out.op.axis) + tuple(out.op.reduce_axis) matmul_i_o_i, matmul_i_i = s[matmul].split(matmul_i, factor=4) matmul_i_o_o_i, matmul_i_o_i = s[matmul].split(matmul_i_o_i, factor=1) matmul_i_o_o_o, matmul_i_o_o_i = s[matmul].split(matmul_i_o_o_i, factor=2) matmul_j_o_i, matmul_j_i = s[matmul].split(matmul_j, factor=8) matmul_j_o_o_i, matmul_j_o_i = s[matmul].split(matmul_j_o_i, factor=1) matmul_j_o_o_o, matmul_j_o_o_i = s[matmul].split(matmul_j_o_o_i, factor=1) matmul_k_o, matmul_k_i = s[matmul].split(matmul_k, factor=4) s[matmul].reorder(matmul_i_o_o_o, matmul_j_o_o_o, matmul_i_o_o_i, matmul_j_o_o_i, matmul_k_o, matmul_i_o_i, matmul_j_o_i, matmul_k_i, matmul_i_i, matmul_j_i) out_i_o_i, out_i_i = s[out].split(out_i, factor=4) out_i_o_o, out_i_o_i = s[out].split(out_i_o_i, factor=2) out_j_o_i, out_j_i = s[out].split(out_j, factor=8) out_j_o_o, out_j_o_i = s[out].split(out_j_o_i, factor=1) s[out].reorder(out_i_o_o, out_j_o_o, out_i_o_i, out_j_o_i, out_i_i, out_j_i) s[matmul].compute_at(s[out], out_j_o_i) out_i_o_o_j_o_o_fused = s[out].fuse(out_i_o_o, out_j_o_o) s[out].parallel(out_i_o_o_j_o_o_fused) s[matmul].pragma(matmul_i_o_o_o, "auto_unroll_max_step", 8) s[matmul].pragma(matmul_i_o_o_o, "unroll_explicit", True) s[matmul].vectorize(matmul_j_i) s[out].vectorize(out_j_i)

一個更復雜的例子是恢復搜索。 在這種情況下，我們需要自己創建搜索策略和成本模型，并通過日志文件恢復搜索策略和成本模型的狀態。 在下面的示例中，我們恢復狀態并再進行 5 次試驗。

def resume_search(task, log_file):print("Resume search:")cost_model = auto_scheduler.XGBModel()cost_model.update_from_file(log_file)search_policy = auto_scheduler.SketchPolicy(task, cost_model, init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file)])tune_option = auto_scheduler.TuningOptions(num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file)])task.tune(tune_option, search_policy=search_policy)resume_search(task, log_file)

此處輸出：

Resume search: /usr/local/lib/python3.6/dist-packages/xgboost/training.py:17: UserWarning: Old style callback is deprecated. See: https://xgboost.readthedocs.io/en/latest/python/callbacks.htmlwarnings.warn(f'Old style callback is deprecated. See: {link}', UserWarning)

總結

在本教程中，我們展示了如何使用 TVM Auto-Scheduler 自動優化矩陣乘法，而無需指定搜索模板。 它結束了一系列從張量表達式 (TE) 語言開始的示例，這些示例演示了 TVM 如何優化計算操作。

Ref：

https://tvm.apache.org/docs/tutorial/auto_scheduler_matmul_x86.html

總結

以上是生活随笔為你收集整理的TVM：使用 Auto-scheduling 来优化算子的全部內容，希望文章能夠幫你解決所遇到的問題。

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