【TVM 學習資料】TensorIR 快速入門

超神經HyperAI發表於2023-03-02

本篇文章譯自英文文件 Blitz Course to TensorIR
作者是 Siyuan Feng。更多 TVM 中文文件可訪問→TVM 中文站

TensorIR 是深度學習領域的特定語言，主要有兩個作用：

  • 在各種硬體後端轉換和最佳化程式。
  • 自動 tensorized 程式最佳化的抽象。
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

IRModule

IRModule 是 TVM 的核心資料結構，它包含深度學習程式，並且是 IR 轉換和模型構建的基礎。

在這裡插入圖片描述
上圖展示的是 IRModule 的生命週期，它可由 TVMScript 建立。轉換 IRModule 的兩種主要方法是 TensorIR 的 schedule 原語轉換和 pass 轉換。此外，也可直接對 IRModule 進行一系列轉換。注意，可以在任何階段將 IRModule 列印到 TVMScript。完成所有轉換和最佳化後，可將 IRModule 構建為可執行模組，從而部署在目標裝置上。

基於 TensorIR 和 IRModule 的設計，可建立一種新的程式設計方法：

  1. 基於 Python-AST 語法，用 TVMScript 編寫程式。
  2. 使用 Python API 轉換和最佳化程式。
  3. 使用命令式轉換 API 互動檢查和提高效能。

建立 IRModule

IRModule 是 TVM IR 的一種可往返語法，可透過編寫 TVMScript 來建立。

與透過張量表示式建立計算表示式（使用張量表示式操作運算元）不同，TensorIR 允許使用者透過 TVMScript（一種嵌在 Python AST 中的語言）進行程式設計。新方法可以編寫複雜的程式並進一步排程和最佳化。

下面是向量加法的示例：

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 我們透過 T.handle 進行資料交換，類似於記憶體指標
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 透過 handle 建立 Buffer
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # block 是針對計算的抽象
            with T.block("B"):
                # 定義一個空間（可並行）block 迭代器，並且將它的值繫結成 i
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0

ir_module = MyModule
print(type(ir_module))
print(ir_module.script())

輸出結果：

<class 'tvm.ir.module.IRModule'>
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[8, "float32"], B: T.Buffer[8, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i in T.serial(8):
            with T.block("B"):
                vi = T.axis.spatial(8, i)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

此外，我們還可以使用張量表示式 DSL 編寫簡單的運算子，並將它們轉換為 IRModule。

from tvm import te

A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")
func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())

輸出結果：

# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[8, "float32"], B: T.Buffer[8, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i0 in T.serial(8):
            with T.block("B"):
                i0_1 = T.axis.spatial(8, i0)
                T.reads(A[i0_1])
                T.writes(B[i0_1])
                B[i0_1] = A[i0_1] + T.float32(1)

構建並執行 IRModule

可將 IRModule 構建為特定 target 後端的可執行模組。

mod = tvm.build(ir_module, target="llvm")  # CPU 後端的模組
print(type(mod))

輸出結果：

<class 'tvm.driver.build_module.OperatorModule'>

準備輸入陣列和輸出陣列，然後執行模組：

a = tvm.nd.array(np.arange(8).astype("float32"))
b = tvm.nd.array(np.zeros((8,)).astype("float32"))
mod(a, b)
print(a)
print(b)

輸出結果：

[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]

轉換 IRModule

IRModule 是程式最佳化的核心資料結構，可透過 Schedule 進行轉換。schedule 包含多個 primitive 方法來互動地轉換程式。每個 primitive 都以特定方式對程式進行轉換，從而最佳化效能。

在這裡插入圖片描述
上圖是最佳化張量程式的典型工作流程。首先，用 TVMScript 或張量表示式建立一個初始 IRModule，然後在這個初始 IRModule 上建立 schedule。接下來，使用一系列排程原語來提高效能。最後，我們可以將其降低並構建成一個可執行模組。

上面只演示了一個簡單的轉換。首先，在輸入 ir_module 上建立 schedule：

sch = tvm.tir.Schedule(ir_module)
print(type(sch))

輸出結果：

<class 'tvm.tir.schedule.schedule.Schedule'>

將巢狀迴圈展開成 3 個迴圈，並列印結果：

# 透過名字獲取 block
block_b = sch.get_block("B")
# 獲取包圍 block 的迴圈
(i,) = sch.get_loops(block_b)
# 展開巢狀迴圈
i_0, i_1, i_2 = sch.split(i, factors=[2, 2, 2])
print(sch.mod.script())

輸出結果：

# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[8, "float32"], B: T.Buffer[8, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0, i_1, i_2 in T.grid(2, 2, 2):
            with T.block("B"):
                vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

還可對迴圈重新排序。例如，將迴圈 i_2 移到 i_1 之外：

sch.reorder(i_0, i_2, i_1)
print(sch.mod.script())

輸出結果

# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[8, "float32"], B: T.Buffer[8, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0, i_2, i_1 in T.grid(2, 2, 2):
            with T.block("B"):
                vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

轉換為 GPU 程式

要在 GPU 上部署模型必須進行執行緒繫結。幸運的是，也可以用原語來增量轉換。

sch.bind(i_0, "blockIdx.x")
sch.bind(i_2, "threadIdx.x")
print(sch.mod.script())

輸出結果：

# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[8, "float32"], B: T.Buffer[8, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(2, thread="blockIdx.x"):
            for i_2 in T.thread_binding(2, thread="threadIdx.x"):
                for i_1 in T.serial(2):
                    with T.block("B"):
                        vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                        T.reads(A[vi])
                        T.writes(B[vi])
                        B[vi] = A[vi] + T.float32(1)

繫結執行緒後，用 cuda 後端來構建 IRModule：

ctx = tvm.cuda(0)
cuda_mod = tvm.build(sch.mod, target="cuda")
cuda_a = tvm.nd.array(np.arange(8).astype("float32"), ctx)
cuda_b = tvm.nd.array(np.zeros((8,)).astype("float32"), ctx)
cuda_mod(cuda_a, cuda_b)
print(cuda_a)
print(cuda_b)

輸出結果：

[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]

下載 Python 原始碼
下載 Jupyter Notebook

以上就是該文件的全部內容，檢視更多 TVM 中文文件，請訪問→TVM 中文站

相關文章