网站域名可以更换吗,市场调研流程,html代码格式,想做个网络推广引言
注#xff1a;本文主要介绍、解释TVM的矩阵优化思想、代码#xff0c;需要配合代码注释一起阅读。
矩阵乘法是计算密集型运算。为了获得良好的 CPU 性能#xff0c;有两个重要的优化措施#xff1a; 提高内存访问的高速缓存命中率。复杂的数值计算和热点内存#x…引言
注本文主要介绍、解释TVM的矩阵优化思想、代码需要配合代码注释一起阅读。
矩阵乘法是计算密集型运算。为了获得良好的 CPU 性能有两个重要的优化措施 提高内存访问的高速缓存命中率。复杂的数值计算和热点内存hot-spot memory访问都可以通过高缓存命中率high cache hit rate来加速。这就要求我们将原内存origin 访问模式转化为符合高速缓存策略的模式。 SIMD单指令多数据也被称为矢量处理单元。在每个周期中SIMD 可以处理一小批数据而不是处理一个单一的值。这就要求我们将循环体中的数据访问模式转化为统一模式以便编译器后端可以将其降低到 SIMD。
1- 基于TVM和Numpy的基线实现
这是numpy 实现矩阵乘法的实现代码作为实现比较的基础benchmark他的速度并不会很慢这是因为其本身就采用了加速矩阵数学计算的算法实现。
import tvm
import tvm.testing
from tvm import te
import numpy# The size of the matrix
# (M, K) x (K, N)
# You are free to try out different shapes, sometimes TVM optimization outperforms numpy with MKL.
M 1024
K 1024
N 1024# The default tensor data type in tvm
dtype float32# You will want to adjust the target to match any CPU vector extensions you
# might have. For example, if youre using using Intel AVX2 (Advanced Vector
# Extensions) ISA for SIMD, you can get the best performance by changing the
# following line to llvm -mcpucore-avx2, or specific type of CPU you use.
# Recall that youre using llvm, you can get this information from the command
# llc --version to get the CPU type, and you can check /proc/cpuinfo
# for additional extensions that your processor might support.target tvm.target.Target(targetllvm, hostllvm)
dev tvm.device(target.kind.name, 0)# Random generated tensor for testing
a tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
b tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)# Repeatedly perform a matrix multiplication to get a performance baseline
# for the default numpy implementation
np_repeat 100
np_running_time timeit.timeit( // 统计numpy.dot(a, b)计算耗时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())现在用 TVM TE 编写基本的矩阵乘法并验证它产生的结果与 numpy 的实现相同。我们还写了一个函数它将帮助衡量调度优化的性能。
# TVM Matrix Multiplication using TE
k te.reduce_axis((0, K), k) // 可以理解为约分轴M*K * K*N M*NK轴维度消失啦
A te.placeholder((M, K), nameA) // 一种tensor通常用于计算图的Input节点使用没有前序节点
B te.placeholder((K, N), nameB)
C te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axisk), nameC) // 定义输出结果。第一个参数(M, N)表示输出矩阵的shapelambda i:则可以理解为 for i: 0-n-1# Default schedule
s te.create_schedule(C.op) // 从一个或多个 Tensor 的 Operation 进行创建
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)让我们来看看上述使用 TVM 低级函数的运算器和默认调度的中间表示。请注意这个实现基本上是矩阵乘法的原始实现在 A 和 B 矩阵的索引上使用三个嵌套循环。
先来看看原始的矩阵乘法实现他基本没有考虑cache连续的问题
for(int x 0; x M; x){for(int y 0; y N; y){for(int k 0; k K; k){C[x][y] A[x][k] * B[k][y];}}
}TVM的计算其实是将二维矩阵平铺为了一维数组进行计算由于边界定义为了1024这里的offset也是1024
main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {for (x: int32, 0, 1024) { // 主要看这里对应关系为 x-M y-N k-Kfor (y: int32, 0, 1024) {C[((x*1024) y)] 0f32for (k: int32, 0, 1024) {let cse_var_2: int32 (x*1024)let cse_var_1: int32 (cse_var_2 y)C[cse_var_1] (C[cse_var_1] (A[(cse_var_2 k)]*B[((k*1024) y)]))}}}
}2-使用TVM调度原语tile分块优化cache阻塞
提高缓冲区命中率的一个重要技巧是阻塞在这个过程中你的内存访问结构是在一个块的内部有一个小的邻域具有很高的内存定位性。在本教程中我们选择一个 32 的块因子。这将导致一个块充满 32 * 32 * sizeof(float) 的内存区域。这相当于一个 4KB 的缓存大小而 L1 缓存的参考缓存大小为 32KB。
对矩阵计算进行分块以加速。分块可以进一步的提升B、C矩阵的空间局部性 我们把C分成多个title然后针对每一个titleA和B中对应的行tile和列tile也会切成相应大小的tile进行多个小矩阵乘法最后加和到C的tile中当我们把tile的大小限定到合适的范围内时就可以把整个tile填充到cache内分块的好处就体现在一个block内的计算小到可以被cache容纳以充分利用cache快速的读写速度。
我们首先为 C 矩阵计算操作创建一个默认的调度然后用指定的块因子对其应用一个 tile 调度原语调度原语返回所产生的循环顺序从最外层到最内层作为一个向量 [x_outer, y_outer, x_inner, y_inner]。然后我们得到操作输出的还原轴并使用 4 的因子对其进行分割操作。这个因子并不直接影响我们现在正在进行的阻塞优化但在以后我们应用矢量化时将会很有用。
现在操作已经被阻塞了我们可以重新调度计算的顺序把较少操作放到计算的最外层循环中帮助保证被阻止的数据仍然在缓存中。这样就完成了调度我们可以建立并测试与原生的调度相比的性能。
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)通过重新安排计算顺序以利用缓存你应该看到计算的性能有了明显的改善。现在打印内部表示并将其与原始表示进行比较。
print(tvm.lower(s, [A, B, C], simple_modeTrue))main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {for (x.outer: int32, 0, 32) { // 外部32个分块for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) { // 内部32*32个浮点数for (y.inner.init: int32, 0, 32) {C[((((x.outer*32768) (x.inner.init*1024)) (y.outer*32)) y.inner.init)] 0f32 // 32768为32*32*32初始化该块内的数为0}}for (k.outer: int32, 0, 256) { // k轴分为4份1024/4256for (k.inner: int32, 0, 4) { // 这里的4主要用于后面的SIMD的使用for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {let cse_var_3: int32 (y.outer*32)let cse_var_2: int32 ((x.outer*32768) (x.inner*1024))let cse_var_1: int32 ((cse_var_2 cse_var_3) y.inner)C[cse_var_1] (C[cse_var_1] (A[((cse_var_2 (k.outer*4)) k.inner)]*B[((((k.outer*4096) (k.inner*1024)) cse_var_3) y.inner)])) // 其实这里B依然是cache不友好读取即跳行读取但可能是考虑到在一个定义的分块cache内}}}}}}
}3-利用SIMD矢量化计算
另一个重要的优化技巧是矢量化。当内存访问模式是统一的编译器可以检测到这种模式并将连续的内存传递给 SIMD 矢量处理器。在 TVM 中我们可以使用 vectorize 接口来提示编译器这种模式利用这一硬件特性。
SIMD 的全称是 Single Instruction Multiple Data中文名“单指令多数据”。顾名思义一条指令处理多个数据。
在本教程中我们选择对内循环的行数据进行矢量化因为在我们之前的优化中它已经是缓存友好的。
# Apply the vectorization optimization
s[C].vectorize(yi) // 对子块内的每一行进行SIMD计算evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationvectorization, loglog)# The generalized IR after vectorization
print(tvm.lower(s, [A, B, C], simple_modeTrue))vectorization: 0.257297
main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C[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) { // 对于y.inner的每一行使用128位SIMD计算32*4一次计算4位let cse_var_3: int32 (y.outer*32)let cse_var_2: int32 ((x.outer*32768) (x.inner*1024))let cse_var_1: int32 (cse_var_2 cse_var_3)C[ramp(cse_var_1, 1, 32)] (C[ramp(cse_var_1, 1, 32)] (broadcast(A[((cse_var_2 (k.outer*4)) k.inner)], 32)*B[ramp((((k.outer*4096) (k.inner*1024)) cse_var_3), 1, 32)]))}}}}}
}4-使用循环优化
如果我们看一下上面的 IR我们可以看到内循环的行数据被矢量化B 被转化为 分块B这从内循环的 (float32x32*)B2 部分可以看出。现在 B的分块 的遍历是顺序的。所以我们要看一下 A 的访问模式。在当前的计划中A 是被逐列访问的每次访问A的128位这对缓冲区不友好。如果我们改变 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
)# Again, print the new generalized IR
print(tvm.lower(s, [A, B, C], simple_modeTrue))loop permutation: 0.130006
main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C[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) { // 调换了这里的顺序let cse_var_3: int32 (y.outer*32)let cse_var_2: int32 ((x.outer*32768) (x.inner*1024))let cse_var_1: int32 (cse_var_2 cse_var_3)C[ramp(cse_var_1, 1, 32)] (C[ramp(cse_var_1, 1, 32)] (broadcast(A[((cse_var_2 (k.outer*4)) k.inner)], 32)*B[ramp((((k.outer*4096) (k.inner*1024)) cse_var_3), 1, 32)]))}}}}}
}5-使用数组打包
另一个重要的技巧是数组打包。这个技巧是对数组的存储维度进行重新排序将某些维度上的连续访问模式转换为扁平化后的顺序模式。 正如上图所示在阻塞计算后我们可以观察到 B 的数组访问模式扁平化后它是有规律的但是不连续的。我们期望经过一些转换后我们可以得到一个连续的访问模式。通过将 [16][16] 数组重新排序为 [16/4][16][4] 数组当从打包的数组中抓取相应的值时B 的访问模式将是连续的。
为了实现这一目标将不得不从新的默认调度开始考虑到 B 的新包装值得花点时间来评论一下。TE 是编写优化运算符的强大而富有表现力的语言但它往往需要对你所编写的底层算法、数据结构和硬件目标有一些了解。在本教程的后面我们将讨论一些让 TVM 承担这一负担的选项。无论如何让我们继续编写新的优化调度。
这里的本质相当于改变了B的按列读取改为按行读取即cache友好。
// 必须重写算法。对计算方式进行重新定义。
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), // TIR 的全称是 Tensor IR, tvm.tir.indexmod 为计算y/bn的余数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) // 执行向量操作的指令。vectorize 是并发思路在CPU里的最细粒度实现。
s[packedB].parallel(x)evaluate_operation(s, [A, B, C], targettarget, namemmult, optimizationarray packing, loglog)# Here is the generated IR after array packing.
print(tvm.lower(s, [A, B, C], simple_modeTrue))array packing: 0.140257
main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope global {for (x: int32, 0, 32) parallel {for (y: int32, 0, 1024) {packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) y)] B[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[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) {let cse_var_3: int32 ((x.outer*32768) (x.inner*1024))let cse_var_2: int32 (k.outer*4)let cse_var_1: int32 (cse_var_3 (y.outer*32))C[ramp(cse_var_1, 1, 32)] (C[ramp(cse_var_1, 1, 32)] (broadcast(A[((cse_var_3 cse_var_2) k.inner)], 32)*packedB_1[(((y.outer*1024) cse_var_2) k.inner)]))}}}}}}
}6-通过缓存优化块的写入
到目前为止我们所有的优化都集中在有效地访问和计算 A 和 B 矩阵的数据以计算 C 矩阵上。在阻塞优化之后运算器将逐块地将结果写入 C而且访问模式不是顺序的。我们可以通过使用一个顺序缓存数组来解决这个问题使用 cache_write、compute_at 和 unroll 的组合来保存块结果并在所有块结果准备好后写入 C。
cache_write和cache_read对应是先在shared memory中存放计算结果最后将结果写回到global memory。当然在真实的场景中我们往往是会将结果先放着register中最后写回。
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) // 默认情况下TVM将在函数的最外层或根计算张量。compute_at指定应在另一个运算符的第一个计算轴上计算一个张量。应该是为了方便后续的在xo维度上进行并行计算。# 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))main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {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_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) y)] B[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_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {let cse_var_4: int32 (k.outer*4)let cse_var_3: int32 (x.c*32)let cse_var_2: int32 ((y.outer*1024) cse_var_4)let cse_var_1: int32 (((x.outer*32768) (x.c*1024)) cse_var_4){C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 1)], 32)*packedB_1[(cse_var_2 1)]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 2)], 32)*packedB_1[(cse_var_2 2)]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 3)], 32)*packedB_1[(cse_var_2 3)]))}}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] C.global_1[((x.inner*32) y.inner)]}}}}}
}7-并行化
到目前为止我们的计算只被设计为使用单核。几乎所有的现代处理器都有多个内核计算可以从并行运行的计算中获益。最后的优化是利用线程级并行化的优势。
# parallel
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
)# Here is the generated IR after parallelization.
print(tvm.lower(s, [A, B, C], simple_modeTrue))parallelization: 0.026403
main 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, [1048576], []),B: Buffer(B_2: Pointer(float32), float32, [1048576], []),C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}buffer_map {A_1: A, B_1: B, C_1: C}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], [])} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope global {for (x: int32, 0, 32) parallel {for (y: int32, 0, 1024) {packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) y)] B[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_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {let cse_var_4: int32 (k.outer*4)let cse_var_3: int32 (x.c*32)let cse_var_2: int32 ((y.outer*1024) cse_var_4)let cse_var_1: int32 (((x.outer*32768) (x.c*1024)) cse_var_4){C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 1)], 32)*packedB_1[(cse_var_2 1)]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 2)], 32)*packedB_1[(cse_var_2 2)]))C.global_1[ramp(cse_var_3, 1, 32)] (C.global_1[ramp(cse_var_3, 1, 32)] (broadcast(A[(cse_var_1 3)], 32)*packedB_1[(cse_var_2 3)]))}}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C[((((x.outer*32768) (x.inner*1024)) (y.outer*32)) y.inner)] C.global_1[((x.inner*32) y.inner)]}}}}}
}总结
如前所述如何使用 TE 和调度原语进行优化可能需要对底层架构和算法有一些了解。然而TE 的设计是作为更复杂的算法的基础可以搜索潜在的优化。
本教程提供了一个 TVM 张量表达TE工作流程的演练使用了一个矢量添加和一个矩阵乘法的例子。一般的工作流程是
通过一系列的操作来描述你的计算。描述我们要如何计算使用调度原语。编译到我们想要的目标函数。可以选择保存该函数以便以后加载。
参考资料
TVM官方教程https://daobook.github.io/tvm/docs/tutorial/tensor_expr_get_started.html
