當(dāng)前位置：首頁 > 人文社科 > 生活经验 >内容正文

生活经验

CPU的自动调度矩阵乘法

發(fā)布時(shí)間：2023/11/28 生活经验 35 豆豆

生活随笔收集整理的這篇文章主要介紹了 CPU的自动调度矩阵乘法小編覺得挺不錯(cuò)的,現(xiàn)在分享給大家,幫大家做個(gè)參考.

CPU的自動調(diào)度矩陣乘法
這是一個(gè)有關(guān)如何對CPU使用自動調(diào)度程序的文檔。
與依靠手動模板定義搜索空間的基于模板的autotvm不同，自動調(diào)度程序不需要任何模板。用戶只需要編寫計(jì)算聲明，而無需任何調(diào)度命令或模板。自動調(diào)度程序可以自動生成較大的搜索空間，并在該空間中找到良好的調(diào)度。
本文以矩陣乘法為例。
注意，本文無法在Windows或最新版本的macOS上運(yùn)行。要使其運(yùn)行，需要將本文的內(nèi)容包裝在一個(gè)if name == “main”:塊中。
import os

import numpy as np
import tvm
from tvm import te, auto_scheduler
定義計(jì)算
首先，定義帶有偏差加法的矩陣的計(jì)算。該函數(shù)應(yīng)返回輸入/輸出張量的列表。通過這些張量，自動調(diào)度器可以獲取整個(gè)計(jì)算圖。
@auto_scheduler.register_workload
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]

創(chuàng)建搜索任務(wù)
然后，創(chuàng)建一個(gè)搜索任務(wù)，其中N = L = M = 1024且dtype =“ float32”。如果計(jì)算機(jī)支持avx指令，可以
? 將下面的“ llvm”替換為“ llvm -mcpu = core-avx2”以啟用AVX2
? 將下面的“ llvm”替換為“ llvm -mcpu = skylake-avx512”以啟用AVX-512
target = tvm.target.Target(“l(fā)lvm”)
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)
出：
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])
接下來，為自動調(diào)度程序設(shè)置參數(shù)。
? num_measure_trials是在搜索過程中可以使用的測量試驗(yàn)的數(shù)量。為了快速演示，在本文中僅進(jìn)行10次試用。實(shí)際上，1000是搜索收斂的一個(gè)很好的值。可以根據(jù)自己的時(shí)間預(yù)算進(jìn)行更多試驗(yàn)。
? 此外，還用RecordToFile將測量記錄轉(zhuǎn)儲到文件matmul.json中。測量記錄可用于最好地查詢歷史記錄，恢復(fù)搜索以及以后進(jìn)行更多分析。
? 查看更多參數(shù)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,
)
運(yùn)行搜索
現(xiàn)在準(zhǔn)備好所有輸入。開始搜索，讓自動調(diào)度程序發(fā)揮作用。經(jīng)過一些測量試驗(yàn)后，可以從日志文件中加載最佳調(diào)度并應(yīng)用它。

Run auto-tuning (search)

task.tune(tune_option)

Apply the best schedule

sch, args = task.apply_best(log_file)
出：
TTTTTTTTTT
可以降低調(diào)度，以便在自動調(diào)度后查看IR。自動調(diào)度程序正確執(zhí)行優(yōu)化，包括多層平鋪，布局轉(zhuǎn)換，并行化，矢量化，展開和運(yùn)算符融合。
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 = {“global_symbol”: “main”, “tir.noalias”: True}
buffers = {out: Buffer(out_2: Pointer(float32), float32, [1024, 1024], []),
C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
buffer_map = {A_1: A, B_1: B, C_1: C, out_1: out} {
attr [auto_scheduler_layout_transform: Pointer(float32)] “storage_scope” = “global”;
allocate(auto_scheduler_layout_transform, float32, [1048576]) {
for (ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused: int32, 0, 131072) “parallel” {
for (ax7: int32, 0, 8) {
auto_scheduler_layout_transform[((ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused8) + ax7)] = (float32)B_2[(((floormod(ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused, 1024)1024) + (floordiv(ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused, 1024)8)) + ax7)]
}
}
for (i.outer.outer.j.outer.outer.fused: int32, 0, 16384) “parallel” {
attr [matmul: Pointer(float32)] “storage_scope” = “global”;
allocate(matmul, float32x8, [4]);
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.inner4096)) + (k.outer4)) + k.inner)], 8)(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)8192) + (k.outer32)) + (k.inner8)), 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.inner4096)) + (k.outer4)) + k.inner) + 1024)], 8)(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)8192) + (k.outer32)) + (k.inner8)), 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.inner4096)) + (k.outer4)) + k.inner) + 2048)], 8)(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)8192) + (k.outer32)) + (k.inner8)), 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.inner4096)) + (k.outer4)) + k.inner) + 3072)], 8)(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)8192) + (k.outer32)) + (k.inner8)), 1, 8)]))
}
}
for (i.inner: int32, 0, 4) {
out_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)8192) + (i.outer.inner4096)) + (i.inner1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)8)), 1, 8)] = ((float32x8)matmul[ramp((i.inner8), 1, 8)] + (float32x8)C_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)8192) + (i.outer.inner4096)) + (i.inner*1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)*8)), 1, 8)])
}
}
}
}
}
檢查正確性并評估性能
構(gòu)建二進(jìn)制文件并檢查其正確性和性能。
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_np

ctx = tvm.cpu()
a_tvm = tvm.nd.array(a_np, ctx=ctx)
b_tvm = tvm.nd.array(b_np, ctx=ctx)
c_tvm = tvm.nd.array(c_np, ctx=ctx)
out_tvm = tvm.nd.empty(out_np.shape, ctx=ctx)
func(a_tvm, b_tvm, c_tvm, out_tvm)

Check results

np.testing.assert_allclose(out_np, out_tvm.asnumpy(), rtol=1e-3)

Evaluate execution time.

evaluator = func.time_evaluator(func.entry_name, ctx, 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: 22.426 ms
使用記錄文件
搜索期間，所有測量記錄都將轉(zhuǎn)儲到記錄文件“ matmul.json”中。測量記錄可用于重新應(yīng)用搜索結(jié)果，繼續(xù)搜索以及執(zhí)行其它分析。
這是一個(gè)示例，其中從文件加載最佳調(diào)度，并打印等效的python調(diào)度API。這可用于調(diào)試和學(xué)習(xí)自動調(diào)度程序的行為。
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)
一個(gè)更復(fù)雜的示例是繼續(xù)搜索。在這種情況下，需要自己創(chuàng)建搜索策略和成本模型，并使用日志文件恢復(fù)搜索策略和成本模型的狀態(tài)。在下面的示例中，恢復(fù)狀態(tài)并進(jìn)行5次以上的試用。
def resume_search(task, log_file_name):
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(log_file_name)
search_policy = auto_scheduler.SketchPolicy(
task,
cost_model,
init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file_name)],
)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file_name)]
)
task.tune(tune_option, search_policy=search_policy)

resume_search(task, log_file)

注意
由于python的多處理和tvm的線程池之間的沖突，因此無法在上面運(yùn)行此行。運(yùn)行tvm生成的二進(jìn)制文件后，python的多處理庫將永遠(yuǎn)掛起。必須確保在調(diào)用auot-scheduler的搜索之前，不要運(yùn)行任何tvm生成的二進(jìn)制文件。要運(yùn)行上面的功能，應(yīng)該注釋掉“檢查正確性和評估性能”部分中的所有代碼。
應(yīng)該在應(yīng)用程序中注意這個(gè)問題。對于此問題，還有其他解決方法。例如，可以啟動新線程/進(jìn)程（使用內(nèi)置的python庫線程或多線程處理），并在新線程/進(jìn)程中運(yùn)行tvm二進(jìn)制文件。這提供了隔離，并避免了主線程/進(jìn)程中的沖突。還可以將auto_scheduler.LocalRPCMeasureContext用于自動調(diào)度程序，如GPU幫助（自動調(diào)度GPU的卷積層）中所示。
腳本的總運(yùn)行時(shí)間：（1分鐘50.410秒）
https://tvm.apache.org/docs/tutorials/auto_scheduler/tune_matmul_x86.html#sphx-glr-tutorials-auto-scheduler-tune-matmul-x86-py

總結(jié)

以上是生活随笔為你收集整理的CPU的自动调度矩阵乘法的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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