本文翻译自Optimizing Operators with Auto-scheduling — tvm 0.9.dev0 documentation
在本教程中，我们将展示TVM的自动调度功能如何在不编写自定义模板的情况下找到最佳调度。
基于模板的AutoTVM通过手写模板来定义搜索空间，而自动调度器不需要任何模板。用户只需要编写计算声明，不需要任何调度命令或模板。自动调度器可以自动生成一个较大的搜索空间，并在这个空间中找到一个好的调度。
我们在本教程中使用矩阵乘法作为例子。
注意：本教程的代码不会在Windows或最新版本的macOS上运行。要让它运行，你需要将本教程的代码主体包在if __name__ == "__main__":中。
import os
import numpy as np
import tvm
from tvm import te, auto_scheduler定义矩阵乘法
首先，我们定义一个带有偏置加法的矩阵乘法。注意，这里使用了TVM张量表达式语言中的标准操作。主要的区别是在函数定义的顶部使用了register_workload装饰器。这个函数应该返回一个输入/输出张量的列表。根据这些张量，自动调度器可以得到整个计算图。
@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, dtype= " float32 "
使用自定义目标提高性能：
为了让TVM充分利用特定的硬件平台，您需要手动指定您的CPU能力。例如:
- 将下面的llvm替换为llvm -mcpu=core-avx2来启用AVX2
- 将下面的llvm替换为llvm -mcpu=skylake-avx512来启用AVX-512
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)输出：
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设置参数
接下来，我们为auto-scheduler设置参数。
Num_measure_trials是我们在搜索过程中可以使用的测量试验的次数。在本教程中，为了快速演示，我们只做10次试验。在实际应用中，1000是一个很好的收敛值。你可以根据你的时间预算做更多的试验。
此外，我们使用RecordToFile将测量记录写入matmull .json文件中。测量记录可用于查询历史最佳，恢复搜索，并在以后进行更多的分析。
更多参数请参阅调优选项tvm.auto_scheduler — tvm 0.9.dev0 documentation
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)检查优化后的调度
自动排程后，我们可以低级化调度以查看 auto-scheduling的IR。自动调度程序正确地执行优化，包括多层tiling、layout、并行化、向量化、展开和算子融合。
print("Lowered TIR:")
print(tvm.lower(sch, args, simple_mode=True))输出：
Lowered TIR:
@main = 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 = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], []),
out: Buffer(out_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C, out_1: out}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], []), out_1: out_3: Buffer(out_2, float32, [1024, 1024], [])} {
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_1: Buffer(auto_scheduler_layout_transform, float32, [1048576], [])[((((ax0.ax1.fused.ax2.fused*8192) + (ax4*32)) + (ax6*8)) + ax7)] = B[((((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_1: Buffer(matmul, float32x8, [4], [])[0] = broadcast(0f32, 8)
matmul_1[1] = broadcast(0f32, 8)
matmul_1[2] = broadcast(0f32, 8)
matmul_1[3] = broadcast(0f32, 8)
for (k.outer: int32, 0, 256) {
for (k.inner: int32, 0, 4) {
let cse_var_2: int32 = (((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8))
let cse_var_1: int32 = ((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner)
{
matmul_1[0] = (matmul_1[0] + (broadcast(A[cse_var_1], 8)*auto_scheduler_layout_transform_1[ramp(cse_var_2, 1, 8)]))
matmul_1[1] = (matmul_1[1] + (broadcast(A[(cse_var_1 + 1024)], 8)*auto_scheduler_layout_transform_1[ramp(cse_var_2, 1, 8)]))
matmul_1[2] = (matmul_1[2] + (broadcast(A[(cse_var_1 + 2048)], 8)*auto_scheduler_layout_transform_1[ramp(cse_var_2, 1, 8)]))
matmul_1[3] = (matmul_1[3] + (broadcast(A[(cse_var_1 + 3072)], 8)*auto_scheduler_layout_transform_1[ramp(cse_var_2, 1, 8)]))
}
}
}
for (i.inner: int32, 0, 4) {
let cse_var_3: int32 = ((((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))
out[ramp(cse_var_3, 1, 8)] = (matmul_1[i.inner] + C[ramp(cse_var_3, 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_np
dev = 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: 93.599 ms使用记录文件
在搜索过程中，所有测量记录都被记录到文件matmull .json 中。测量记录可用于重新应用搜索结果，恢复搜索，并执行其他分析。
下面是一个例子，我们从一个文件加载最佳调度，并打印等效的python调度API。这可以用于调试和学习auto-scheduler的行为。
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.7/dist-packages/xgboost/training.py:17: UserWarning: Old style callback is deprecated. See: https://xgboost.readthedocs.io/en/latest/python/callbacks.html
warnings.warn(f'Old style callback is deprecated. See: {link}', UserWarning)总结
在本教程中，我们展示了如何使用TVM Auto-Scheduler自动优化矩阵乘法，而不需要指定搜索模板。本文结束了一系列例子，这些例子从张量表达式(TE)语言开始，展示了TVM如何优化计算操作。
扫一扫
709
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
