CPU的自动调度矩阵乘法
这是一个有关如何对CPU使用自动调度程序的文档。
与依靠手动模板定义搜索空间的基于模板的autotvm不同，自动调度程序不需要任何模板。用户只需要编写计算声明，而无需任何调度命令或模板。自动调度程序可以自动生成较大的搜索空间，并在该空间中找到良好的调度。
本文以矩阵乘法为例。
注意，本文无法在Windows或最新版本的macOS上运行。要使其运行，需要将本文的内容包装在一个if name == “main”:块中。
import os

import numpy as np
import tvm
from tvm import te, auto_scheduler
定义计算
首先，定义带有偏差加法的矩阵的计算。该函数应返回输入/输出张量的列表。通过这些张量，自动调度器可以获取整个计算图。
@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]

创建搜索任务
然后，创建一个搜索任务，其中N = L = M = 1024且dtype =“ float32”。如果计算机支持avx指令，可以
• 将下面的“ 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])
接下来，为自动调度程序设置参数。
• 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)
出：
TTTTTTTTTT
可以降低调度，以便在自动调度后查看IR。自动调度程序正确执行优化，包括多层平铺，布局转换，并行化，矢量化，展开和运算符融合。
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)])
}
}
}
}
}
检查正确性并评估性能
构建二进制文件并检查其正确性和性能。
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
使用记录文件
搜索期间，所有测量记录都将转储到记录文件“ matmul.json”中。测量记录可用于重新应用搜索结果，继续搜索以及执行其它分析。
这是一个示例，其中从文件加载最佳调度，并打印等效的python调度API。这可用于调试和学习自动调度程序的行为。
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_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的线程池之间的冲突，因此无法在上面运行此行。运行tvm生成的二进制文件后，python的多处理库将永远挂起。必须确保在调用auot-scheduler的搜索之前，不要运行任何tvm生成的二进制文件。要运行上面的功能，应该注释掉“检查正确性和评估性能”部分中的所有代码。
应该在应用程序中注意这个问题。对于此问题，还有其他解决方法。例如，可以启动新线程/进程（使用内置的python库线程或多线程处理），并在新线程/进程中运行tvm二进制文件。这提供了隔离，并避免了主线程/进程中的冲突。还可以将auto_scheduler.LocalRPCMeasureContext用于自动调度程序，如GPU帮助（自动调度GPU的卷积层）中所示。
脚本的总运行时间：（1分钟50.410秒）
https://tvm.apache.org/docs/tutorials/auto_scheduler/tune_matmul_x86.html#sphx-glr-tutorials-auto-scheduler-tune-matmul-x86-py

CPU的自动调度矩阵乘法相关推荐

如何在CPU上优化GEMM矩阵乘法
如何在CPU上优化GEMM矩阵乘法 How to optimize GEMM on CPU (TL;DR) TVM 提供抽象接口,允许用户分别描述算法和算法的实现组织(所谓的调度).通常,在高性能调度 ...
ARM CPU神经网络自动调度
ARM CPU神经网络自动调度对特定设备和工作负载进行自动调整对于获得最佳性能至关重要.这是一个有关如何通过RPC使用自动调度器为ARM CPU调整整个神经网络的教程. 为了自动调整神经网络,将网络 ...
大佬是怎么优雅实现矩阵乘法的？
点击上方"视学算法",选择加"星标"或"置顶" 重磅干货,第一时间送达作者丨立交桥跳水冠军来源丨https://zhuanlan.zhi ...
ARM CPU自动调度神经网络
ARM CPU自动调度神经网络对特定设备和工作负载进行自动调度,对于获得最佳性能至关重要.通过RPC使用自动调度器为ARM CPU调度整个神经网络. 为了自动调度神经网络,将网络划分为小的子图,进行 ...
为x86 CPU自动调度神经网络
为x86 CPU自动调度神经网络对特定设备和工作负载进行自动调试对于获得最佳性能至关重要.这是有关如何使用自动调度器为x86 CPU调试整个神经网络的文档. 为了自动调试神经网络,将网络划分为小的子 ...
GPU与CPU版本的矩阵乘法对比
由于刚刚开始学习Cuda,还没有整理出一个完整的Cuda类,只是在Nvidia提供的kenerl架构上做修改. 但用于初体验GPU给我们带来的好处也绰绰有余了. 直接贴代码: /*矩阵乘法,CPU版本 ...
Java实现矩阵运算——矩阵乘法、矩阵转置、自动填充矩阵行
在做大数据或人工智能开发的过程做难免会遇到矩阵运算,本文在这里给大家实现一个简单的矩阵运算,请看下代码: package test;/*** 矩阵运算* * @author Administrator ...
智源青年科学家梁云：异构系统中张量计算的自动调度和优化框架
与6位图灵奖得主和100多位专家共同探讨人工智能的下一个十年长按图片,内行盛会,首次免费注册北京智源大会倒计时:9天计算机体系结构领域国际顶级会议每次往往仅录用几十篇论文,录用率在20%左右 ...
HLS：矩阵乘法单元设计与SDK测试
目录一.引言二.程序框架三.初步设计四.报告分析五.优化操作六.接口优化七.上板测试八.补充部分九.时间与参考一.引言矩阵乘法,涉及数组优化.循环优化和接口优化等.是一个学习HL ...

CPU的自动调度矩阵乘法