网站建设 价格,网站域名用公司注册信息查询,网站注册转化率,培训的网站建设TVM#xff1a;使用Tensor Expression (TE)来处理算子
在本教程中#xff0c;我们将聚焦于在 TVM 中使用张量表达式#xff08;TE#xff09;来定义张量计算和实现循环优化。TE用纯函数语言描述张量计算#xff08;即每个表达式都没有副作用#xff09;。当在 TVM 的整体…TVM使用Tensor Expression (TE)来处理算子
在本教程中我们将聚焦于在 TVM 中使用张量表达式TE来定义张量计算和实现循环优化。TE用纯函数语言描述张量计算即每个表达式都没有副作用。当在 TVM 的整体上下文中查看时Relay 将计算描述为一组算子并且其中每一个算子都可以表示为 TE 表达式每个 TE 表达式获取输入张量并生成输出张量。 本文是TVM中 TE 语言的入门教程。TVM 使用领域专用domain specific的张量表达式来高效地构造内核。我们以两个使用 TE 语言的为例来演示基本工作流。第一个示例介绍了 TE 和带有向量加法的 schedule。第二个示例通过逐步优化矩阵与 TE 的乘法来扩展这些概念。这个矩阵乘法示例将作为未来涵盖更高级的 TVM 特性的教程的对比基础。
示例一使用TE为CPU编写和调度向量加法
初始化 tvm环境
我们的第一个例子是使用 Python 来为向量加法实现一个 TE然后是一个针对 CPU 的 schedule我们从初始化 tvm 环境开始
import tvm
import tvm.testing
from tvm import te
import numpy as np# 如果能够指定目标 CPU那么将会得到更好地性能
# 如果用的是llvm可以通过 llc --version 来查看 CPU 类型
# 可以通过查看 /proc/cpuinfo 来查看你的处理器可能支持的其他扩展
# 比如如果你的 CPU 有 AVX-512 指令集那么你可以使用 llvm -mcpuskylake-avx512 选项tgt tvm.target.Target(targetllvm, hostllvm)描述向量计算
我们首先描述向量加法计算。TVM 采用张量语义每个中间结果表示为一个多维数组。我们需要描述规则来得到张量。我们首先定义一个符号变量 n 来表示形状。然后我们定义两个 placeholder 张量A、B它们的形状是 (n,)。然后我们通过一个 compute 操作得到结果张量 C。compute 定义了一种计算其输出符合指定的张量形状并在由 lambda 函数定义的张量中的每个位置执行计算。注意虽然 n 是一个变量但它定义了A、B 和 C 张量之间的一致形状。请注意在这个阶段没有实际的计算发生因为我们只是声明应该如何进行计算。
n te.var(n)
A te.placeholder((n,), nameA)
B te.placeholder((n,), nameB)
C te.compute(A.shape, lambda i: A[i] B[i], nameC)注意Lambda函数 te.compute方法的第二个参数是执行计算的函数。在本例中我们使用一个匿名函数也称为lambda函数来定义计算在本例中是对 a 和 B 的第 i 个元素的加法。 为计算创建一个默认的Schedule
虽然上面几行描述了计算规则但我们可以用许多不同的方法计算 C 以适应不同的设备。对于具有多个 axis 的张量您可以选择首先迭代哪个 axis 另外计算可以跨不同的线程拆分。TVM要求用户提供一个 schedule来描述应如何执行计算。TE 中的 schedule 操作可以更改循环顺序、跨不同线程拆分计算、将数据块分组在一起以及其他操作。schedule 背后的一个重要概念是它们只描述如何执行计算因此相同 TE 的不同 schedule 一定会产生相同的结果。
在 TVM 中我们可以创建一种朴素的 schedule 按照行优先的顺序来计算 C。
for (int i 0; i n; i) {C[i] A[i] B[i];
}s te.create_schedule(C.op)编译并验证默认的 schedule
通过 TE 表达式和 schedule我们可以为目标语言和体系结构生成可运行的代码在本例中是 LLVM 和 CPU 。我们向 TVM 提供 schedule、schedule 中的TE表达式列表、目标和主机以及我们正在生成的函数的名称。输出的结果是可以直接从 Python 调用 type-erased 函数。
在下一行中我们使用 tvm.build 创建一个函数。build 函数接受 schedule、函数所需的签名包括输入和输出以及我们要编译到的目标语言。
fadd tvm.build(s, [A, B, C], tgt, namemyadd)我们运行该函数并将输出与 numpy 中的相同计算进行比较。编译后的 TVM 函数提供了一个简明的C API可以被任何语言调用。我们首先创建一个设备在本例中为CPU这是一个 TVM 可以编译 schedule 的设备。在本例中设备是LLVM CPU target。然后我们可以在设备中初始化张量并执行自定义的加法操作。为了验证计算的正确性我们可以将c张量的输出结果与 numpy 执行的相同计算进行比较。
dev tvm.device(tgt.kind.name, 0)n 1024
a tvm.nd.array(np.random.uniform(sizen).astype(A.dtype), dev)
b tvm.nd.array(np.random.uniform(sizen).astype(B.dtype), dev)
c tvm.nd.array(np.zeros(n, dtypeC.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())为了对比这个朴素版本的自定义向量加法与 numpy 的速度差异创建一个辅助函数来运行 TVM 生成代码的 profile。
import timeitnp_repeat 100
np_running_time timeit.timeit(setupimport numpy\nn 32768\ndtype float32\na numpy.random.rand(n, 1).astype(dtype)\nb numpy.random.rand(n, 1).astype(dtype)\n,stmtanswer a b,numbernp_repeat,
)
print(Numpy running time: %f % (np_running_time / np_repeat))def evaluate_addition(func, target, optimization, log):dev tvm.device(target.kind.name, 0)n 32768a tvm.nd.array(np.random.uniform(sizen).astype(A.dtype), dev)b tvm.nd.array(np.random.uniform(sizen).astype(B.dtype), dev)c tvm.nd.array(np.zeros(n, dtypeC.dtype), dev)evaluator func.time_evaluator(func.entry_name, dev, number10)mean_time evaluator(a, b, c).meanprint(%s: %f % (optimization, mean_time))log.append((optimization, mean_time))log [(numpy, np_running_time / np_repeat)]
evaluate_addition(fadd, tgt, naive, loglog)此处输出
Numpy running time: 0.000008
naive: 0.000006使用并行性paralleism来优化 schedule
我们已经说明了 TE 的基本原理现在让我们更深入地了解 schedule 的作用以及它们如何用于优化不同体系结构的张量表达式。schedule 是应用于表达式的一系列步骤用于以多种不同方式对其进行转换。当一个 schedule 应用于TE中的一个表达式时输入和输出保持不变但在编译时表达式的实现可能会改变。在默认 schedule 中这个张量加法是串行运行的但该操作其实是很容易在所有处理器线程之间并行。我们可以将我们的操作并行调度到计算中
s[C].parallel(C.op.axis[0])tvm.lower 命令将生成 TE 的中间表示IR以及相应的 schedule 。通过在执行不同的 schedule 操作时 lowing 表达式我们可以看到 schedule 对计算顺序的影响。我们使用标志 simple_modeTrue 返回可读的 C 风格语句。
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], typeauto),A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], typeauto),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], typeauto)}buffer_map {A_1: A, B_1: B, C_1: C} {for (i: int32, 0, n) parallel {C_2[(i*stride)] ((float32*)A_2[(i*stride_1)] (float32*)B_2[(i*stride_2)])}
}TVM现在可以在独立的线程上运行这些块。我们在执行并行操作的情况下编译并运行这个新的 schedule
fadd_parallel tvm.build(s, [A, B, C], tgt, namemyadd_parallel)
fadd_parallel(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())evaluate_addition(fadd_parallel, tgt, parallel, loglog)此处输出
parallel: 0.000005使用矢量化vectorization来优化 schedule
现代 CPU 能够对浮点数进行 SIMD 操作我们可以对计算表达式使用另一个 schedule 来利用这一点。实现这一点需要多个步骤首先我们必须使用 split scheduling 原语将 schedule 拆分为内部循环和外部循环。内部循环可以使用向量化来使用使用向量化调度原语的 SIMD 指令然后外部循环可以使用并行调度原语进行并行化。选择分割因子作为CPU上的线程数。 注SIMD全称 Single Instruction Multiple Data单指令多数据流能够复制多个操作数并把它们打包在大型寄存器的一组指令集。 # 由于我们需要修改之前例子中的并行操作因此这里要重建 schedule
n te.var(n)
A te.placeholder((n,), nameA)
B te.placeholder((n,), nameB)
C te.compute(A.shape, lambda i: A[i] B[i], nameC)s te.create_schedule(C.op)# factor 的选择需要适合你的线程数这取决于架构
# 建议将此系数设置为等于可用CPU核心数。
factor 4outer, inner s[C].split(C.op.axis[0], factorfactor)
s[C].parallel(outer)
s[C].vectorize(inner)fadd_vector tvm.build(s, [A, B, C], tgt, namemyadd_parallel)evaluate_addition(fadd_vector, tgt, vector, loglog)print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
vector: 0.000016
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {A: Buffer(A_2: Pointer(float32), float32, [n: int32], [stride: int32], typeauto),C: Buffer(C_2: Pointer(float32), float32, [n], [stride_1: int32], typeauto),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], typeauto)}buffer_map {A_1: A, B_1: B, C_1: C} {for (i.outer: int32, 0, floordiv((n 3), 4)) parallel {for (i.inner.s: int32, 0, 4) {if tir.likely((((i.outer*4) i.inner.s) n), dtypebool) {C_2[(((i.outer*4) i.inner.s)*stride_1)] ((float32*)A_2[(((i.outer*4) i.inner.s)*stride)] (float32*)B_2[(((i.outer*4) i.inner.s)*stride_2)])}}}
}对比不同的 schedule
下面我们来对比以下之前提到的不同 schedule
baseline log[0][1]
print(%s\t%s\t%s % (Operator.rjust(20), Timing.rjust(20), Performance.rjust(20)))
for result in log:print(%s\t%s\t%s% (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20)))此处输出
Operator Timing Performancenumpy 7.98278022557497e-06 1.0naive 5.9189e-06 0.7414584684465222
parallel 4.9771999999999995e-06 0.6234920490550659vector 1.6127399999999997e-05 2.0202735819196875注意Code Specialization 代码专门化 正如我们所看到的A、B 和 C 的声明都采用相同的形状参数 n。TVM将利用这一点只向 kernel 传递一个 shape 参数我们在打印的设备代码中找到它。这是专门化化的一种形式。 在 host 端TVM 将自动生成检查代码以检查参数中的约束。因此如果将具有不同形状的数组传递到 fadd 中将引发错误。 我们可以做更多的专门化。例如我们可以在计算声明中写入ntvm.runtime.convert1024而不是 nte.var“n”。生成的函数将只获取长度为1024的向量。 我们已经定义、调度并编译了一个向量加法运算符然后可以在 TVM Runtime 执行它。我们可以将算子保存为库稍后可以使用 TVM Runtime 加载该库。
针对GPU的矩阵加法可选
在介绍保存与加载自定义算子库的方法之前我们先来看一下如何针对 GPU 做矩阵加法。
TVM能够针对多种体系结构。在本例我们将针对GPU中矢量加法的编译。
# 本段代码默认不运行如果想要运行的话请将 run_cuda Truerun_cuda False
if run_cuda:# 这里的 target 需要根据自己的 GPU 类型修改# NVIDIAcuda# Radeonrocm# OpenCLopencltgt_gpu tvm.target.Target(targetcuda, hostllvm)# 重建 schedulen te.var(n)A te.placeholder((n,), nameA)B te.placeholder((n,), nameB)C te.compute(A.shape, lambda i: A[i] B[i], nameC)print(type(C))s te.create_schedule(C.op)bx, tx s[C].split(C.op.axis[0], factor64)################################################################################# 最后我们必须将迭代轴bx和tx绑定到GPU计算网格中的线程。# 朴素的 schedule 对GPU无效这些是允许我们生成在GPU上运行的代码的特定构造。s[C].bind(bx, te.thread_axis(blockIdx.x))s[C].bind(tx, te.thread_axis(threadIdx.x))####################################################################### 编译# -----------# 在指定完 schdule 之后我们可以将其编译成一个 TVM 函数。默认情况下TVM编译成一个 type-erased 函 # 数可以从python端直接调用该函数。# 在下一行中我们使用 tvm.build 来创建一个函数。build 函数采用 schedule、函数所需的签名包括输如和输出以及我们要编译到的目标语言。# 编译 fadd 的结果是一个GPU设备函数如果涉及GPU以及一个调用 GPU 函数的 host wrapper。fadd是生成的主机包装函数它在内部包含对生成的设备函数的引用。fadd tvm.build(s, [A, B, C], targettgt_gpu, namemyadd)################################################################################# 编译过的 TVM 函数会有一个简洁的 C API它可以被任意的语言调用## 我们提供一个 Python 的最小的数组 API 来帮助快速的测试和原型化# 该数组 API 是基于 DLPack https://github.com/dmlc/dlpack_ 标准.## - 我们首先创建一个 GPU 设备# - 然后 tvm.nd.array 从 GPU 拷贝数据# - fadd 运行真正的计算# - numpy() 从 GPU 数组拷贝回 CPU (这样我们就能验证正确性了).## 请注意将数据复制到 GPU 上的内存和从中复制数据是必需的步骤。dev tvm.device(tgt_gpu.kind.name, 0)n 1024a tvm.nd.array(np.random.uniform(sizen).astype(A.dtype), dev)b tvm.nd.array(np.random.uniform(sizen).astype(B.dtype), dev)c tvm.nd.array(np.zeros(n, dtypeC.dtype), dev)fadd(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())################################################################################# 检查生成的 GPU 代码# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~# 我们可以检查在 TVM 中生成的代码tvm.build 的结果是一个 TVM 模块。fadd 是一个 host 模块其中包含 # host wrapper 的 host module它同样包含一个CUDAGPU设备模块## 下面的代码取得设备模块并打印内容代码if (tgt_gpu.kind.name cudaor tgt_gpu.kind.name rocmor tgt_gpu.kind.name.startswith(opencl)):dev_module fadd.imported_modules[0]print(-----GPU code-----)print(dev_module.get_source())else:print(fadd.get_source())保存和加载编译过的模块
保存编译过的模块
除了运行时编译之外我们还可以将编译后的模块保存到一个文件中并在以后重新加载。下面的代码执行以下步骤
它将编译后的主机模块保存到一个对象文件中。然后将设备模块保存到 ptx 文件中。cc.create_shared 调用编译器gcc来创建共享库
from tvm.contrib import cc
from tvm.contrib import utilstemp utils.tempdir()
fadd.save(temp.relpath(myadd.o))
if tgt.kind.name cuda:fadd.imported_modules[0].save(temp.relpath(myadd.ptx))
if tgt.kind.name rocm:fadd.imported_modules[0].save(temp.relpath(myadd.hsaco))
if tgt.kind.name.startswith(opencl):fadd.imported_modules[0].save(temp.relpath(myadd.cl))
cc.create_shared(temp.relpath(myadd.so), [temp.relpath(myadd.o)])
print(temp.listdir())此处输出
[myadd.o, myadd.so]注意Module Storage Format 模块存储格式 CPUHost模块直接保存为共享库.so。设备代码可以有多种自定义格式。在我们的示例中设备代码存储在 ptx 中元数据在 json 文件中。它们可以通过导入单独加载和链接。 加载编译过的模块
我们可以从文件系统加载已编译的模块并运行代码。以下代码分别加载主机和设备模块并将它们链接在一起。我们可以验证新加载的函数是否有效。
fadd1 tvm.runtime.load_module(temp.relpath(myadd.so))
if tgt.kind.name cuda:fadd1_dev tvm.runtime.load_module(temp.relpath(myadd.ptx))fadd1.import_module(fadd1_dev)if tgt.kind.name rocm:fadd1_dev tvm.runtime.load_module(temp.relpath(myadd.hsaco))fadd1.import_module(fadd1_dev)if tgt.kind.name.startswith(opencl):fadd1_dev tvm.runtime.load_module(temp.relpath(myadd.cl))fadd1.import_module(fadd1_dev)fadd1(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())将所有东西打包在一个库中
在上面的示例中我们分别存储设备和主机代码。TVM 还支持将所有内容导出为一个共享库。在 hood 下我们将设备模块打包成二进制blob并将它们与主机代码链接在一起。目前我们支持Metal、OpenCL和CUDA模块的包装。
fadd.export_library(temp.relpath(myadd_pack.so))
fadd2 tvm.runtime.load_module(temp.relpath(myadd_pack.so))
fadd2(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())注意Runtime API and Thread Safety 运行时API与线程安全 TVM 的编译模块并不依赖于 TVM 编译器。它们只依赖于最小 Runtime Library。TVM Runtime Library 包装设备驱动程序并向编译函数提供线程安全和设备无关调用。 这意味着我们可以从任何GPU上的任何线程调用已编译的TVM函数前提是您已经为该GPU编译了代码。 生成OpenCL代码
TVM 为多种后端提供代码生成功能。我们还可以生成在 CPU 后端上运行的 OpenCL 代码或 LLVM 代码。
下面的代码可以生成OpenCL代码在OpenCL设备上创建数组并验证代码的正确性。
if tgt.kind.name.startswith(opencl):fadd_cl tvm.build(s, [A, B, C], tgt, namemyadd)print(------opencl code------)print(fadd_cl.imported_modules[0].get_source())dev tvm.cl(0)n 1024a tvm.nd.array(np.random.uniform(sizen).astype(A.dtype), dev)b tvm.nd.array(np.random.uniform(sizen).astype(B.dtype), dev)c tvm.nd.array(np.zeros(n, dtypeC.dtype), dev)fadd_cl(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() b.numpy())注意TE Scheduling Primitives TE 调度原语 TVM 包括许多不同的调度原语 split按定义的因子将指定轴拆分为两个轴。 tile平铺将按定义的因子沿两个轴分割计算。 fuse融合一次计算的两个连续轴。 reorder可以将计算轴重新排序为定义的顺序。 bind可以将计算绑定到特定线程在GPU编程中很有用。 compute_at默认情况下TVM将在函数的最外层或根计算张量。compute_at指定应在另一个运算符的第一个计算轴上计算一个张量。 compute_inline当标记为inline时计算将展开然后插入到需要张量的地址中。 compute_root将计算移动到函数的最外层或根。这意味着在进入下一个阶段之前将对计算阶段进行完全计算。 可以在Schedule primitives 文档页面中找到这些原语的完整描述。 示例二用TE手动优化矩阵乘
现在我们将考虑第二个更高级一些的示例演示如何用 18 行 Python 代码 TVM 加速一个共同的矩阵乘法运算 18倍。
矩阵乘法是一种计算密集型运算。要获得良好的CPU性能有两个重要的优化
提高内存访问的缓存命中率。高缓存命中率可以加速复杂的数值计算和热点内存访问。这要求我们将源内存访问模式转换为适合缓存策略的模式。SIMD单指令多数据也称为矢量处理单元。在每个循环中SIMD 都可以处理一小批数据而不是处理单个值。这要求我们以统一模式转换循环体中的数据访问模式以便LLVM 后端可以将其 lower 到 SIMD。
本教程中使用的技术是这个仓库中提到的技巧的一部分。其中一些已被 TVM 抽象自动使用但由于 TVM 的一些约束有一些无法自动使用。
准备工作和性能baseline
我们首先采集 numpy 实现的矩阵乘的数据
import tvm
import tvm.testing
from tvm import te
import numpy# 矩阵的尺寸
# (M, K) x (K, N)
# 你可以自己试一些不同的尺寸有时候 TVM 的优化结果会好于含 MKL 的numpy
M 1024
K 1024
N 1024# tvm 中默认的数据类型
dtype float32# 与之前一样这里可以根据自己的处理器及其是否支持某些指令集来改变 targettarget tvm.target.Target(targetllvm, hostllvm)
dev tvm.device(target.kind.name, 0)# 随机生成一些 tensor 用于测试
a tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
b tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)# 重复实验得到 numpy 的矩阵乘实现的 baseline
np_repeat 100
np_running_time timeit.timeit(setupimport numpy\nM str(M) \nK str(K) \nN str(N) \ndtype float32\na numpy.random.rand(M, K).astype(dtype)\nb numpy.random.rand(K, N).astype(dtype)\n,stmtanswer numpy.dot(a, b),numbernp_repeat,
)
print(Numpy running time: %f % (np_running_time / np_repeat))answer numpy.dot(a.numpy(), b.numpy())此处输出
Numpy running time: 0.009308现在我们用 TVM TE 编写一个基本矩阵乘法并验证它产生的结果与numpy实现相同。我们还编写了一个函数它将帮助我们度量进度优化的性能。
# 使用 TE 实现的 TVM 的矩阵乘
k te.reduce_axis((0, K), k)
A te.placeholder((M, K), nameA)
B te.placeholder((K, N), nameB)
C te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axisk), nameC)# 默认 schedule
s te.create_schedule(C.op)
func tvm.build(s, [A, B, C], targettarget, namemmult)c tvm.nd.array(numpy.zeros((M, N), dtypedtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol1e-5)def evaluate_operation(s, vars, target, name, optimization, log):func tvm.build(s, [A, B, C], targettarget, namemmult)assert funcc tvm.nd.array(numpy.zeros((M, N), dtypedtype), dev)func(a, b, c)tvm.testing.assert_allclose(c.numpy(), answer, rtol1e-5)evaluator func.time_evaluator(func.entry_name, dev, number10)mean_time evaluator(a, b, c).meanprint(%s: %f % (optimization, mean_time))log.append((optimization, mean_time))log []evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationnone, loglog)此处输出
none: 3.109406让我们看一下使用 TVM lower 函数的算子和默认 schedule 的中间表示 IR。请注意该实现本质上是矩阵乘法的简单实现在 A 和 B 矩阵的索引上使用三个嵌套循环。
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {for (x: int32, 0, 1024) {for (y: int32, 0, 1024) {C_2[((x*1024) y)] 0f32for (k: int32, 0, 1024) {C_2[((x*1024) y)] ((float32*)C_2[((x*1024) y)] ((float32*)A_2[((x*1024) k)]*(float32*)B_2[((k*1024) y)]))}}}
}优化1blocking阻塞
提高缓存命中率的一个重要技巧是阻塞在这种情况下我们可以构造内存访问使块内部是具有高内存局部性的小邻域。在本教程中我们选择块因子 32。这会使得一个块填充内存的 32*32*sizeoffloat区域。这对应于 4KB 的缓存大小和一级缓存 32KB 的参考缓存大小。
我们首先为 C 操作创建一个默认的调度然后使用指定的块因子对其应用一个 tile 调度原语调度原语以向量 [x_-outery_-outerx_-innery_-inner] 的形式返回从最外层到最内层的循环顺序。然后我们得到操作输出的缩减轴并使用因子4对其执行拆分操作。这个因素不会直接影响我们现在正在进行的阻塞优化但在以后应用矢量化时会很有用。 现在操作已被阻塞我们可以对计算进行重新排序将简化操作放入计算的最外层循环中从而帮助确保被阻塞的数据仍保留在缓存中。这就完成了 schedule我们可以构建和测试与原始 schedule 相比的性能。
bn 32# Blocking by loop tiling
xo, yo, xi, yi s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) s[C].op.reduce_axis
ko, ki s[C].split(k, factor4)# Hoist reduction domain outside the blocking loop
s[C].reorder(xo, yo, ko, ki, xi, yi)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationblocking, loglog)此处输出
blocking: 0.291928通过重新排序计算以利用缓存我们可以看到计算性能的显著提高。现在打印内部表示并将其与原始表示进行比较
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {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} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {for (y.inner.init: int32, 0, 32) {C_2[((((x.outer*32768) (x.inner.init*1024)) (y.outer*32)) y.inner.init)] 0f32}}for (k.outer: int32, 0, 256) {for (k.inner: int32, 0, 4) {for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] ((float32*)C_2[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] ((float32*)A_2[((((x.outer*32768) (x.inner*1024)) (k.outer*4)) k.inner)]*(float32*)B_2[((((k.outer*4096) (k.inner*1024)) (y.outer*32)) y.inner)]))}}}}}}
}优化2 vectorization矢量化
另一个重要的优化技巧是矢量化。当内存访问模式一致时编译器可以检测到该模式并将连续内存传递给 SIMD 向量处理器。在TVM中我们可以利用这个硬件特性使用矢量化接口来提示编译器这个模式。
在本教程中我们选择对内部循环行数据进行矢量化因为它已经是我们之前优化中的缓存友好型数据。
# 应用矢量化的优化方式
s[C].vectorize(yi)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationvectorization, loglog)# 矢量化之后生成的 IR
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
vectorization: 0.331263
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C_2[ramp((((x.outer*32768) (x.inner.init*1024)) (y.outer*32)), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (k.inner: int32, 0, 4) {for (x.inner: int32, 0, 32) {C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] ((float32x32*)C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.inner*1024)) (k.outer*4)) k.inner)], 32)*(float32x32*)B_2[ramp((((k.outer*4096) (k.inner*1024)) (y.outer*32)), 1, 32)]))}}}}}
}优化3Loop Permutation循环置换
如果我们看一下上面的 IR我们可以看到内环行数据被矢量化B 被转换成 PackedB这在内环的float32x32B2部分中很明显。PackedB 的遍历现在是顺序的。因此我们将研究 A 的访问模式。在当前 schdule中A 是逐列访问的这对缓存不友好。如果我们改变嵌套循环顺序 ki 和内部轴 xi对 A 的访问模式将变得更加缓存友好。
s te.create_schedule(C.op)
xo, yo, xi, yi s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) s[C].op.reduce_axis
ko, ki s[C].split(k, factor4)# re-ordering
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationloop permutation, loglog
)# 再一次打印新生成的 IR
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
loop permutation: 0.113750
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C_2[ramp((((x.outer*32768) (x.inner.init*1024)) (y.outer*32)), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.inner: int32, 0, 32) {for (k.inner: int32, 0, 4) {C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] ((float32x32*)C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.inner*1024)) (k.outer*4)) k.inner)], 32)*(float32x32*)B_2[ramp((((k.outer*4096) (k.inner*1024)) (y.outer*32)), 1, 32)]))}}}}}
}优化4Array Packing数组打包
另一个重要技巧是数组打包。此技巧是对阵列的存储维度重新排序以便在展平后将特定维度上的连续访问模式转换为序列模式。 如上图所示在阻塞计算后我们可以观察到 B 的阵列访问模式平坦后它是规则的但不连续的。我们希望经过一些转换后我们可以得到一个连续的访问模式。通过将[16][16]数组重新排序为[16/4][16][4]数组在从压缩数组中获取相应值时B 的访问模式将是顺序的。
为了实现这一点我们必须从一个新的默认 schedule 开始考虑到 B 的新 wrapper。花点时间对此进行讨论是值得的TE 是一种用于编写优化算子的功能强大的表达性语言但它通常需要一些底层算法、数据结构以及您正在编写的硬件 target。在本教程的后面我们将讨论让 TVM 承担这一负担的一些选择。不管怎样让我们继续新的优化 schedule。
# 我们要轻微地重写算法
packedB te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn z], namepackedB)
C te.compute((M, N),lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axisk),nameC,
)s te.create_schedule(C.op)xo, yo, xi, yi s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) s[C].op.reduce_axis
ko, ki s[C].split(k, factor4)s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)x, y, z s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationarray packing, loglog)# 这里是数组打包之后生成的 IR
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
array packing: 0.224114
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope global {for (x: int32, 0, 32) parallel {for (y: int32, 0, 1024) {packedB[ramp(((x*32768) (y*32)), 1, 32)] (float32x32*)B_2[ramp(((y*1024) (x*32)), 1, 32)]}}for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C_2[ramp((((x.outer*32768) (x.inner.init*1024)) (y.outer*32)), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.inner: int32, 0, 32) {for (k.inner: int32, 0, 4) {C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] ((float32x32*)C_2[ramp((((x.outer*32768) (x.inner*1024)) (y.outer*32)), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.inner*1024)) (k.outer*4)) k.inner)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) (k.inner*32)), 1, 32)]))}}}}}}
}优化5Optimizing Block Writing Through Caching通过缓存优化块写入
到目前为止我们所有的优化都集中在高效地访问和计算来自 A 和 B 矩阵的数据以计算C矩阵。阻塞优化后操作员将结果逐块写入 C并且访问模式不是顺序的。我们可以通过使用顺序缓存数组来解决这个问题使用cache_write、compute_at 和 unroll 的组合来保存块结果并在所有块结果就绪时写入到 C。
s te.create_schedule(C.op)# Allocate write cache
CC s.cache_write(C, global)xo, yo, xi, yi s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)# Write cache is computed at yo
s[CC].compute_at(s[C], yo)# New inner axes
xc, yc s[CC].op.axis(k,) s[CC].op.reduce_axis
ko, ki s[CC].split(k, factor4)
s[CC].reorder(ko, xc, ki, yc)
s[CC].unroll(ki)
s[CC].vectorize(yc)x, y, z s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationblock caching, loglog)# Here is the generated IR after write cache blocking.
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
block caching: 0.224213
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope global;allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope global {for (x: int32, 0, 32) parallel {for (y: int32, 0, 1024) {packedB[ramp(((x*32768) (y*32)), 1, 32)] (float32x32*)B_2[ramp(((y*1024) (x*32)), 1, 32)]}}for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.c.init: int32, 0, 32) {C.global[ramp((x.c.init*32), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[(((x.outer*32768) (x.c*1024)) (k.outer*4))], 32)*(float32x32*)packedB[ramp(((y.outer*32768) (k.outer*128)), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 1)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 32), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 2)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 64), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 3)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 96), 1, 32)]))}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] (float32*)C.global[((x.inner*32) y.inner)]}}}}}
}优化6Parallelization并行化
# 并行
s[C].parallel(xo)x, y, z s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationparallelization, loglog
)# 这里是并行化之后的 IR
print(tvm.lower(s, [A, B, C], simple_modeTrue))此处输出
parallelization: 0.067949
primfn(A_1: handle, B_1: handle, C_1: handle) - ()attr {from_legacy_te_schedule: True, global_symbol: main, tir.noalias: True}buffers {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),A: Buffer(A_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} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope global {for (x: int32, 0, 32) parallel {for (y: int32, 0, 1024) {packedB[ramp(((x*32768) (y*32)), 1, 32)] (float32x32*)B_2[ramp(((y*1024) (x*32)), 1, 32)]}}for (x.outer: int32, 0, 32) parallel {allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope global;for (y.outer: int32, 0, 32) {for (x.c.init: int32, 0, 32) {C.global[ramp((x.c.init*32), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[(((x.outer*32768) (x.c*1024)) (k.outer*4))], 32)*(float32x32*)packedB[ramp(((y.outer*32768) (k.outer*128)), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 1)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 32), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 2)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 64), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] ((float32x32*)C.global[ramp((x.c*32), 1, 32)] (broadcast((float32*)A_2[((((x.outer*32768) (x.c*1024)) (k.outer*4)) 3)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) (k.outer*128)) 96), 1, 32)]))}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] (float32*)C.global[((x.inner*32) y.inner)]}}}}}
}矩阵乘例子的总结
在仅用 18 行代码应用上述简单优化之后我们生成的代码就可以得到与使用数学内核库MKL的 numpy 接近的性能。我们刚才一直都记录了性能因此在这里可以直接比较结果
baseline log[0][1]
print(%s\t%s\t%s % (Operator.rjust(20), Timing.rjust(20), Performance.rjust(20)))
for result in log:print(%s\t%s\t%s% (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20)))此处输出 Operator Timing Performancenone 3.1094061458 1.0blocking 0.29192816779999997 0.09388550549895809vectorization 0.3312631714 0.10653583220302389
loop permutation 0.1137497149 0.036582456445468314array packing 0.2241142794 0.07207623221003798block caching 0.22421289339999997 0.07210794694763607parallelization 0.0679485881 0.021852593361526892请注意以上的输出反映的是非独占 Docker 容器上的运行时间因此并不可靠。强烈建议您自己运行本教程观察 TVM 实现的性能增益并仔细阅读每个示例以了解矩阵乘法运算的迭代改进。
总结
如前所述如何使用 TE 和调度原语应用优化可能需要一些底层架构和算法的知识。然而TE 设计为更复杂的算法是为了可以搜索潜在的优化。有了本 TE 简介中的知识我们现在可以开始探索 TVM 如何自动化进度优化过程。
本教程提供了使用向量加法和矩阵乘法示例的TVM张量表达式TE工作流演练。一般的工作流程是 通过一系列操作描述您的计算。 描述我们希望如何计算和使用调度原语。 编译到我们想要的目标函数。 保存要稍后加载的函数可选。
接下来的教程将扩展矩阵乘法示例并展示如何使用可调参数构建矩阵乘法和其他操作的通用模板这些参数使得我们能够自动优化特定平台的计算。
Ref
https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html
