當(dāng)前位置：首頁(yè) > 人文社科 > 生活经验 >内容正文

生活经验

如何在 CPU 上优化 GEMM

發(fā)布時(shí)間：2023/11/28 生活经验 35 豆豆

生活随笔收集整理的這篇文章主要介紹了如何在 CPU 上优化 GEMM 小編覺(jué)得挺不錯(cuò)的,現(xiàn)在分享給大家,幫大家做個(gè)參考.

如何在 CPU 上優(yōu)化 GEMM
(TL;DR) TVM 提供抽象接口，允許用戶(hù)分別描述算法和算法的實(shí)施組織（所謂的調(diào)度）。通常，在高性能調(diào)度中編寫(xiě)算法，會(huì)破壞算法的可讀性和模塊化。嘗試各種看似有前途的調(diào)度也很耗時(shí)。在 TVM 的幫助下，可以有效地嘗試這些調(diào)度，提高性能。
將演示如何使用 TVM 優(yōu)化矩陣乘法，通過(guò)簡(jiǎn)單地添加 18 行額外代碼，實(shí)現(xiàn)比基線(xiàn)快 200 倍。
在 CPU 上執(zhí)行的密集計(jì)算應(yīng)用程序，有兩個(gè)重要的優(yōu)化：

提高內(nèi)存訪(fǎng)問(wèn)的緩存命中率。復(fù)雜的數(shù)值計(jì)算和熱點(diǎn)內(nèi)存訪(fǎng)問(wèn)，都可以通過(guò)高緩存命中率加速。需要將原始內(nèi)存訪(fǎng)問(wèn)模式，轉(zhuǎn)換為適合緩存策略的模式。
SIMD（單指令多數(shù)據(jù)），或者稱(chēng)向量處理單元。每次都會(huì)處理一小批數(shù)據(jù)，不是單個(gè)網(wǎng)格。需要統(tǒng)一模式，轉(zhuǎn)換循環(huán)體中的數(shù)據(jù)訪(fǎng)問(wèn)模式， LLVM 后端可以降低為 SIMD。
實(shí)際上，所有方法都是repo 中提到的一個(gè)子技巧。一些已被 TVM 抽象自動(dòng)應(yīng)用，由于 TVM 的限制，一些不能簡(jiǎn)單地應(yīng)用。
下面所有實(shí)驗(yàn)結(jié)果，都是在配備 Intel i7-4770HQ CPU 的，15’ MacBook 上實(shí)現(xiàn)的。對(duì)于所有 x86 CPU，緩存行大小應(yīng)為 64 字節(jié)。
準(zhǔn)備和基線(xiàn)
將演示如何使用 TVM 優(yōu)化矩陣乘法。在實(shí)際演示前，先定義這些變量。然后編寫(xiě)一個(gè)基線(xiàn)實(shí)現(xiàn)，這是在 TVM 中編寫(xiě)矩陣乘法的最簡(jiǎn)單方法。
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 = “l(fā)lvm”
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.009345
Baseline: 3.291115
在 TVM 中，檢查較低級(jí)別的 IR，調(diào)試或優(yōu)化調(diào)度。這是使用基線(xiàn)調(diào)度生成的 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)]))
}
}
}
}
阻塞
提高緩存命中率的一個(gè)重要技巧是阻塞——數(shù)據(jù)塊將逐塊計(jì)算。塊內(nèi)部的內(nèi)存訪(fǎng)問(wèn)是一個(gè)具有高內(nèi)存局部性的小鄰域。選擇了 32 作為分塊因子。該塊將填充 32 * 32 * sizeof(float) 即 4KB 的緩存，總大小為 32KB（L1 數(shù)據(jù)緩存）。
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.310688
這是阻塞后，生成的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)]))
}
}
}
}
}
}
}
向量化
另一個(gè)重要的技巧是向量化。當(dāng)內(nèi)存訪(fǎng)問(wèn)模式一致時(shí)，編譯器可以檢測(cè)到這種模式，將連續(xù)內(nèi)存?zhèn)鬟f給向量處理器。在 TVM 中，可以使用vectorize接口，提示編譯器這種模式，這樣就可以大大加速。
選擇向量化內(nèi)循環(huán)行數(shù)據(jù)，這是緩存友好的。
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.341067
向量化后，生成的 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)]))
}
}
}
}
}
}
循環(huán)排列
查看上面的 IR，可以看到 B 和 C 的內(nèi)循環(huán)行數(shù)據(jù)，都進(jìn)行了向量化。接下來(lái)，查看 A 的訪(fǎng)問(wèn)模式。在當(dāng)前調(diào)度中，A 是逐列訪(fǎng)問(wèn)的，這對(duì)緩存不友好. 如果改變 ki 和內(nèi)軸 mi 的嵌套循環(huán)順序，A 矩陣的訪(fǎng)問(wèn)模式，對(duì)緩存更友好。
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.111449
循環(huán)排列后，生成的 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)]))
}
}
}
}
}
}
陣列封裝
另一個(gè)重要的技巧是數(shù)組打包。訣竅是對(duì)多維數(shù)組的存儲(chǔ)，進(jìn)行重新排序，展平存儲(chǔ)在一維內(nèi)存中后，按順序訪(fǎng)問(wèn)。
可以使用數(shù)組打包，解決 B 的訪(fǎng)問(wèn)模式。觀(guān)察扁平化后 B 的數(shù)組訪(fǎng)問(wèn)模式，在 K 維度上迭代時(shí)，這不是連續(xù)的。可以用維度 [K][N] 重新排序 B，使其具有維度 [N/bn][K][bn]，bn 是阻塞因子，也是內(nèi)循環(huán)中 B 的向量大小。這種重新排序，將 N 分成兩個(gè)維度 — bigN (N/bn) 和 littleN (bn) —新維度 [N/bn][K][bn] 匹配 B，從外循環(huán)到內(nèi)循環(huán)的索引（no, ko, ki, ni) ，在展平后，導(dǎo)致 B 的順序訪(fǎng)問(wèn)模式。

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.217310
陣列打包后，生成的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)]))
}
}
}
}
}
}
}
塊的寫(xiě)緩存
阻塞后，程序?qū)⒔Y(jié)果逐塊寫(xiě)入C，訪(fǎng)問(wèn)模式不是順序的。可以使用一個(gè)順序緩存數(shù)組，保存塊結(jié)果，在所有塊結(jié)果準(zhǔn)備好時(shí)，寫(xiě)入 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.215912
阻塞后，生成的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)]
}
}
}
}
}
}
并行化
可以利用多核處理器，進(jìn)行線(xiàn)程級(jí)并行化。
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.066558
并行化后，生成的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) “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)]
}
}
}
}
}
}
總結(jié)
僅用 18 行代碼，應(yīng)用上述簡(jiǎn)單優(yōu)化后，生成的代碼，可以使用 MKL實(shí)現(xiàn)numpy性能的60% 。輸出反映了非排他性 Docker 容器上的運(yùn)行時(shí)間，是不可靠的。強(qiáng)烈建議自己運(yùn)行，觀(guān)察 TVM 實(shí)現(xiàn)的性能提升。

參考鏈接：
https://tvm.apache.org/docs/tutorials/optimize/opt_gemm.html

總結(jié)

以上是生活随笔為你收集整理的如何在 CPU 上优化 GEMM的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。

上一篇：一些量化(quantization)技巧
下一篇： Conda安装Glossary词汇表

日韩av黄I国产麻豆传媒I国产91av视频在线观看I日韩一区二区三区在线看I美女国产在线I麻豆视频国产在线观看I成人黄色短片

生活经验

如何在 CPU 上优化 GEMM

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.

The default tensor type in tvm

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

Random generated tensor for testing

Algorithm

Default schedule

Blocking by loop tiling

Hoist reduction domain outside the blocking loop

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

we can see big speedup compared with the baseline.

Vectorization

re-ordering

We have to re-write the algorithm slightly.

Allocate write cache

Write cache is computed at no

New inner axes

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

parallel

總結(jié)