【TVM 教程】如何在 CPU 上最佳化 GEMM

超神经HyperAI發表於2024-12-04

Apache TVM 是一個端到端的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速晶片。更多 TVM 中文文件可訪問 → Apache TVM 是一個端到端的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速晶片。 | Apache TVM 中文站

TVM 提供抽象介面，允許使用者分別描述演算法和演算法的實現（所謂的排程）。通常，以高效能排程編寫演算法會破壞演算法的可讀性和模組化。此外，嘗試各種看似有希望的 schedules 非常耗時。在 TVM 的幫助下，可以有效地嘗試這些 schedules 以提高效能。

本教程將演示如何用 TVM 最佳化矩陣乘法，並透過 18 行程式碼實現比 baseline 快 200 倍的效能。

在 CPU 上執行的密集計算應用程式有兩個重要的最佳化：

  1. 提高記憶體訪問的 cache 命中率。高 cache 命中率可以加速複雜的數值計算和熱點記憶體訪問。需要將原始記憶體訪問模式轉換為適合 cache 策略的模式。
  2. SIMD（單指令多資料），或者稱之為向量處理單元，每次都會處理一小批資料，而不是單個網格。需要將迴圈體中的資料訪問模式轉換為統一模式，以便 LLVM 後端可以將其降級到 SIMD。

實際上，本教程中使用的所有方法都在這個 repo 中提到了。其中一些已被 TVM 抽象自動應用，但有一些由於 TVM 的限制，不能被簡單地應用。

下面提到的所有實驗結果，都是在配備 Intel i7-4770HQ CPU 的 2015 年 15 英寸 MacBook 上執行的，所有 x86 CPU 的快取記憶體行大小應為 64 位元組。

準備和 baseline​

本教程演示如何使用 TVM 最佳化矩陣乘法。實際演示前，首先定義這些變數。然後編寫一個 baseline 實現，這是在 TVM 中編寫矩陣乘法的最簡單方法。

import tvm
import tvm.testing
from tvm import te
import numpy
import timeit

# 矩陣的大小
# (M, K) x (K, N)
# 可自由嘗試不同的 shapes，有時 TVM 最佳化在 MKL 中的表現優於 numpy。
M = 1024
K = 1024
N = 1024

# tvm 中的預設張量型別
dtype = "float32"

# 為 SIMD 使用英特爾 AVX2（高階向量擴充套件）ISA
# 要獲得最佳效能，更改以下行
# 為 llvm -mcpu=core-avx2，或者使用的特定型別的 CPU
target = "llvm"
dev = tvm.device(target, 0)

# 用於測試的隨機生成張量
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())

# 演算法
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")

# 預設 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.018437
Baseline: 3.336375

在 TVM 中，始終可以檢查較低階別的 IR 以除錯或最佳化 schedule。這是使用 baseline schedule 生成的 IR。

print(tvm.lower(s, [A, B, C], simple_mode=True))

輸出結果：

@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 (m: int32, 0, 1024) {
 for (n: int32, 0, 1024) {
      C[((m*1024) + n)] = 0f32
      for (k: int32, 0, 1024) {
        let cse_var_2: int32 = (m*1024)
        let cse_var_1: int32 = (cse_var_2 + n)
        C[cse_var_1] = (C[cse_var_1] + (A[(cse_var_2 + k)]*B[((k*1024) + n)]))
 }
 }
 }
}

分塊​

提高快取命中率的一個重要技巧是分塊——資料塊將逐塊計算。塊內的記憶體訪問是一個區域性性的小鄰域。本教程選擇 32 作為分塊因子，因此該塊將填充 32 32 sizeof(float) ，即總大小為 32KB 的快取中的 4KB（L1 資料快取）。

bn = 32
kfactor = 4
s = te.create_schedule(C.op)

# 透過迴圈 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)

# 將 reduction 域提升到分塊迴圈之外
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)

# 透過簡單地將迴圈 32x32 分塊，並將 ko、ki 提升到分塊迴圈之外，
# 可以看到與 baseline 相比，加速有很大提升。
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print("Opt1: %f" % evaluator(a, b, c).mean)

輸出結果：

Opt1: 0.307321

分塊後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))
@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 (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[((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)) + 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) {
 let cse_var_3: int32 = (n.outer*32)
 let cse_var_2: int32 = ((m.outer*32768) + (m.inner*1024))
 let cse_var_1: int32 = ((cse_var_2 + cse_var_3) + n.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) + n.inner)]))
 }
 }
 }
 }
 }
 }
}

向量化​

另一個重要技巧是向量化，當記憶體訪問模式一致時，編譯器可以檢測到這種模式並將連續記憶體傳遞給向量處理器。TVM 中可以用 vectorize 介面來提示編譯器這種模式，這樣就可以進行加速。

本教程選擇向量化內部迴圈 row data（對快取更友好）。

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)

# 向量化
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.349439

向量化後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))

輸出結果：

@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 (m.outer: int32, 0, 32) {
 for (n.outer: int32, 0, 32) {
 for (m.inner.init: int32, 0, 32) {
        C[ramp((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)), 1, 32)] = broadcast(0f32, 32)
 }
 for (k.outer: int32, 0, 256) {
 for (k.inner: int32, 0, 4) {
 for (m.inner: int32, 0, 32) {
 let cse_var_3: int32 = (n.outer*32)
 let cse_var_2: int32 = ((m.outer*32768) + (m.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)]))
 }
 }
 }
 }
 }
}

迴圈置換​

檢視上面的 IR，可以看到內部迴圈的 row data 對於 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)

# 重新排序
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.115375

迴圈置換後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))

輸出結果：

@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 (m.outer: int32, 0, 32) {
 for (n.outer: int32, 0, 32) {
 for (m.inner.init: int32, 0, 32) {
        C[ramp((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)), 1, 32)] = broadcast(0f32, 32)
 }
 for (k.outer: int32, 0, 256) {
 for (m.inner: int32, 0, 32) {
 for (k.inner: int32, 0, 4) {
 let cse_var_3: int32 = (n.outer*32)
 let cse_var_2: int32 = ((m.outer*32768) + (m.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)]))
 }
 }
 }
 }
 }
}

陣列打包​

另一個重要的技巧是陣列打包，對多維陣列的儲存進行重新排序，展平並儲存在一維記憶體中，方便順序訪問。

可以用陣列打包來解決 B 的訪問模式。觀察展平後 B 的陣列訪問模式，當迭代 K 維時，它不是順序的。可以用維度 K 對 B 重新排序，使其具有 N/bn[bn] 維度，其中 bn 是分塊因子，也是內迴圈中 B 的向量大小。

這種重新排序將 N 拆分為兩個維度——bigN（N/bn）和 littleN（bn）——新維度 N/bn[bn] 匹配 B 從外部到內部迴圈的索引（no, ko, ki, ni) 在展平後導致 B 的順序訪問模式。

# 我們必須稍微重新編寫演算法。
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.109499

陣列打包後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))

輸出結果：

@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 (bigN: int32, 0, 32) "parallel" {
 for (k: int32, 0, 1024) {
        packedB_1: Buffer(packedB, float32x32, [32768], [])[((bigN*1024) + k)] = B[ramp(((k*1024) + (bigN*32)), 1, 32)]
 }
 }
 for (m.outer: int32, 0, 32) {
 for (n.outer: int32, 0, 32) {
 for (m.inner.init: int32, 0, 32) {
          C[ramp((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)), 1, 32)] = broadcast(0f32, 32)
 }
 for (k.outer: int32, 0, 256) {
 for (m.inner: int32, 0, 32) {
 for (k.inner: int32, 0, 4) {
 let cse_var_3: int32 = ((m.outer*32768) + (m.inner*1024))
 let cse_var_2: int32 = (k.outer*4)
 let cse_var_1: int32 = (cse_var_3 + (n.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[(((n.outer*1024) + cse_var_2) + k.inner)]))
 }
 }
 }
 }
 }
 }
}

塊的寫快取​

分塊後，程式會逐塊將結果寫入 C（訪問模式不是順序的），因此，可以使用順序快取陣列來儲存塊結果，並在所有塊結果準備好時寫入 C。

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)

# 寫快取在 no 被計算
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)

# TODO: 新增單獨的最佳化步驟，來討論迴圈展開
# unrolling 是一種迴圈最佳化策略，可以減少分支
# 預測失敗，以及增加併發執行的機會
# 展開 kfactor 迴圈
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.110823

分塊後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))
@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 (bigN: int32, 0, 32) "parallel" {
 for (k: int32, 0, 1024) {
        packedB_1: Buffer(packedB, float32x32, [32768], [])[((bigN*1024) + k)] = B[ramp(((k*1024) + (bigN*32)), 1, 32)]
 }
 }
 for (m.outer: int32, 0, 32) {
 for (n.outer: int32, 0, 32) {
 for (m.c.init: int32, 0, 32) {
          C.global_1: Buffer(C.global, float32, [1024], [])[ramp((m.c.init*32), 1, 32)] = broadcast(0f32, 32)
        }
        for (k.outer: int32, 0, 256) {
          for (m.c: int32, 0, 32) {
            let cse_var_4: int32 = (k.outer*4)
            let cse_var_3: int32 = (m.c*32)
            let cse_var_2: int32 = ((n.outer*1024) + cse_var_4)
            let cse_var_1: int32 = (((m.outer*32768) + (m.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 (m.inner: int32, 0, 32) {
 for (n.inner: int32, 0, 32) {
            C[((((m.outer*32768) + (m.inner*1024)) + (n.outer*32)) + n.inner)] = C.global_1[((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)

# 並行
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.144875

並行化後生成的 IR：

print(tvm.lower(s, [A, B, C], simple_mode=True))

輸出結果：

@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 (bigN: int32, 0, 32) "parallel" {
 for (k: int32, 0, 1024) {
        packedB_1: Buffer(packedB, float32x32, [32768], [])[((bigN*1024) + k)] = B[ramp(((k*1024) + (bigN*32)), 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_1: Buffer(C.global, float32, [1024], [])[ramp((m.c.init*32), 1, 32)] = broadcast(0f32, 32)
        }
        for (k.outer: int32, 0, 256) {
          for (m.c: int32, 0, 32) {
            let cse_var_4: int32 = (k.outer*4)
            let cse_var_3: int32 = (m.c*32)
            let cse_var_2: int32 = ((n.outer*1024) + cse_var_4)
            let cse_var_1: int32 = (((m.outer*32768) + (m.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 (m.inner: int32, 0, 32) {
 for (n.inner: int32, 0, 32) {
            C[((((m.outer*32768) + (m.inner*1024)) + (n.outer*32)) + n.inner)] = C.global_1[((m.inner*32) + n.inner)]
 }
 }
 }
 }
 }
}

總結​

應用上述簡單最佳化後，僅用 18 行程式碼，就可以達到使用 MKL numpy 效能的 60%。注意，網頁上的輸出反映了非專有 Docker 容器上的執行時間，是不可靠的。推薦自己執行本教程，觀察 TVM 的效能提升。

下載 Python 原始碼：opt_gemm.py

下載 Jupyter Notebook：opt_gemm.ipynb

相關文章