如何在CPU上优化GEMM矩阵乘法
How to optimize GEMM on CPU
(TL;DR) TVM 提供抽象接口，允许用户分别描述算法和算法的实现组织（所谓的调度）。通常，在高性能调度中编写算法会破坏算法的可读性和模块化。此外，尝试各种看似有希望的调度也很耗时。在 TVM 的帮助下，可以有效地尝试这些调度以提高性能。
在本文中，将演示如何使用 TVM 优化方阵乘法，通过简单地添加 18 行额外代码实现比基线快 200 倍。
在 CPU 上执行的密集计算应用程序有两个重要的优化：
 提高内存访问的缓存命中率。复杂的数值计算和热点内存访问都可以通过高缓存命中率加速。这需要将原始内存访问模式转换为适合缓存策略的模式。
 SIMD（单指令多数据），或者称之为向量处理单元。每次都会处理一小批数据，不是单个网格。这需要以统一模式转换循环体中的数据访问模式，以便 LLVM 后端可以将其降低为 SIMD。
实际上，本文中使用的所有方法，都是本repo 中提到的技巧的一个子集。其中一些已被 TVM 抽象自动应用，但由于 TVM 的限制，其中一些不能简单地应用。
下面提到的所有实验结果，都是在配备 Intel i7-4770HQ CPU 的 2015 年的 15’ MacBook 上执行的。对于所有 x86 CPU，缓存行大小应为 64 字节。
准备和基线
在本文中，将演示如何使用 TVM 优化矩阵乘法。在实际演示前，先定义这些变量。然后编写一个基线实现，这是在 TVM 中编写矩阵乘法的最简单方法。
import tvm
import tvm.testing
from tvm import te
import numpy
import timeit

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 type in tvm

dtype = “float32”

using Intel AVX2(Advanced Vector Extensions) ISA for SIMD

To get the best performance, please change the following line

to llvm -mcpu=core-avx2, or specific type of CPU you use

target = “llvm”
dev = tvm.device(target, 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)
np_repeat = 100
np_runing_time = timeit.timeit(
setup=“import numpy\n”
"M = " + str(M) + “\n”
"K = " + str(K) + “\n”
"N = " + str(N) + “\n”
‘dtype = “float32”\n’
“a = numpy.random.rand(M, K).astype(dtype)\n”
“b = numpy.random.rand(K, N).astype(dtype)\n”,
stmt=“answer = numpy.dot(a, b)”,
number=np_repeat,
)
print(“Numpy running time: %f” % (np_runing_time / np_repeat))
answer = numpy.dot(a.numpy(), b.numpy())

Algorithm

k = te.reduce_axis((0, K), “k”)
A = te.placeholder((M, K), name=“A”)
B = te.placeholder((K, N), name=“B”)
C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name=“C”)

Default schedule

s = te.create_schedule(C.op)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func

c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)

evaluator = func.time_evaluator(func.entry_name, dev, number=1)
print(“Baseline: %f” % evaluator(a, b, c).mean)
输出：
Numpy running time: 0.009229
Baseline: 3.340634
在 TVM 中，始终可以检查较低级别的 IR 以调试或优化调度。这是使用基线调度生成的 IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (m: int32, 0, 1024) {
for (n: int32, 0, 1024) {
C_2[((m1024) + n)] = 0f32
for (k: int32, 0, 1024) {
C_2[((m1024) + n)] = ((float32*)C_2[((m1024) + n)] + ((float32)A_2[((m1024) + k)](float32*)B_2[((k*1024) + n)]))
}
}
}
}
阻塞
提高缓存命中率的一个重要技巧是阻塞——数据块将逐块计算。块内部的内存访问是一个具有高内存局部性的小邻域。在本文中，选择了 32 作为分块因子。所以该块将填充32 * 32 * sizeof(float)，即总大小为32KB的缓存中的4KB（L1数据缓存）
bn = 32
kfactor = 4
s = te.create_schedule(C.op)

Blocking by loop tiling

mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(kaxis,) = s[C].op.reduce_axis
ko, ki = s[C].split(kaxis, factor=kfactor)

Hoist reduction domain outside the blocking loop

s[C].reorder(mo, no, ko, ki, mi, ni)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)

By simply tiling the loop 32x32, and hoisting ko, ki outside the blocking loops,

we can see big speedup compared with the baseline.

evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print(“Opt1: %f” % evaluator(a, b, c).mean)
输出：
Opt1: 0.295032
这是阻塞后生成的IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (m.outer: int32, 0, 32) {
for (n.outer: int32, 0, 32) {
for (m.inner.init: int32, 0, 32) {
for (n.inner.init: int32, 0, 32) {
C_2[((((m.outer32768) + (m.inner.init1024)) + (n.outer32)) + n.inner.init)] = 0f32
}
}
for (k.outer: int32, 0, 256) {
for (k.inner: int32, 0, 4) {
for (m.inner: int32, 0, 32) {
for (n.inner: int32, 0, 32) {
C_2[((((m.outer32768) + (m.inner1024)) + (n.outer32)) + n.inner)] = ((float32*)C_2[((((m.outer32768) + (m.inner1024)) + (n.outer32)) + n.inner)] + ((float32)A_2[((((m.outer32768) + (m.inner1024)) + (k.outer4)) + k.inner)](float32*)B_2[((((k.outer4096) + (k.inner1024)) + (n.outer*32)) + n.inner)]))
}
}
}
}
}
}
}
矢量化
另一个重要的技巧是矢量化。当内存访问模式一致时，编译器可以检测到这种模式，将连续内存传递给向量处理器。在 TVM 中，可以使用vectorize接口，提示编译器这种模式，这样就可以大大加速。
在本文中，选择矢量化内循环行数据，因为对缓存友好。
s = te.create_schedule(C.op)
mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(kaxis,) = s[C].op.reduce_axis
ko, ki = s[C].split(kaxis, factor=kfactor)
s[C].reorder(mo, no, ko, ki, mi, ni)

Vectorization

s[C].vectorize(ni)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print(“Opt2: %f” % evaluator(a, b, c).mean)
输出：
Opt2: 0.331193
这是矢量化后生成的 IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (m.outer: int32, 0, 32) {
for (n.outer: int32, 0, 32) {
for (m.inner.init: int32, 0, 32) {
C_2[ramp((((m.outer32768) + (m.inner.init1024)) + (n.outer32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (k.inner: int32, 0, 4) {
for (m.inner: int32, 0, 32) {
C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] = ((float32x32*)C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.inner1024)) + (k.outer4)) + k.inner)], 32)(float32x32*)B_2[ramp((((k.outer4096) + (k.inner1024)) + (n.outer*32)), 1, 32)]))
}
}
}
}
}
}
循环排列
如果查看上面的 IR，可以看到 B 和 C 的内循环行数据，都进行了向量化。接下来将查看 A 的访问模式。在当前调度中，A 是逐列访问的，这对缓存不友好. 如果改变 ki 和内轴 mi 的嵌套循环顺序，A 矩阵的访问模式对缓存更友好。
s = te.create_schedule(C.op)
mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(kaxis,) = s[C].op.reduce_axis
ko, ki = s[C].split(kaxis, factor=kfactor)

re-ordering

s[C].reorder(mo, no, ko, mi, ki, ni)
s[C].vectorize(ni)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print(“Opt3: %f” % evaluator(a, b, c).mean)
输出：
Opt3: 0.113024
这是循环排列后生成的 IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (m.outer: int32, 0, 32) {
for (n.outer: int32, 0, 32) {
for (m.inner.init: int32, 0, 32) {
C_2[ramp((((m.outer32768) + (m.inner.init1024)) + (n.outer32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (m.inner: int32, 0, 32) {
for (k.inner: int32, 0, 4) {
C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] = ((float32x32*)C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.inner1024)) + (k.outer4)) + k.inner)], 32)(float32x32*)B_2[ramp((((k.outer4096) + (k.inner1024)) + (n.outer*32)), 1, 32)]))
}
}
}
}
}
}
数组打包
另一个重要的技巧是数组打包。诀窍是对多维数组的存储进行重新排序，以便在展平，存储在一维内存中后按顺序访问。

注意：此图是阵列打包工作原理的一般说明。
可以使用数组打包来解决 B 的访问模式。观察扁平化后 B 的数组访问模式，当在 K 维度上迭代时，这不是连续的。可以用维度 [K][N] 重新排序 B，具有维度 [N/bn][K][bn]，其中 bn 是阻塞因子，也是内循环中 B 的向量大小。这种重新排序将 N 分成两个维度 — bigN (N/bn) 和 littleN (bn) — 新维度 [N/bn][K][bn] 匹配 B 从外循环到内循环的索引（no, ko, ki, ni) ，在展平后导致 B 的顺序访问模式。

We have to re-write the algorithm slightly.

packedB = te.compute(
(N / bn, K, bn), lambda bigN, k, littleN: B[k, bigN * bn + littleN], name=“packedB”
)
C = te.compute(
(M, N),
lambda m, n: te.sum(A[m, k] * packedB[n // bn, k, tvm.tir.indexmod(n, bn)], axis=k),
name=“C”,
)
s = te.create_schedule(C.op)
mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(kaxis,) = s[C].op.reduce_axis
ko, ki = s[C].split(kaxis, factor=kfactor)
s[C].reorder(mo, no, ko, mi, ki, ni)
s[C].vectorize(ni)
bigN, _, littleN = s[packedB].op.axis
s[packedB].vectorize(littleN)
s[packedB].parallel(bigN)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print(“Opt4: %f” % evaluator(a, b, c).mean)
输出：
Opt4: 0.232269
这是阵列打包后生成的IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (bigN: int32, 0, 32) “parallel” {
for (k: int32, 0, 1024) {
packedB[ramp(((bigN32768) + (k32)), 1, 32)] = (float32x32*)B_2[ramp(((k1024) + (bigN32)), 1, 32)]
}
}
for (m.outer: int32, 0, 32) {
for (n.outer: int32, 0, 32) {
for (m.inner.init: int32, 0, 32) {
C_2[ramp((((m.outer32768) + (m.inner.init1024)) + (n.outer32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (m.inner: int32, 0, 32) {
for (k.inner: int32, 0, 4) {
C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] = ((float32x32*)C_2[ramp((((m.outer32768) + (m.inner1024)) + (n.outer32)), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.inner1024)) + (k.outer4)) + k.inner)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + (k.inner*32)), 1, 32)]))
}
}
}
}
}
}
}
块的写缓存
阻塞后，程序将结果逐块写入C，访问模式不是顺序的。使用一个顺序缓存数组保存块结果，在所有块结果准备好时写入 C。
s = te.create_schedule(C.op)

Allocate write cache

CC = s.cache_write(C, “global”)
mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)

Write cache is computed at no

s[CC].compute_at(s[C], no)

New inner axes

mc, nc = s[CC].op.axis
(kaxis,) = s[CC].op.reduce_axis
ko, ki = s[CC].split(kaxis, factor=kfactor)
s[CC].reorder(ko, mc, ki, nc)
s[CC].vectorize(nc)

TODO: Add separate optimization step to discuss loop unrolloing

unrolling is a loop optimization strategy which can reduce branch

prediction failures and increases the chance of concurrent execution

unroll kfactor loops

s[CC].unroll(ki)
bigN, _, littleN = s[packedB].op.axis
s[packedB].vectorize(littleN)
s[packedB].parallel(bigN)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print(“Opt5: %f” % evaluator(a, b, c).mean)
输出：
Opt5: 0.225938
这是阻塞后生成的IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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 (bigN: int32, 0, 32) “parallel” {
for (k: int32, 0, 1024) {
packedB[ramp(((bigN32768) + (k32)), 1, 32)] = (float32x32*)B_2[ramp(((k1024) + (bigN32)), 1, 32)]
}
}
for (m.outer: int32, 0, 32) {
for (n.outer: int32, 0, 32) {
for (m.c.init: int32, 0, 32) {
C.global[ramp((m.c.init32), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (m.c: int32, 0, 32) {
C.global[ramp((m.c32), 1, 32)] = ((float32x32*)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[(((m.outer32768) + (m.c1024)) + (k.outer4))], 32)(float32x32*)packedB[ramp(((n.outer32768) + (k.outer128)), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 1)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 32), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 2)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 64), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 3)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 96), 1, 32)]))
}
}
for (m.inner: int32, 0, 32) {
for (n.inner: int32, 0, 32) {
C_2[((((m.outer32768) + (m.inner1024)) + (n.outer32)) + n.inner)] = (float32)C.global[((m.inner*32) + n.inner)]
}
}
}
}
}
}
并行化
此外，还可以利用多核处理器进行线程级并行化。
s = te.create_schedule(C.op)
CC = s.cache_write(C, “global”)
mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
s[CC].compute_at(s[C], no)
mc, nc = s[CC].op.axis
(kaxis,) = s[CC].op.reduce_axis
ko, ki = s[CC].split(kaxis, factor=kfactor)
s[CC].reorder(ko, mc, ki, nc)
s[CC].vectorize(nc)
s[CC].unroll(ki)

parallel

s[C].parallel(mo)
bigN, _, littleN = s[packedB].op.axis
s[packedB].vectorize(littleN)
s[packedB].parallel(bigN)
func = tvm.build(s, [A, B, C], target=target, name=“mmult”)
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
evaluator = func.time_evaluator(func.entry_name, dev, number=50)
opt6_time = evaluator(a, b, c).mean
print(“Opt6: %f” % opt6_time)
输出：
Opt6: 0.068730
这是并行化后生成的IR。
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
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} {
allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {
for (bigN: int32, 0, 32) “parallel” {
for (k: int32, 0, 1024) {
packedB[ramp(((bigN32768) + (k32)), 1, 32)] = (float32x32*)B_2[ramp(((k1024) + (bigN32)), 1, 32)]
}
}
for (m.outer: int32, 0, 32) “parallel” {
allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global;
for (n.outer: int32, 0, 32) {
for (m.c.init: int32, 0, 32) {
C.global[ramp((m.c.init32), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (m.c: int32, 0, 32) {
C.global[ramp((m.c32), 1, 32)] = ((float32x32*)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[(((m.outer32768) + (m.c1024)) + (k.outer4))], 32)(float32x32*)packedB[ramp(((n.outer32768) + (k.outer128)), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 1)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 32), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 2)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 64), 1, 32)]))
C.global[ramp((m.c32), 1, 32)] = ((float32x32)C.global[ramp((m.c32), 1, 32)] + (broadcast((float32)A_2[((((m.outer32768) + (m.c1024)) + (k.outer4)) + 3)], 32)(float32x32*)packedB[ramp((((n.outer32768) + (k.outer128)) + 96), 1, 32)]))
}
}
for (m.inner: int32, 0, 32) {
for (n.inner: int32, 0, 32) {
C_2[((((m.outer32768) + (m.inner1024)) + (n.outer32)) + n.inner)] = (float32)C.global[((m.inner*32) + n.inner)]
}
}
}
}
}
}
概括
在仅用 18 行代码应用上述简单优化后，生成的代码可以使用 MKL实现numpy性能的60%。网页上的输出反映了非排他性 Docker 容器上的运行时间，因此不可靠。强烈建议运行本文，观察 TVM 实现的性能提升。

参考链接：
https://tvm.apache.org/docs/how_to/optimize_operators/opt_gemm.html

如何在CPU上优化GEMM矩阵乘法相关推荐

如何在 CPU 上优化 GEMM
如何在 CPU 上优化 GEMM (TL;DR) TVM 提供抽象接口,允许用户分别描述算法和算法的实施组织(所谓的调度).通常,在高性能调度中编写算法,会破坏算法的可读性和模块化.尝试各种看似有前途 ...
如何在CPU上优化GEMM（上）
如何在CPU上优化GEMM(上) (TL:DR)TVM提供了抽象接口,用户分别描述算法和算法的实现组织(所谓的调度).通常,在高性能调度中编写算法会破坏算法的可读性和模块性.尝试各种看似有希望的时间表 ...
如何在CPU上优化GEMM（下）
如何在CPU上优化GEMM(下) Array Packing 另一个重要的技巧是数组打包.这个技巧是对数组的存储维度进行重新排序,将某个维度上的连续访问模式在平滑后转换为顺序模式. 如上图所示,在阻塞 ...
如何在 GPU 上优化卷积
如何在 GPU 上优化卷积将演示如何在 TVM 中编写高性能卷积实现.正方形大小的输入张量和过滤器为例,假设卷积的输入具有大batch批量.在这个例子中,使用不同的布局存储数据,实现更好的数据局部性 ...
如何在GPU上优化卷积
如何在GPU上优化卷积本文将演示如何在TVM中编写高性能的卷积实现.以平方大小的输入张量和滤波器为例,并假设卷积的输入量很大.使用不同的布局来存储数据,以实现更好的数据局部性.缓冲区布局为HWCN, ...
CPU的自动调度矩阵乘法
CPU的自动调度矩阵乘法这是一个有关如何对CPU使用自动调度程序的文档. 与依靠手动模板定义搜索空间的基于模板的autotvm不同,自动调度程序不需要任何模板.用户只需要编写计算声明,而无需任何调度 ...
DeepMind攻克50年数学难题！AlphaZero史上最快矩阵乘法算法登Nature封面
新智元报道编辑:David Joey [新智元导读]DeepMind碾压人类高手的AI围棋大师AlphaZero,下一个目标是数学算法!现已发现50年以来最快的矩阵乘法算法. 下围棋碾压人类 ...
性能比拼！超详细的Tengine GEMM矩阵乘法汇编教程
点击我爱计算机视觉标星,更快获取CVML新技术 Tengine 是OPEN AI LAB 针对前端智能设备开发的软件开发包,核心部分是一个轻量级,模块化,高性能的AI 推断引擎,并支持用DLA.GPU ...
详解如何在STM32上使用4x4矩阵键盘
矩阵键盘的使用在日常生活中非常常见,驱动矩阵键盘其实就是单纯的使用GPIO去获取电平的状态,但如何获取其实有很多方法,常见的方法有轮询和中断,本次介绍一个比较常用的方法,通过定时器轮询来进行读取(采用 ...

如何在CPU上优化GEMM矩阵乘法