使用Tensor Expression張量表達式處理算子
這是TVM中Tensor表達語言的入門教程。TVM使用特定于域的張量表達式來進行有效的內核構造。

本文將演示使用張量表達式語言的基本工作流程。
from future import absolute_import, print_function
import tvm
import numpy as np

Global declarations of environment.

tgt_host=“llvm”

Change it to respective GPU if gpu is enabled Ex: cuda, opencl

tgt=“cuda”
Vector Add Example
將使用向量添加示例來演示工作流程。
描述計算
作為第一步，需要描述計算。TVM采用Tensor語義，每個中間結果表示為多維數組。用戶需要描述生成Tensor的計算規則。
先定義一個符號變量n表示形狀。再定義兩個占位符Tensor A和B，給定形狀（n，）
然后用計算操作描述結果Tensor C. 計算函數采用張量的形狀，以及描述張量的每個位置的計算規則的lambda函數。
在此階段沒有計算，因為只是聲明應該如何進行計算。
n = tvm.var(“n”)
A = tvm.placeholder((n,), name=‘A’)
B = tvm.placeholder((n,), name=‘B’)
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name=“C”)
print(type?)
輸出：
<class ‘tvm.tensor.Tensor’>
計算圖 Schedule the Computation
雖然上面的行描述了計算規則，但可以以多種方式計算C，因為C軸可以以數據并行方式計算。TVM要求用戶提供的計算描述，稱為一個schedule。
schedule是一組計算轉換，轉換程序中的計算循環。
在構造schedule后，默認情況下，schedule以行主要順序以串行方式計算C.
for (int i = 0; i < n; ++i) {
C[i] = A[i] + B[i];
}
s = tvm.create_schedule(C.op)
使用split構造，分割C的第一個軸，這將原始迭代軸分成兩次迭代的乘積。這相當于以下代碼:
for (int bx = 0; bx < ceil(n / 64); ++bx) {
for (int tx = 0; tx < 64; ++tx) {
int i = bx * 64 + tx;
if (i < n) {
C[i] = A[i] + B[i];
}
}
}
bx, tx = s[C].split(C.op.axis[0], factor=64)
最后，將迭代軸bx和tx綁定到GPU計算網格中的線程。這些是GPU特定的構造，允許生成在GPU上運行的代碼。
if tgt == “cuda” or tgt.startswith(‘opencl’):
s[C].bind(bx, tvm.thread_axis(“blockIdx.x”))
s[C].bind(tx, tvm.thread_axis(“threadIdx.x”))
Compilation
在完成指定Schedule后，可以編譯為TVM函數。默認情況下，TVM編譯成一個與類型無關的函數，可以從python端直接調用。
在下面的行中，使用tvm.build創建一個函數。構建函數采用調度，函數的期望簽名（包括輸入和輸出）以及要編譯的目標語言。
編譯fadd的結果是GPU設備功能（如果涉及GPU）以及調用GPU功能的主機包裝器。fadd是生成的主機包裝器函數，包含對內部生成的設備函數的引用。
fadd = tvm.build(s, [A, B, C], tgt, target_host=tgt_host, name=“myadd”)
運行功能
編譯的TVM函數公開了一個可以從任何語言調用的簡潔C API。
在python中提供了一個最小的數組API，幫助快速測試和原型設計。陣列API基于DLPack標準。
? 首先創建一個GPU上下文。
? 然后tvm.nd.array將數據復制到GPU。
? fadd運行實際計算。
? asnumpy（）將GPU陣列復制回CPU，可以驗證正確性
ctx = tvm.context(tgt, 0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
檢查生成的代碼
您可以在TVM中檢查生成的代碼。tvm.build的結果是一個TVM模塊。fadd是包含主機包裝器的主機模塊，它還包含用于CUDA（GPU）功能的設備模塊。
以下代碼獲取設備模塊并打印內容代碼。
if tgt == “cuda” or tgt.startswith(‘opencl’):
dev_module = fadd.imported_modules[0]
print("-----GPU code-----")
print(dev_module.get_source())
else:
print(fadd.get_source())
日期：
-----GPU code-----
extern “C” global void myadd_kernel0( float* restrict C, float* restrict A, float* restrict B, int n) {
if ((((int)blockIdx.x) < (((n - 64) / 64) + 1)) && (((int)blockIdx.x) < (((n + 63) / 64) - 1))) {
C[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = (A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] + B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
C[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = (A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] + B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
}
}
}
}
注意
代碼定制化
可能已經注意到，A，B和C的聲明都采用相同的形狀參數，n。正如將在打印的設備代碼中找到的那樣，TVM將利用此功能僅將單個形狀參數傳遞給內核。這是一種專業化形式。
在主機端，TVM將自動生成檢查參數中的約束的檢查代碼。因此，如果將具有不同形狀的數組傳遞給fadd，會引發錯誤。
可以做更多的專業化。例如，可以在計算聲明中編寫 n = tvm.convert(1024)，而不是n = tvm.var(“n”)，使生成的函數僅采用長度為1024的向量。
保存編譯模塊
除了運行時編譯外，可以將已編譯的模塊保存到文件中，在以后加載。稱為提前編譯。
以下代碼首先執行以下步驟：
? 將已編譯的主機模塊保存到目標文件中。
? 然后將設備模塊保存到ptx文件中。
? cc.create_shared調用編譯器（gcc），創建共享庫
from tvm.contrib import cc
from tvm.contrib import util

temp = util.tempdir()
fadd.save(temp.relpath(“myadd.o”))
if tgt == “cuda”:
fadd.imported_modules[0].save(temp.relpath(“myadd.ptx”))
if tgt.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.so’, ‘myadd.ptx’, ‘myadd.o’, ‘myadd.tvm_meta.json’]
注意
模塊存儲格式
CPU（主機）模塊直接保存為共享庫（.so）。可以有多種自定義格式的設備代碼。在示例中，設備代碼存儲在ptx中，以及元數據json文件中。可以通過導入實現單獨加載和鏈接。
加載編譯模塊
可以從文件系統加載已編譯的模塊并運行代碼。以下代碼分別加載主機和設備模塊，并將重新鏈接在一起。可以驗證新加載的函數是否有效。
fadd1 = tvm.module.load(temp.relpath(“myadd.so”))
if tgt == “cuda”:
fadd1_dev = tvm.module.load(temp.relpath(“myadd.ptx”))
fadd1.import_module(fadd1_dev)

if tgt.startswith(‘opencl’):
fadd1_dev = tvm.module.load(temp.relpath(“myadd.cl”))
fadd1.import_module(fadd1_dev)

fadd1(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
將所有內容打包到一個庫中
在上面的示例中，分別存儲設備和主機代碼。TVM還支持將所有內容導出為一個共享庫。將設備模塊打包成二進制blob，將與主機代碼鏈接在一起。目前，支持Metal，OpenCL和CUDA模塊的包裝。
fadd.export_library(temp.relpath(“myadd_pack.so”))
fadd2 = tvm.module.load(temp.relpath(“myadd_pack.so”))
fadd2(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())

注意，運行時API和線程安全
TVM的編譯模塊不依賴于TVM編譯器。相反，僅依賴于最小運行時庫。TVM運行時庫包裝設備驅動程序，并為編譯的函數提供線程安全和設備無關的調用。
可以從任何GPU上的任何線程，調用已編譯的TVM函數。
生成OpenCL代碼
TVM為多個后端提供代碼生成功能，可以生成在CPU后端上運行的OpenCL代碼或LLVM代碼。
以下代碼塊生成OpenCL代碼，在OpenCL設備上創建數組，驗證代碼的正確性。
if tgt.startswith(‘opencl’):
fadd_cl = tvm.build(s, [A, B, C], tgt, name=“myadd”)
print("------opencl code------")
print(fadd_cl.imported_modules[0].get_source())
ctx = tvm.cl(0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd_cl(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
Summary
本文使用向量添加示例介紹TVM工作流程。一般工作流程是
? 通過一系列操作描述計算。
? 描述如何計算使用schedule原函數。
? 編譯到想要的目標函數。
? （可選）保存稍后要加載的功能。
查看其它示例和文件，了解有關TVM中支持的操作，調度原語和其它功能的更多信息。

參考文獻：
https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html
https://blog.csdn.net/xxradon/article/details/98213214

總結

以上是生活随笔為你收集整理的使用Tensor Expression张量表达式处理算子的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。