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)

出:

*T*T*T*T*T*T*T*T*T*T

可以降低调度,以便在自动调度后查看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.fused*8) + 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.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_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的自动调度矩阵乘法的更多相关文章

  1. ARM CPU自动调度神经网络

    ARM CPU自动调度神经网络 对特定设备和工作负载进行自动调度,对于获得最佳性能至关重要.通过RPC使用自动调度器为ARM CPU调度整个神经网络. 为了自动调度神经网络,将网络划分为小的子图,进行 ...

  2. 为x86 CPU自动调度神经网络

    为x86 CPU自动调度神经网络 对特定设备和工作负载进行自动调试对于获得最佳性能至关重要.这是有关如何使用自动调度器为x86 CPU调试整个神经网络的文档. 为了自动调试神经网络,将网络划分为小的子 ...

  3. OpenGL学习进程(12)第九课:矩阵乘法实现3D变换

    本节是OpenGL学习的第九个课时,下面将详细介绍OpenGL的多种3D变换和如何操作矩阵堆栈.     (1)3D变换: OpenGL中绘制3D世界的空间变换包括:模型变换.视图变换.投影变换和视口 ...

  4. 4-2.矩阵乘法的Strassen算法详解

    题目描述 请编程实现矩阵乘法,并考虑当矩阵规模较大时的优化方法. 思路分析 根据wikipedia上的介绍:两个矩阵的乘法仅当第一个矩阵B的列数和另一个矩阵A的行数相等时才能定义.如A是m×n矩阵和B ...

  5. 2.3CUDA矩阵乘法

    CPU 矩阵乘法 能相乘的两个矩阵,必须满足一个矩阵的行数和第二个矩阵的列数相同. A(N*P) * B(P*M) = C(N*M). 其中P是行数,N是列数, 从宽高的角度来说,即 A的宽度和B的高 ...

  6. 矩阵乘法 and BIOS loads MBR into 0x7C00?

    tianpeng <再谈矩阵与矩阵乘法> 讲的也好 矩阵乘矩阵 这个结果是怎么算出来的? 第一个矩阵第一行的每个数字(2和1),各自乘以第二个矩阵第一列对应位置的数字(1和1),然后将乘积 ...

  7. POJ 2778 DNA Sequence (AC自动机,矩阵乘法)

    题意:给定n个不能出现的模式串,给定一个长度m,要求长度为m的合法串有多少种. 思路:用AC自动机,利用AC自动机上的节点做矩阵乘法. #include<iostream> #includ ...

  8. 矩阵乘法的MPI并行计算

    1.问题描述 矩阵乘法问题描述如下: 给定矩阵A和B,其中A是m*p大小矩阵,B是p*n大小的矩阵.求C = A*B. 求解这个问题最简单的算法是遍历A的行和B的列,求得C的相应元素,时间复杂度O(m ...

  9. 有关CUBLAS中的矩阵乘法函数

    关于cuBLAS库中矩阵乘法相关的函数及其输入输出进行详细讨论. ▶ 涨姿势: ● cuBLAS中能用于运算矩阵乘法的函数有4个,分别是 cublasSgemm(单精度实数).cublasDgemm( ...

随机推荐

  1. hdu4791水题

    题意:       打印东西,给你区间和每个区间的价格,然后输入任务张数,输出最少花费..    题解:       昨晚的小测试就有这个题目,当时蒙B了,怎么也读不懂题目,一直纠结怎么把150拆成1 ...

  2. 洛谷P1089 津津的储蓄计划

    题目描述 津津的零花钱一直都是自己管理.每个月的月初妈妈给津津300元钱,津津会预算这个月的花销,并且总能做到实际花销和预算的相同. 为了让津津学习如何储蓄,妈妈提出,津津可以随时把整百的钱存在她那里 ...

  3. hdu4370 比较抽象的最短路

    题意:       给你一个n*n的矩阵,然后让咱们构造另一个n*n的矩阵,构造的矩阵有如下要求, 1.X12+X13+...X1n=1. 2.X1n+X2n+...Xn-1n=1. 3.for ea ...

  4. hdu5246超级赛亚ACMer

    题意(中文题意直接粘吧)                             超级赛亚ACMer Problem Description   百小度是一个ACMer,也是一个超级赛亚人,每个ACM ...

  5. 在进程空间使用虚拟内存(Windows 核心编程)

    虚拟内存空间 如今的 Windows 操作系统不仅可以运行多个应用程序,还可以让每一个应用程序享受到约 4 GB 的虚拟内存空间(包括系统占用),假如内存为 4 GB 的话.那为什么 Window 可 ...

  6. 解决在Vim中鼠标右键不能粘贴问题

    最近维护一台服务器,使用putty登录后,用vim时,鼠标右键不能 粘贴而是进入了visual模式.网上查找一番找到了解决方法: 方 法一:在普通模式下键入" :set mouse-=a&q ...

  7. iwrite复制攻略

    打开iwrite,一提交作业,发现: 这可咋办啊! 那就跟着步骤来呗: 按F12打开元素审查 点一下左上角 再点一下文本框,就能定位到HTML中的位置 在文本框中写几个字母,康康具体位置: 那就复制进 ...

  8. pandas(10):数据增删改

    目录 一.对索引进行操作 1 操作索引值df.rename() 二.指定数据替换.修改df.replace() 三.特殊值--缺失值处理 四.新增行列 1 直接赋值添加新列 2 df.assign() ...

  9. Blazor实现未登录重定向到登录页的方法

    今天研究了一下blazor,发现他默认启动就是类似于后台管理系统的界面,看到这个页面我就想给他写个登录,有登录就涉及到未登录重定向的问题,但是我没有找到blazor全局路由的设置,知道的老哥可以告诉我 ...

  10. 通过SQL注入获得网站后台用户密码

    通过 SQL 注入攻击,掌握网站的工作机制,认识到 SQL 注入攻击的防范措施,加强对 Web 攻击的防范. 一.实验环境 下载所需代码及软件:获取链接:链接:https://pan.baidu.co ...