TVM：使用Tensor Expression (TE)來(lái)處理算子

在本教程中，我們將聚焦于在 TVM 中使用張量表達(dá)式（TE）來(lái)定義張量計(jì)算和實(shí)現(xiàn)循環(huán)優(yōu)化。TE用純函數(shù)語(yǔ)言描述張量計(jì)算（即每個(gè)表達(dá)式都沒(méi)有副作用）。當(dāng)在 TVM 的整體上下文中查看時(shí)，Relay 將計(jì)算描述為一組算子，并且其中每一個(gè)算子都可以表示為 TE 表達(dá)式，每個(gè) TE 表達(dá)式獲取輸入張量并生成輸出張量。
本文是TVM中 TE 語(yǔ)言的入門(mén)教程。TVM 使用領(lǐng)域?qū)Ｓ?#xff08;domain specific）的張量表達(dá)式來(lái)高效地構(gòu)造內(nèi)核。我們以?xún)蓚€(gè)使用 TE 語(yǔ)言的為例來(lái)演示基本工作流。第一個(gè)示例介紹了 TE 和帶有向量加法的 schedule。第二個(gè)示例通過(guò)逐步優(yōu)化矩陣與 TE 的乘法來(lái)擴(kuò)展這些概念。這個(gè)矩陣乘法示例將作為未來(lái)涵蓋更高級(jí)的 TVM 特性的教程的對(duì)比基礎(chǔ)。

示例一：使用TE為CPU編寫(xiě)和調(diào)度向量加法

初始化 tvm環(huán)境

我們的第一個(gè)例子是使用 Python 來(lái)為向量加法實(shí)現(xiàn)一個(gè) TE，然后是一個(gè)針對(duì) CPU 的 schedule，我們從初始化 tvm 環(huán)境開(kāi)始：

import tvm import tvm.testing from tvm import te import numpy as np# 如果能夠指定目標(biāo) CPU，那么將會(huì)得到更好地性能 # 如果用的是llvm，可以通過(guò) `llc --version` 來(lái)查看 CPU 類(lèi)型 # 可以通過(guò)查看 /proc/cpuinfo 來(lái)查看你的處理器可能支持的其他擴(kuò)展， # 比如，如果你的 CPU 有 AVX-512 指令集，那么你可以使用 `llvm -mcpu=skylake-avx512` 選項(xiàng)tgt = tvm.target.Target(target="llvm", host="llvm")

描述向量計(jì)算

我們首先描述向量加法計(jì)算。TVM 采用張量語(yǔ)義，每個(gè)中間結(jié)果表示為一個(gè)多維數(shù)組。我們需要描述規(guī)則來(lái)得到張量。我們首先定義一個(gè)符號(hào)變量 n 來(lái)表示形狀。然后我們定義兩個(gè) placeholder 張量：A、B，它們的形狀是 (n,)。然后我們通過(guò)一個(gè) compute 操作，得到結(jié)果張量 C。compute 定義了一種計(jì)算，其輸出符合指定的張量形狀，并在由 lambda 函數(shù)定義的張量中的每個(gè)位置執(zhí)行計(jì)算。注意，雖然 n 是一個(gè)變量，但它定義了A、B 和 C 張量之間的一致形狀。請(qǐng)注意，在這個(gè)階段沒(méi)有實(shí)際的計(jì)算發(fā)生，因?yàn)槲覀冎皇锹暶鲬?yīng)該如何進(jìn)行計(jì)算。

n = te.var("n") A = te.placeholder((n,), name="A") B = te.placeholder((n,), name="B") C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")

注意：Lambda函數(shù)

te.compute方法的第二個(gè)參數(shù)是執(zhí)行計(jì)算的函數(shù)。在本例中，我們使用一個(gè)匿名函數(shù)（也稱(chēng)為lambda函數(shù)）來(lái)定義計(jì)算，在本例中是對(duì) a 和 B 的第 i 個(gè)元素的加法。

為計(jì)算創(chuàng)建一個(gè)默認(rèn)的Schedule

雖然上面幾行描述了計(jì)算規(guī)則，但我們可以用許多不同的方法計(jì)算 C 以適應(yīng)不同的設(shè)備。對(duì)于具有多個(gè) axis 的張量，您可以選擇首先迭代哪個(gè) axis ，另外計(jì)算可以跨不同的線程拆分。TVM要求用戶提供一個(gè) schedule，來(lái)描述應(yīng)如何執(zhí)行計(jì)算。TE 中的 schedule 操作可以更改循環(huán)順序、跨不同線程拆分計(jì)算、將數(shù)據(jù)塊分組在一起，以及其他操作。schedule 背后的一個(gè)重要概念是，它們只描述如何執(zhí)行計(jì)算，因此相同 TE 的不同 schedule 一定會(huì)產(chǎn)生相同的結(jié)果。

在 TVM 中，我們可以創(chuàng)建一種樸素的 schedule ，按照行優(yōu)先的順序來(lái)計(jì)算 C。

for (int i = 0; i < n; ++i) {C[i] = A[i] + B[i]; } s = te.create_schedule(C.op)

編譯并驗(yàn)證默認(rèn)的 schedule

通過(guò) TE 表達(dá)式和 schedule，我們可以為目標(biāo)語(yǔ)言和體系結(jié)構(gòu)生成可運(yùn)行的代碼，在本例中是 LLVM 和 CPU 。我們向 TVM 提供 schedule、schedule 中的TE表達(dá)式列表、目標(biāo)和主機(jī)，以及我們正在生成的函數(shù)的名稱(chēng)。輸出的結(jié)果是可以直接從 Python 調(diào)用 type-erased 函數(shù)。

在下一行中，我們使用 tvm.build 創(chuàng)建一個(gè)函數(shù)。build 函數(shù)接受 schedule、函數(shù)所需的簽名（包括輸入和輸出）以及我們要編譯到的目標(biāo)語(yǔ)言。

fadd = tvm.build(s, [A, B, C], tgt, name="myadd")

我們運(yùn)行該函數(shù)，并將輸出與 numpy 中的相同計(jì)算進(jìn)行比較。編譯后的 TVM 函數(shù)提供了一個(gè)簡(jiǎn)明的C API，可以被任何語(yǔ)言調(diào)用。我們首先創(chuàng)建一個(gè)設(shè)備（在本例中為CPU），這是一個(gè) TVM 可以編譯 schedule 的設(shè)備。在本例中，設(shè)備是LLVM CPU target。然后，我們可以在設(shè)備中初始化張量并執(zhí)行自定義的加法操作。為了驗(yàn)證計(jì)算的正確性，我們可以將c張量的輸出結(jié)果與 numpy 執(zhí)行的相同計(jì)算進(jìn)行比較。

dev = tvm.device(tgt.kind.name, 0)n = 1024 a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev) b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev) c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev) fadd(a, b, c) tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

為了對(duì)比這個(gè)樸素版本的自定義向量加法與 numpy 的速度差異，創(chuàng)建一個(gè)輔助函數(shù)來(lái)運(yùn)行 TVM 生成代碼的 profile。

import timeitnp_repeat = 100 np_running_time = timeit.timeit(setup="import numpy\n""n = 32768\n"'dtype = "float32"\n'"a = numpy.random.rand(n, 1).astype(dtype)\n""b = numpy.random.rand(n, 1).astype(dtype)\n",stmt="answer = a + b",number=np_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(size=n).astype(A.dtype), dev)b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)evaluator = func.time_evaluator(func.entry_name, dev, number=10)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", log=log)

此處輸出：

Numpy running time: 0.000008 naive: 0.000006

使用并行性（paralleism）來(lái)優(yōu)化 schedule

我們已經(jīng)說(shuō)明了 TE 的基本原理，現(xiàn)在讓我們更深入地了解 schedule 的作用，以及它們?nèi)绾斡糜趦?yōu)化不同體系結(jié)構(gòu)的張量表達(dá)式。schedule 是應(yīng)用于表達(dá)式的一系列步驟，用于以多種不同方式對(duì)其進(jìn)行轉(zhuǎn)換。當(dāng)一個(gè) schedule 應(yīng)用于TE中的一個(gè)表達(dá)式時(shí)，輸入和輸出保持不變，但在編譯時(shí)，表達(dá)式的實(shí)現(xiàn)可能會(huì)改變。在默認(rèn) schedule 中，這個(gè)張量加法是串行運(yùn)行的，但該操作其實(shí)是很容易在所有處理器線程之間并行。我們可以將我們的操作并行調(diào)度到計(jì)算中：

s[C].parallel(C.op.axis[0])

tvm.lower 命令將生成 TE 的中間表示（IR）以及相應(yīng)的 schedule 。通過(guò)在執(zhí)行不同的 schedule 操作時(shí) lowing 表達(dá)式，我們可以看到 schedule 對(duì)計(jì)算順序的影響。我們使用標(biāo)志 simple_mode=True 返回可讀的 C 風(fēng)格語(yǔ)句。

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, [n: int32], [stride: int32], type="auto"),A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}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現(xiàn)在可以在獨(dú)立的線程上運(yùn)行這些塊。我們?cè)趫?zhí)行并行操作的情況下編譯并運(yùn)行這個(gè)新的 schedule：

fadd_parallel = tvm.build(s, [A, B, C], tgt, name="myadd_parallel") fadd_parallel(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())evaluate_addition(fadd_parallel, tgt, "parallel", log=log)

此處輸出：

parallel: 0.000005

使用矢量化（vectorization）來(lái)優(yōu)化 schedule

現(xiàn)代 CPU 能夠?qū)Ω↑c(diǎn)數(shù)進(jìn)行 SIMD 操作，我們可以對(duì)計(jì)算表達(dá)式使用另一個(gè) schedule 來(lái)利用這一點(diǎn)。實(shí)現(xiàn)這一點(diǎn)需要多個(gè)步驟：首先，我們必須使用 split scheduling 原語(yǔ)將 schedule 拆分為內(nèi)部循環(huán)和外部循環(huán)。內(nèi)部循環(huán)可以使用向量化來(lái)使用使用向量化調(diào)度原語(yǔ)的 SIMD 指令，然后外部循環(huán)可以使用并行調(diào)度原語(yǔ)進(jìn)行并行化。選擇分割因子作為CPU上的線程數(shù)。

注：SIMD，全稱(chēng) Single Instruction Multiple Data，單指令多數(shù)據(jù)流，能夠復(fù)制多個(gè)操作數(shù)，并把它們打包在大型寄存器的一組指令集。

# 由于我們需要修改之前例子中的并行操作，因此這里要重建 schedule n = te.var("n") A = te.placeholder((n,), name="A") B = te.placeholder((n,), name="B") C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")s = te.create_schedule(C.op)# factor 的選擇需要適合你的線程數(shù)，這取決于架構(gòu)， # 建議將此系數(shù)設(shè)置為等于可用CPU核心數(shù)。 factor = 4outer, inner = s[C].split(C.op.axis[0], factor=factor) s[C].parallel(outer) s[C].vectorize(inner)fadd_vector = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")evaluate_addition(fadd_vector, tgt, "vector", log=log)print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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], type="auto"),C: Buffer(C_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}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), dtype=bool) {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)])}}} }

對(duì)比不同的 schedule

下面我們來(lái)對(duì)比以下之前提到的不同 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

代碼專(zhuān)門(mén)化

正如我們所看到的，A、B 和 C 的聲明都采用相同的形狀參數(shù) n。TVM將利用這一點(diǎn)，只向 kernel 傳遞一個(gè) shape 參數(shù)，我們?cè)诖蛴〉脑O(shè)備代碼中找到它。這是專(zhuān)門(mén)化化的一種形式。

在 host 端，TVM 將自動(dòng)生成檢查代碼，以檢查參數(shù)中的約束。因此，如果將具有不同形狀的數(shù)組傳遞到 fadd 中，將引發(fā)錯(cuò)誤。

我們可以做更多的專(zhuān)門(mén)化。例如，我們可以在計(jì)算聲明中寫(xiě)入n=tvm.runtime.convert（1024）而不是 n=te.var（“n”）。生成的函數(shù)將只獲取長(zhǎng)度為1024的向量。

我們已經(jīng)定義、調(diào)度并編譯了一個(gè)向量加法運(yùn)算符，然后可以在 TVM Runtime 執(zhí)行它。我們可以將算子保存為庫(kù)，稍后可以使用 TVM Runtime 加載該庫(kù)。

針對(duì)GPU的矩陣加法（可選）

在介紹保存與加載自定義算子庫(kù)的方法之前，我們先來(lái)看一下如何針對(duì) GPU 做矩陣加法。

TVM能夠針對(duì)多種體系結(jié)構(gòu)。在本例，我們將針對(duì)GPU中矢量加法的編譯。

# 本段代碼默認(rèn)不運(yùn)行，如果想要運(yùn)行的話，請(qǐng)將 ``run_cuda = True``run_cuda = False if run_cuda:# 這里的 target 需要根據(jù)自己的 GPU 類(lèi)型修改：# NVIDIA：cuda# Radeon：rocm# OpenCL：opencltgt_gpu = tvm.target.Target(target="cuda", host="llvm")# 重建 schedulen = te.var("n")A = te.placeholder((n,), name="A")B = te.placeholder((n,), name="B")C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")print(type(C))s = te.create_schedule(C.op)bx, tx = s[C].split(C.op.axis[0], factor=64)################################################################################# 最后，我們必須將迭代軸bx和tx綁定到GPU計(jì)算網(wǎng)格中的線程。# 樸素的 schedule 對(duì)GPU無(wú)效，這些是允許我們生成在GPU上運(yùn)行的代碼的特定構(gòu)造。s[C].bind(bx, te.thread_axis("blockIdx.x"))s[C].bind(tx, te.thread_axis("threadIdx.x"))####################################################################### 編譯# -----------# 在指定完 schdule 之后，我們可以將其編譯成一個(gè) TVM 函數(shù)。默認(rèn)情況下，TVM編譯成一個(gè) type-erased 函 # 數(shù)，可以從python端直接調(diào)用該函數(shù)。# 在下一行中，我們使用 tvm.build 來(lái)創(chuàng)建一個(gè)函數(shù)。build 函數(shù)采用 schedule、函數(shù)所需的簽名（包括輸如和輸出）以及我們要編譯到的目標(biāo)語(yǔ)言。# 編譯 fadd 的結(jié)果是一個(gè)GPU設(shè)備函數(shù)（如果涉及GPU）以及一個(gè)調(diào)用 GPU 函數(shù)的 host wrapper。fadd是生成的主機(jī)包裝函數(shù)，它在內(nèi)部包含對(duì)生成的設(shè)備函數(shù)的引用。fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")################################################################################# 編譯過(guò)的 TVM 函數(shù)會(huì)有一個(gè)簡(jiǎn)潔的 C API，它可以被任意的語(yǔ)言調(diào)用## 我們提供一個(gè) Python 的最小的數(shù)組 API 來(lái)幫助快速的測(cè)試和原型化# 該數(shù)組 API 是基于 `DLPack <https://github.com/dmlc/dlpack>`_ 標(biāo)準(zhǔn).## - 我們首先創(chuàng)建一個(gè) GPU 設(shè)備# - 然后 tvm.nd.array 從 GPU 拷貝數(shù)據(jù)# - ``fadd`` 運(yùn)行真正的計(jì)算# - ``numpy()`` 從 GPU 數(shù)組拷貝回 CPU (這樣我們就能驗(yàn)證正確性了).## 請(qǐng)注意，將數(shù)據(jù)復(fù)制到 GPU 上的內(nèi)存和從中復(fù)制數(shù)據(jù)是必需的步驟。dev = tvm.device(tgt_gpu.kind.name, 0)n = 1024a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)fadd(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())################################################################################# 檢查生成的 GPU 代碼# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~# 我們可以檢查在 TVM 中生成的代碼，tvm.build 的結(jié)果是一個(gè) TVM 模塊。fadd 是一個(gè) host 模塊其中包含 # host wrapper 的 host module，它同樣包含一個(gè)CUDA（GPU）設(shè)備模塊## 下面的代碼取得設(shè)備模塊并打印內(nèi)容代碼if (tgt_gpu.kind.name == "cuda"or tgt_gpu.kind.name == "rocm"or 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())

保存和加載編譯過(guò)的模塊

保存編譯過(guò)的模塊

除了運(yùn)行時(shí)編譯之外，我們還可以將編譯后的模塊保存到一個(gè)文件中，并在以后重新加載。下面的代碼執(zhí)行以下步驟：

• 它將編譯后的主機(jī)模塊保存到一個(gè)對(duì)象文件中。
• 然后將設(shè)備模塊保存到 ptx 文件中。
• cc.create_shared 調(diào)用編譯器（gcc）來(lái)創(chuàng)建共享庫(kù)

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

模塊存儲(chǔ)格式

CPU（Host）模塊直接保存為共享庫(kù)（.so）。設(shè)備代碼可以有多種自定義格式。在我們的示例中，設(shè)備代碼存儲(chǔ)在 ptx 中，元數(shù)據(jù)在 json 文件中。它們可以通過(guò)導(dǎo)入單獨(dú)加載和鏈接。

加載編譯過(guò)的模塊

我們可以從文件系統(tǒng)加載已編譯的模塊并運(yùn)行代碼。以下代碼分別加載主機(jī)和設(shè)備模塊，并將它們鏈接在一起。我們可以驗(yàn)證新加載的函數(shù)是否有效。

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())

將所有東西打包在一個(gè)庫(kù)中

在上面的示例中，我們分別存儲(chǔ)設(shè)備和主機(jī)代碼。TVM 還支持將所有內(nèi)容導(dǎo)出為一個(gè)共享庫(kù)。在 hood 下，我們將設(shè)備模塊打包成二進(jìn)制blob，并將它們與主機(jī)代碼鏈接在一起。目前我們支持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

運(yùn)行時(shí)API與線程安全

TVM 的編譯模塊并不依賴(lài)于 TVM 編譯器。它們只依賴(lài)于最小 Runtime Library。TVM Runtime Library 包裝設(shè)備驅(qū)動(dòng)程序，并向編譯函數(shù)提供線程安全和設(shè)備無(wú)關(guān)調(diào)用。

這意味著我們可以從任何GPU上的任何線程調(diào)用已編譯的TVM函數(shù)，前提是您已經(jīng)為該GPU編譯了代碼。

生成OpenCL代碼

TVM 為多種后端提供代碼生成功能。我們還可以生成在 CPU 后端上運(yùn)行的 OpenCL 代碼或 LLVM 代碼。

下面的代碼可以生成OpenCL代碼，在OpenCL設(shè)備上創(chuàng)建數(shù)組，并驗(yàn)證代碼的正確性。

if tgt.kind.name.startswith("opencl"):fadd_cl = tvm.build(s, [A, B, C], tgt, name="myadd")print("------opencl code------")print(fadd_cl.imported_modules[0].get_source())dev = tvm.cl(0)n = 1024a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)fadd_cl(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

注意：TE Scheduling Primitives

TE 調(diào)度原語(yǔ)

TVM 包括許多不同的調(diào)度原語(yǔ)：
split：按定義的因子將指定軸拆分為兩個(gè)軸。
tile：平鋪將按定義的因子沿兩個(gè)軸分割計(jì)算。
fuse：融合一次計(jì)算的兩個(gè)連續(xù)軸。
reorder：可以將計(jì)算軸重新排序?yàn)槎x的順序。
bind：可以將計(jì)算綁定到特定線程，在GPU編程中很有用。
compute_at：默認(rèn)情況下，TVM將在函數(shù)的最外層或根計(jì)算張量。compute_at指定應(yīng)在另一個(gè)運(yùn)算符的第一個(gè)計(jì)算軸上計(jì)算一個(gè)張量。
compute_inline：當(dāng)標(biāo)記為inline時(shí)，計(jì)算將展開(kāi)，然后插入到需要張量的地址中。
compute_root：將計(jì)算移動(dòng)到函數(shù)的最外層或根。這意味著，在進(jìn)入下一個(gè)階段之前，將對(duì)計(jì)算階段進(jìn)行完全計(jì)算。

可以在Schedule primitives 文檔頁(yè)面中找到這些原語(yǔ)的完整描述。

示例二：用TE手動(dòng)優(yōu)化矩陣乘

現(xiàn)在，我們將考慮第二個(gè)更高級(jí)一些的示例，演示如何用 18 行 Python 代碼 TVM 加速一個(gè)共同的矩陣乘法運(yùn)算 18倍。

矩陣乘法是一種計(jì)算密集型運(yùn)算。要獲得良好的CPU性能，有兩個(gè)重要的優(yōu)化：

• 提高內(nèi)存訪問(wèn)的緩存命中率。高緩存命中率可以加速?gòu)?fù)雜的數(shù)值計(jì)算和熱點(diǎn)內(nèi)存訪問(wèn)。這要求我們將源內(nèi)存訪問(wèn)模式轉(zhuǎn)換為適合緩存策略的模式。
• SIMD（單指令多數(shù)據(jù)），也稱(chēng)為矢量處理單元。在每個(gè)循環(huán)中，SIMD 都可以處理一小批數(shù)據(jù)，而不是處理單個(gè)值。這要求我們以統(tǒng)一模式轉(zhuǎn)換循環(huán)體中的數(shù)據(jù)訪問(wèn)模式，以便LLVM 后端可以將其 lower 到 SIMD。

本教程中使用的技術(shù)是這個(gè)倉(cāng)庫(kù)中提到的技巧的一部分。其中一些已被 TVM 抽象自動(dòng)使用，但由于 TVM 的一些約束，有一些無(wú)法自動(dòng)使用。

準(zhǔn)備工作和性能baseline

我們首先采集 numpy 實(shí)現(xiàn)的矩陣乘的數(shù)據(jù)：

import tvm import tvm.testing from tvm import te import numpy# 矩陣的尺寸： # (M, K) x (K, N) # 你可以自己試一些不同的尺寸，有時(shí)候 TVM 的優(yōu)化結(jié)果會(huì)好于含 MKL 的numpy M = 1024 K = 1024 N = 1024# tvm 中默認(rèn)的數(shù)據(jù)類(lèi)型 dtype = "float32"# 與之前一樣，這里可以根據(jù)自己的處理器及其是否支持某些指令集來(lái)改變 targettarget = tvm.target.Target(target="llvm", host="llvm") dev = tvm.device(target.kind.name, 0)# 隨機(jī)生成一些 tensor 用于測(cè)試 a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev) b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)# 重復(fù)實(shí)驗(yàn)，得到 numpy 的矩陣乘實(shí)現(xiàn)的 baseline np_repeat = 100 np_running_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_running_time / np_repeat))answer = numpy.dot(a.numpy(), b.numpy())

此處輸出：

Numpy running time: 0.009308

現(xiàn)在，我們用 TVM TE 編寫(xiě)一個(gè)基本矩陣乘法，并驗(yàn)證它產(chǎn)生的結(jié)果與numpy實(shí)現(xiàn)相同。我們還編寫(xiě)了一個(gè)函數(shù)，它將幫助我們度量進(jìn)度優(yōu)化的性能。

# 使用 TE 實(shí)現(xiàn)的 TVM 的矩陣乘 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 x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")# 默認(rèn) schedule s = te.create_schedule(C.op) func = tvm.build(s, [A, B, C], target=target, name="mmult")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)def evaluate_operation(s, vars, target, name, optimization, log):func = tvm.build(s, [A, B, C], target=target, name="mmult")assert funcc = 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)mean_time = evaluator(a, b, c).meanprint("%s: %f" % (optimization, mean_time))log.append((optimization, mean_time))log = []evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="none", log=log)

此處輸出：

none: 3.109406

讓我們看一下使用 TVM lower 函數(shù)的算子和默認(rèn) schedule 的中間表示 IR。請(qǐng)注意，該實(shí)現(xiàn)本質(zhì)上是矩陣乘法的簡(jiǎn)單實(shí)現(xiàn)，在 A 和 B 矩陣的索引上使用三個(gè)嵌套循環(huán)。

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 (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)]))}}} }

優(yōu)化1：blocking阻塞

提高緩存命中率的一個(gè)重要技巧是阻塞，在這種情況下，我們可以構(gòu)造內(nèi)存訪問(wèn)，使塊內(nèi)部是具有高內(nèi)存局部性的小鄰域。在本教程中，我們選擇塊因子 32。這會(huì)使得一個(gè)塊填充內(nèi)存的 32*32*sizeof（float）區(qū)域。這對(duì)應(yīng)于 4KB 的緩存大小，和一級(jí)緩存 32KB 的參考緩存大小。

我們首先為 C 操作創(chuàng)建一個(gè)默認(rèn)的調(diào)度，然后使用指定的塊因子對(duì)其應(yīng)用一個(gè) tile 調(diào)度原語(yǔ)，調(diào)度原語(yǔ)以向量 [x_-outer，y_-outer，x_-inner，y_-inner] 的形式返回從最外層到最內(nèi)層的循環(huán)順序。然后，我們得到操作輸出的縮減軸，并使用因子4對(duì)其執(zhí)行拆分操作。這個(gè)因素不會(huì)直接影響我們現(xiàn)在正在進(jìn)行的阻塞優(yōu)化，但在以后應(yīng)用矢量化時(shí)會(huì)很有用。
現(xiàn)在操作已被阻塞，我們可以對(duì)計(jì)算進(jìn)行重新排序，將簡(jiǎn)化操作放入計(jì)算的最外層循環(huán)中，從而幫助確保被阻塞的數(shù)據(jù)仍保留在緩存中。這就完成了 schedule，我們可以構(gòu)建和測(cè)試與原始 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, factor=4)# Hoist reduction domain outside the blocking loop s[C].reorder(xo, yo, ko, ki, xi, yi)evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="blocking", log=log)

此處輸出：

blocking: 0.291928

通過(guò)重新排序計(jì)算以利用緩存，我們可以看到計(jì)算性能的顯著提高。現(xiàn)在，打印內(nèi)部表示并將其與原始表示進(jìn)行比較：

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} {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)]))}}}}}} }

優(yōu)化2： vectorization矢量化

另一個(gè)重要的優(yōu)化技巧是矢量化。當(dāng)內(nèi)存訪問(wèn)模式一致時(shí)，編譯器可以檢測(cè)到該模式并將連續(xù)內(nèi)存?zhèn)鬟f給 SIMD 向量處理器。在TVM中，我們可以利用這個(gè)硬件特性，使用矢量化接口來(lái)提示編譯器這個(gè)模式。

在本教程中，我們選擇對(duì)內(nèi)部循環(huán)行數(shù)據(jù)進(jìn)行矢量化，因?yàn)樗呀?jīng)是我們之前優(yōu)化中的緩存友好型數(shù)據(jù)。

# 應(yīng)用矢量化的優(yōu)化方式 s[C].vectorize(yi)evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="vectorization", log=log)# 矢量化之后生成的 IR print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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)]))}}}}} }

優(yōu)化3：Loop Permutation循環(huán)置換

如果我們看一下上面的 IR，我們可以看到內(nèi)環(huán)行數(shù)據(jù)被矢量化，B 被轉(zhuǎn)換成 PackedB（這在內(nèi)環(huán)的（float32x32）B2部分中很明顯）。PackedB 的遍歷現(xiàn)在是順序的。因此，我們將研究 A 的訪問(wèn)模式。在當(dāng)前 schdule中，A 是逐列訪問(wèn)的，這對(duì)緩存不友好。如果我們改變嵌套循環(huán)順序 ki 和內(nèi)部軸 xi，對(duì) A 的訪問(wèn)模式將變得更加緩存友好。

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, factor=4)# re-ordering s[C].reorder(xo, yo, ko, xi, ki, yi) s[C].vectorize(yi)evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="loop permutation", log=log )# 再一次打印新生成的 IR print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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)]))}}}}} }

優(yōu)化4：Array Packing數(shù)組打包

另一個(gè)重要技巧是數(shù)組打包。此技巧是對(duì)陣列的存儲(chǔ)維度重新排序，以便在展平后將特定維度上的連續(xù)訪問(wèn)模式轉(zhuǎn)換為序列模式。

如上圖所示，在阻塞計(jì)算后，我們可以觀察到 B 的陣列訪問(wèn)模式（平坦后），它是規(guī)則的但不連續(xù)的。我們希望經(jīng)過(guò)一些轉(zhuǎn)換后，我們可以得到一個(gè)連續(xù)的訪問(wèn)模式。通過(guò)將[16][16]數(shù)組重新排序?yàn)閇16/4][16][4]數(shù)組，在從壓縮數(shù)組中獲取相應(yīng)值時(shí)，B 的訪問(wèn)模式將是順序的。

為了實(shí)現(xiàn)這一點(diǎn)，我們必須從一個(gè)新的默認(rèn) schedule 開(kāi)始，考慮到 B 的新 wrapper。花點(diǎn)時(shí)間對(duì)此進(jìn)行討論是值得的：TE 是一種用于編寫(xiě)優(yōu)化算子的功能強(qiáng)大的表達(dá)性語(yǔ)言，但它通常需要一些底層算法、數(shù)據(jù)結(jié)構(gòu)，以及您正在編寫(xiě)的硬件 target。在本教程的后面，我們將討論讓 TVM 承擔(dān)這一負(fù)擔(dān)的一些選擇。不管怎樣，讓我們繼續(xù)新的優(yōu)化 schedule。

# 我們要輕微地重寫(xiě)算法 packedB = te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn + z], name="packedB") C = te.compute((M, N),lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axis=k),name="C", )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, factor=4)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], target=target, name="mmult", optimization="array packing", log=log)# 這里是數(shù)組打包之后生成的 IR print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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)]))}}}}}} }

優(yōu)化5：Optimizing Block Writing Through Caching通過(guò)緩存優(yōu)化塊寫(xiě)入

到目前為止，我們所有的優(yōu)化都集中在高效地訪問(wèn)和計(jì)算來(lái)自 A 和 B 矩陣的數(shù)據(jù)，以計(jì)算C矩陣。阻塞優(yōu)化后，操作員將結(jié)果逐塊寫(xiě)入 C，并且訪問(wèn)模式不是順序的。我們可以通過(guò)使用順序緩存數(shù)組來(lái)解決這個(gè)問(wèn)題，使用cache_write、compute_at 和 unroll 的組合來(lái)保存塊結(jié)果，并在所有塊結(jié)果就緒時(shí)寫(xiě)入到 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, factor=4) 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], target=target, name="mmult", optimization="block caching", log=log)# Here is the generated IR after write cache blocking. print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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)]}}}}} }

優(yōu)化6：Parallelization并行化

# 并行 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], target=target, name="mmult", optimization="parallelization", log=log )# 這里是并行化之后的 IR print(tvm.lower(s, [A, B, C], simple_mode=True))

此處輸出：

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)]}}}}} }

矩陣乘例子的總結(jié)

在僅用 18 行代碼應(yīng)用上述簡(jiǎn)單優(yōu)化之后，我們生成的代碼就可以得到與使用數(shù)學(xué)內(nèi)核庫(kù)（MKL）的 numpy 接近的性能。我們剛才一直都記錄了性能，因此在這里可以直接比較結(jié)果：

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

請(qǐng)注意，以上的輸出反映的是非獨(dú)占 Docker 容器上的運(yùn)行時(shí)間，因此并不可靠。強(qiáng)烈建議您自己運(yùn)行本教程，觀察 TVM 實(shí)現(xiàn)的性能增益，并仔細(xì)閱讀每個(gè)示例，以了解矩陣乘法運(yùn)算的迭代改進(jìn)。

總結(jié)

如前所述，如何使用 TE 和調(diào)度原語(yǔ)應(yīng)用優(yōu)化可能需要一些底層架構(gòu)和算法的知識(shí)。然而，TE 設(shè)計(jì)為更復(fù)雜的算法是為了可以搜索潛在的優(yōu)化。有了本 TE 簡(jiǎn)介中的知識(shí)，我們現(xiàn)在可以開(kāi)始探索 TVM 如何自動(dòng)化進(jìn)度優(yōu)化過(guò)程。

本教程提供了使用向量加法和矩陣乘法示例的TVM張量表達(dá)式（TE）工作流演練。一般的工作流程是：

• 通過(guò)一系列操作描述您的計(jì)算。

• 描述我們希望如何計(jì)算和使用調(diào)度原語(yǔ)。

• 編譯到我們想要的目標(biāo)函數(shù)。

• 保存要稍后加載的函數(shù)（可選）。

接下來(lái)的教程將擴(kuò)展矩陣乘法示例，并展示如何使用可調(diào)參數(shù)構(gòu)建矩陣乘法和其他操作的通用模板，這些參數(shù)使得我們能夠自動(dòng)優(yōu)化特定平臺(tái)的計(jì)算。

Ref：

https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html

總結(jié)

以上是生活随笔為你收集整理的TVM：使用Tensor Expression (TE)来处理算子的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

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