TVM：使用Tensor Expression (TE)来处理算子
TVM：使用Tensor Expression (TE)來處理算子
在本教程中，我們將聚焦于在 TVM 中使用張量表達式（TE）來定義張量計算和實現(xiàn)循環(huán)優(yōu)化。TE用純函數(shù)語言描述張量計算（即每個表達式都沒有副作用）。當(dāng)在 TVM 的整體上下文中查看時，Relay 將計算描述為一組算子，并且其中每一個算子都可以表示為 TE 表達式，每個 TE 表達式獲取輸入張量并生成輸出張量。
本文是TVM中 TE 語言的入門教程。TVM 使用領(lǐng)域?qū)Ｓ?#xff08;domain specific）的張量表達式來高效地構(gòu)造內(nèi)核。我們以兩個使用 TE 語言的為例來演示基本工作流。第一個示例介紹了 TE 和帶有向量加法的 schedule。第二個示例通過逐步優(yōu)化矩陣與 TE 的乘法來擴展這些概念。這個矩陣乘法示例將作為未來涵蓋更高級的 TVM 特性的教程的對比基礎(chǔ)。
示例一：使用TE為CPU編寫和調(diào)度向量加法
初始化 tvm環(huán)境
我們的第一個例子是使用 Python 來為向量加法實現(xiàn)一個 TE，然后是一個針對 CPU 的 schedule，我們從初始化 tvm 環(huán)境開始：
import tvm import tvm.testing from tvm import te import numpy as np# 如果能夠指定目標(biāo) CPU，那么將會得到更好地性能 # 如果用的是llvm，可以通過 `llc --version` 來查看 CPU 類型 # 可以通過查看 /proc/cpuinfo 來查看你的處理器可能支持的其他擴展， # 比如，如果你的 CPU 有 AVX-512 指令集，那么你可以使用 `llvm -mcpu=skylake-avx512` 選項tgt = tvm.target.Target(target="llvm", host="llvm")描述向量計算
我們首先描述向量加法計算。TVM 采用張量語義，每個中間結(jié)果表示為一個多維數(shù)組。我們需要描述規(guī)則來得到張量。我們首先定義一個符號變量 n 來表示形狀。然后我們定義兩個 placeholder 張量：A、B，它們的形狀是 (n,)。然后我們通過一個 compute 操作，得到結(jié)果張量 C。compute 定義了一種計算，其輸出符合指定的張量形狀，并在由 lambda 函數(shù)定義的張量中的每個位置執(zhí)行計算。注意，雖然 n 是一個變量，但它定義了A、B 和 C 張量之間的一致形狀。請注意，在這個階段沒有實際的計算發(fā)生，因為我們只是聲明應(yīng)該如何進行計算。
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方法的第二個參數(shù)是執(zhí)行計算的函數(shù)。在本例中，我們使用一個匿名函數(shù)（也稱為lambda函數(shù)）來定義計算，在本例中是對 a 和 B 的第 i 個元素的加法。
為計算創(chuàng)建一個默認的Schedule
雖然上面幾行描述了計算規(guī)則，但我們可以用許多不同的方法計算 C 以適應(yīng)不同的設(shè)備。對于具有多個 axis 的張量，您可以選擇首先迭代哪個 axis ，另外計算可以跨不同的線程拆分。TVM要求用戶提供一個 schedule，來描述應(yīng)如何執(zhí)行計算。TE 中的 schedule 操作可以更改循環(huán)順序、跨不同線程拆分計算、將數(shù)據(jù)塊分組在一起，以及其他操作。schedule 背后的一個重要概念是，它們只描述如何執(zhí)行計算，因此相同 TE 的不同 schedule 一定會產(chǎn)生相同的結(jié)果。
在 TVM 中，我們可以創(chuàng)建一種樸素的 schedule ，按照行優(yōu)先的順序來計算 C。
for (int i = 0; i < n; ++i) {C[i] = A[i] + B[i]; } s = te.create_schedule(C.op)編譯并驗證默認的 schedule
通過 TE 表達式和 schedule，我們可以為目標(biāo)語言和體系結(jié)構(gòu)生成可運行的代碼，在本例中是 LLVM 和 CPU 。我們向 TVM 提供 schedule、schedule 中的TE表達式列表、目標(biāo)和主機，以及我們正在生成的函數(shù)的名稱。輸出的結(jié)果是可以直接從 Python 調(diào)用 type-erased 函數(shù)。
在下一行中，我們使用 tvm.build 創(chuàng)建一個函數(shù)。build 函數(shù)接受 schedule、函數(shù)所需的簽名（包括輸入和輸出）以及我們要編譯到的目標(biāo)語言。
fadd = tvm.build(s, [A, B, C], tgt, name="myadd")我們運行該函數(shù)，并將輸出與 numpy 中的相同計算進行比較。編譯后的 TVM 函數(shù)提供了一個簡明的C API，可以被任何語言調(diào)用。我們首先創(chuàng)建一個設(shè)備（在本例中為CPU），這是一個 TVM 可以編譯 schedule 的設(shè)備。在本例中，設(shè)備是LLVM CPU target。然后，我們可以在設(shè)備中初始化張量并執(zhí)行自定義的加法操作。為了驗證計算的正確性，我們可以將c張量的輸出結(jié)果與 numpy 執(zhí)行的相同計算進行比較。
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())為了對比這個樸素版本的自定義向量加法與 numpy 的速度差異，創(chuàng)建一個輔助函數(shù)來運行 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）來優(yōu)化 schedule
我們已經(jīng)說明了 TE 的基本原理，現(xiàn)在讓我們更深入地了解 schedule 的作用，以及它們?nèi)绾斡糜趦?yōu)化不同體系結(jié)構(gòu)的張量表達式。schedule 是應(yīng)用于表達式的一系列步驟，用于以多種不同方式對其進行轉(zhuǎn)換。當(dāng)一個 schedule 應(yīng)用于TE中的一個表達式時，輸入和輸出保持不變，但在編譯時，表達式的實現(xiàn)可能會改變。在默認 schedule 中，這個張量加法是串行運行的，但該操作其實是很容易在所有處理器線程之間并行。我們可以將我們的操作并行調(diào)度到計算中：
s[C].parallel(C.op.axis[0])tvm.lower 命令將生成 TE 的中間表示（IR）以及相應(yīng)的 schedule 。通過在執(zhí)行不同的 schedule 操作時 lowing 表達式，我們可以看到 schedule 對計算順序的影響。我們使用標(biāo)志 simple_mode=True 返回可讀的 C 風(fēng)格語句。
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)在可以在獨立的線程上運行這些塊。我們在執(zhí)行并行操作的情況下編譯并運行這個新的 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）來優(yōu)化 schedule
現(xiàn)代 CPU 能夠?qū)Ω↑c數(shù)進行 SIMD 操作，我們可以對計算表達式使用另一個 schedule 來利用這一點。實現(xiàn)這一點需要多個步驟：首先，我們必須使用 split scheduling 原語將 schedule 拆分為內(nèi)部循環(huán)和外部循環(huán)。內(nèi)部循環(huán)可以使用向量化來使用使用向量化調(diào)度原語的 SIMD 指令，然后外部循環(huán)可以使用并行調(diào)度原語進行并行化。選擇分割因子作為CPU上的線程數(shù)。
注：SIMD，全稱 Single Instruction Multiple Data，單指令多數(shù)據(jù)流，能夠復(fù)制多個操作數(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)])}}} }對比不同的 schedule
下面我們來對比以下之前提到的不同 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
代碼專門化
正如我們所看到的，A、B 和 C 的聲明都采用相同的形狀參數(shù) n。TVM將利用這一點，只向 kernel 傳遞一個 shape 參數(shù)，我們在打印的設(shè)備代碼中找到它。這是專門化化的一種形式。
在 host 端，TVM 將自動生成檢查代碼，以檢查參數(shù)中的約束。因此，如果將具有不同形狀的數(shù)組傳遞到 fadd 中，將引發(fā)錯誤。
我們可以做更多的專門化。例如，我們可以在計算聲明中寫入n=tvm.runtime.convert（1024）而不是 n=te.var（“n”）。生成的函數(shù)將只獲取長度為1024的向量。
我們已經(jīng)定義、調(diào)度并編譯了一個向量加法運算符，然后可以在 TVM Runtime 執(zhí)行它。我們可以將算子保存為庫，稍后可以使用 TVM Runtime 加載該庫。
針對GPU的矩陣加法（可選）
在介紹保存與加載自定義算子庫的方法之前，我們先來看一下如何針對 GPU 做矩陣加法。
TVM能夠針對多種體系結(jié)構(gòu)。在本例，我們將針對GPU中矢量加法的編譯。
# 本段代碼默認不運行，如果想要運行的話，請將 ``run_cuda = True``run_cuda = False if run_cuda:# 這里的 target 需要根據(jù)自己的 GPU 類型修改：# 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計算網(wǎng)格中的線程。# 樸素的 schedule 對GPU無效，這些是允許我們生成在GPU上運行的代碼的特定構(gòu)造。s[C].bind(bx, te.thread_axis("blockIdx.x"))s[C].bind(tx, te.thread_axis("threadIdx.x"))####################################################################### 編譯# -----------# 在指定完 schdule 之后，我們可以將其編譯成一個 TVM 函數(shù)。默認情況下，TVM編譯成一個 type-erased 函 # 數(shù)，可以從python端直接調(diào)用該函數(shù)。# 在下一行中，我們使用 tvm.build 來創(chuàng)建一個函數(shù)。build 函數(shù)采用 schedule、函數(shù)所需的簽名（包括輸如和輸出）以及我們要編譯到的目標(biāo)語言。# 編譯 fadd 的結(jié)果是一個GPU設(shè)備函數(shù)（如果涉及GPU）以及一個調(diào)用 GPU 函數(shù)的 host wrapper。fadd是生成的主機包裝函數(shù)，它在內(nèi)部包含對生成的設(shè)備函數(shù)的引用。fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")################################################################################# 編譯過的 TVM 函數(shù)會有一個簡潔的 C API，它可以被任意的語言調(diào)用## 我們提供一個 Python 的最小的數(shù)組 API 來幫助快速的測試和原型化# 該數(shù)組 API 是基于 `DLPack <https://github.com/dmlc/dlpack>`_ 標(biāo)準(zhǔn).## - 我們首先創(chuàng)建一個 GPU 設(shè)備# - 然后 tvm.nd.array 從 GPU 拷貝數(shù)據(jù)# - ``fadd`` 運行真正的計算# - ``numpy()`` 從 GPU 數(shù)組拷貝回 CPU (這樣我們就能驗證正確性了).## 請注意，將數(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é)果是一個 TVM 模塊。fadd 是一個 host 模塊其中包含 # host wrapper 的 host module，它同樣包含一個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())保存和加載編譯過的模塊
保存編譯過的模塊
除了運行時編譯之外，我們還可以將編譯后的模塊保存到一個文件中，并在以后重新加載。下面的代碼執(zhí)行以下步驟：
- 它將編譯后的主機模塊保存到一個對象文件中。
- 然后將設(shè)備模塊保存到 ptx 文件中。
- cc.create_shared 調(diào)用編譯器（gcc）來創(chuàng)建共享庫
此處輸出：
['myadd.o', 'myadd.so']注意：Module Storage Format
模塊存儲格式
CPU（Host）模塊直接保存為共享庫（.so）。設(shè)備代碼可以有多種自定義格式。在我們的示例中，設(shè)備代碼存儲在 ptx 中，元數(shù)據(jù)在 json 文件中。它們可以通過導(dǎo)入單獨加載和鏈接。
加載編譯過的模塊
我們可以從文件系統(tǒng)加載已編譯的模塊并運行代碼。以下代碼分別加載主機和設(shè)備模塊，并將它們鏈接在一起。我們可以驗證新加載的函數(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())將所有東西打包在一個庫中
在上面的示例中，我們分別存儲設(shè)備和主機代碼。TVM 還支持將所有內(nèi)容導(dǎo)出為一個共享庫。在 hood 下，我們將設(shè)備模塊打包成二進制blob，并將它們與主機代碼鏈接在一起。目前我們支持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
運行時API與線程安全
TVM 的編譯模塊并不依賴于 TVM 編譯器。它們只依賴于最小 Runtime Library。TVM Runtime Library 包裝設(shè)備驅(qū)動程序，并向編譯函數(shù)提供線程安全和設(shè)備無關(guān)調(diào)用。
這意味著我們可以從任何GPU上的任何線程調(diào)用已編譯的TVM函數(shù)，前提是您已經(jīng)為該GPU編譯了代碼。
生成OpenCL代碼
TVM 為多種后端提供代碼生成功能。我們還可以生成在 CPU 后端上運行的 OpenCL 代碼或 LLVM 代碼。
下面的代碼可以生成OpenCL代碼，在OpenCL設(shè)備上創(chuàng)建數(shù)組，并驗證代碼的正確性。
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)度原語
TVM 包括許多不同的調(diào)度原語：
split：按定義的因子將指定軸拆分為兩個軸。
tile：平鋪將按定義的因子沿兩個軸分割計算。
fuse：融合一次計算的兩個連續(xù)軸。
reorder：可以將計算軸重新排序為定義的順序。
bind：可以將計算綁定到特定線程，在GPU編程中很有用。
compute_at：默認情況下，TVM將在函數(shù)的最外層或根計算張量。compute_at指定應(yīng)在另一個運算符的第一個計算軸上計算一個張量。
compute_inline：當(dāng)標(biāo)記為inline時，計算將展開，然后插入到需要張量的地址中。
compute_root：將計算移動到函數(shù)的最外層或根。這意味著，在進入下一個階段之前，將對計算階段進行完全計算。
可以在Schedule primitives 文檔頁面中找到這些原語的完整描述。
示例二：用TE手動優(yōu)化矩陣乘
現(xiàn)在，我們將考慮第二個更高級一些的示例，演示如何用 18 行 Python 代碼 TVM 加速一個共同的矩陣乘法運算 18倍。
矩陣乘法是一種計算密集型運算。要獲得良好的CPU性能，有兩個重要的優(yōu)化：
- 提高內(nèi)存訪問的緩存命中率。高緩存命中率可以加速復(fù)雜的數(shù)值計算和熱點內(nèi)存訪問。這要求我們將源內(nèi)存訪問模式轉(zhuǎn)換為適合緩存策略的模式。
- SIMD（單指令多數(shù)據(jù)），也稱為矢量處理單元。在每個循環(huán)中，SIMD 都可以處理一小批數(shù)據(jù)，而不是處理單個值。這要求我們以統(tǒng)一模式轉(zhuǎn)換循環(huán)體中的數(shù)據(jù)訪問模式，以便LLVM 后端可以將其 lower 到 SIMD。
本教程中使用的技術(shù)是這個倉庫中提到的技巧的一部分。其中一些已被 TVM 抽象自動使用，但由于 TVM 的一些約束，有一些無法自動使用。
準(zhǔn)備工作和性能baseline
我們首先采集 numpy 實現(xiàn)的矩陣乘的數(shù)據(jù)：
import tvm import tvm.testing from tvm import te import numpy# 矩陣的尺寸： # (M, K) x (K, N) # 你可以自己試一些不同的尺寸，有時候 TVM 的優(yōu)化結(jié)果會好于含 MKL 的numpy M = 1024 K = 1024 N = 1024# tvm 中默認的數(shù)據(jù)類型 dtype = "float32"# 與之前一樣，這里可以根據(jù)自己的處理器及其是否支持某些指令集來改變 targettarget = tvm.target.Target(target="llvm", host="llvm") dev = tvm.device(target.kind.name, 0)# 隨機生成一些 tensor 用于測試 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ù)實驗，得到 numpy 的矩陣乘實現(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 編寫一個基本矩陣乘法，并驗證它產(chǎn)生的結(jié)果與numpy實現(xiàn)相同。我們還編寫了一個函數(shù)，它將幫助我們度量進度優(yōu)化的性能。
# 使用 TE 實現(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")# 默認 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ù)的算子和默認 schedule 的中間表示 IR。請注意，該實現(xiàn)本質(zhì)上是矩陣乘法的簡單實現(xiàn)，在 A 和 B 矩陣的索引上使用三個嵌套循環(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òu)造內(nèi)存訪問，使塊內(nèi)部是具有高內(nèi)存局部性的小鄰域。在本教程中，我們選擇塊因子 32。這會使得一個塊填充內(nèi)存的 32*32*sizeof（float）區(qū)域。這對應(yīng)于 4KB 的緩存大小，和一級緩存 32KB 的參考緩存大小。
我們首先為 C 操作創(chuàng)建一個默認的調(diào)度，然后使用指定的塊因子對其應(yīng)用一個 tile 調(diào)度原語，調(diào)度原語以向量 [x_-outer，y_-outer，x_-inner，y_-inner] 的形式返回從最外層到最內(nèi)層的循環(huán)順序。然后，我們得到操作輸出的縮減軸，并使用因子4對其執(zhí)行拆分操作。這個因素不會直接影響我們現(xiàn)在正在進行的阻塞優(yōu)化，但在以后應(yīng)用矢量化時會很有用。
現(xiàn)在操作已被阻塞，我們可以對計算進行重新排序，將簡化操作放入計算的最外層循環(huán)中，從而幫助確保被阻塞的數(shù)據(jù)仍保留在緩存中。這就完成了 schedule，我們可以構(gòu)建和測試與原始 schedule 相比的性能。
此處輸出：
blocking: 0.291928通過重新排序計算以利用緩存，我們可以看到計算性能的顯著提高。現(xiàn)在，打印內(nèi)部表示并將其與原始表示進行比較：
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矢量化
另一個重要的優(yōu)化技巧是矢量化。當(dāng)內(nèi)存訪問模式一致時，編譯器可以檢測到該模式并將連續(xù)內(nèi)存?zhèn)鬟f給 SIMD 向量處理器。在TVM中，我們可以利用這個硬件特性，使用矢量化接口來提示編譯器這個模式。
在本教程中，我們選擇對內(nèi)部循環(huán)行數(shù)據(jù)進行矢量化，因為它已經(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 的訪問模式。在當(dāng)前 schdule中，A 是逐列訪問的，這對緩存不友好。如果我們改變嵌套循環(huán)順序 ki 和內(nèi)部軸 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, 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ù)組打包
另一個重要技巧是數(shù)組打包。此技巧是對陣列的存儲維度重新排序，以便在展平后將特定維度上的連續(xù)訪問模式轉(zhuǎn)換為序列模式。
如上圖所示，在阻塞計算后，我們可以觀察到 B 的陣列訪問模式（平坦后），它是規(guī)則的但不連續(xù)的。我們希望經(jīng)過一些轉(zhuǎn)換后，我們可以得到一個連續(xù)的訪問模式。通過將[16][16]數(shù)組重新排序為[16/4][16][4]數(shù)組，在從壓縮數(shù)組中獲取相應(yīng)值時，B 的訪問模式將是順序的。
為了實現(xiàn)這一點，我們必須從一個新的默認 schedule 開始，考慮到 B 的新 wrapper?；c時間對此進行討論是值得的：TE 是一種用于編寫優(yōu)化算子的功能強大的表達性語言，但它通常需要一些底層算法、數(shù)據(jù)結(jié)構(gòu)，以及您正在編寫的硬件 target。在本教程的后面，我們將討論讓 TVM 承擔(dān)這一負擔(dān)的一些選擇。不管怎樣，讓我們繼續(xù)新的優(yōu)化 schedule。
# 我們要輕微地重寫算法 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通過緩存優(yōu)化塊寫入
到目前為止，我們所有的優(yōu)化都集中在高效地訪問和計算來自 A 和 B 矩陣的數(shù)據(jù)，以計算C矩陣。阻塞優(yōu)化后，操作員將結(jié)果逐塊寫入 C，并且訪問模式不是順序的。我們可以通過使用順序緩存數(shù)組來解決這個問題，使用cache_write、compute_at 和 unroll 的組合來保存塊結(jié)果，并在所有塊結(jié)果就緒時寫入到 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)用上述簡單優(yōu)化之后，我們生成的代碼就可以得到與使用數(shù)學(xué)內(nèi)核庫（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請注意，以上的輸出反映的是非獨占 Docker 容器上的運行時間，因此并不可靠。強烈建議您自己運行本教程，觀察 TVM 實現(xiàn)的性能增益，并仔細閱讀每個示例，以了解矩陣乘法運算的迭代改進。
總結(jié)
如前所述，如何使用 TE 和調(diào)度原語應(yīng)用優(yōu)化可能需要一些底層架構(gòu)和算法的知識。然而，TE 設(shè)計為更復(fù)雜的算法是為了可以搜索潛在的優(yōu)化。有了本 TE 簡介中的知識，我們現(xiàn)在可以開始探索 TVM 如何自動化進度優(yōu)化過程。
本教程提供了使用向量加法和矩陣乘法示例的TVM張量表達式（TE）工作流演練。一般的工作流程是：
-
通過一系列操作描述您的計算。
-
描述我們希望如何計算和使用調(diào)度原語。
-
編譯到我們想要的目標(biāo)函數(shù)。
-
保存要稍后加載的函數(shù)（可選）。
接下來的教程將擴展矩陣乘法示例，并展示如何使用可調(diào)參數(shù)構(gòu)建矩陣乘法和其他操作的通用模板，這些參數(shù)使得我們能夠自動優(yōu)化特定平臺的計算。
Ref：
https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html
總結(jié)
以上是生活随笔為你收集整理的TVM：使用Tensor Expression (TE)来处理算子的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 信用卡能转账到微信吗
- 下一篇： gcc参数 -i, -L, -l, -i