自定義pass編寫
TVM是一個框架，抽象了機器學習加速器的異質性。有時，用戶可能需要自定義一些分析和IR轉換，使TVM適應自己的專用硬件。本文可幫助用戶在TVM中編寫自定義pass。
先決條件
在閱讀本文之前，假設讀者已經熟悉以下主題：
? 在TVM中編寫算法并進行調度。否則，請參見示例教程，例如 如何在CPU上優(yōu)化GEMM。
? HalideIR的基本結構。否則，請參閱HalideIR/src/ir/IR.h以了解定義了IR節(jié)點的哪些屬性。
? 訪客設計模式。否則，請檢查 Python AST模塊以查看AST訪問者的實現(xiàn)方式。
? 如何將Schedule降低為IRModule類或LLVM模塊。否則，請參考python/tvm/build_module.py以獲得一些基礎知識。
import tvm
from tvm import te
import numpy as np
首先編寫一個非常簡單的矢量加法，并使用默認調度對其進行構建。然后，使用定制的下降通道來直接操縱IR，而不是使用調度原語。
n = tvm.tir.const(128, “int32”)
a = te.placeholder((n,), name=“a”)
b = te.placeholder((n,), name=“b”)
c = te.compute((n,), lambda i: a[i] + b[i], name=“c”)

sch = te.create_schedule(c.op)
ir = tvm.lower(sch, [a, b, c])
print(ir)
輸出：
primfn(a_1: handle, b_1: handle, c_1: handle) -> ()
attr = {“global_symbol”: “main”, “tir.noalias”: True}
buffers = {c: Buffer(c_2: Pointer(float32), float32, [128], []),
b: Buffer(b_2: Pointer(float32), float32, [128], []),
a: Buffer(a_2: Pointer(float32), float32, [128], [])}
buffer_map = {a_1: a, b_1: b, c_1: c} {
for (i: int32, 0, 128) {
c_2[i] = ((float32*)a_2[i] + (float32*)b_2[i])
}
}
寫pass
本質上，“ IR轉換遍歷”是將語句映射到新語句的功能。因此，定義此向量化函數(shù)并逐步實現(xiàn)它。
TVM已經為用戶提供了兩類來分析和轉換IR。
IR Vistor訪客
可以用tvm.tir.stmt_functor.post_order_visit(stmt, func)funcfunc來從Halide IR收集信息。這是一個函數(shù)回調。在退出當前IR節(jié)點之前，即在后訂單訪問之前，將調用此函數(shù)。然后，利用side effects 副作用來存儲IR訪問的結果，返回值將被忽略。
必須使用一些數(shù)組來存儲IR訪問的結果。該值甚至是一個變量。這主要是由于Python-C運行時中的限制。每次遞歸都會刷新變量值，但會保留數(shù)組值。
loops = []
def find_width8(op):
“”" Find all the ‘tir.For’ nodes whose extent can be divided by 8. “”"
if isinstance(op, tvm.tir.For):
if isinstance(op.extent, tvm.tir.IntImm):
if op.extent.value % 8 == 0:
loops.append(op)
IR轉換
轉換界面與訪問者界面略有不同。訪問者中僅存在一個后回調，但是轉換訪問者既支持前回調又支持后回調。如果要保留原始IR節(jié)點，只需返回None。如果要將當前節(jié)點更改為某個節(jié)點，請使用TVM IR maker界面進行構建并返回此值。
筆記
如果調用了預訂功能并返回了非“無”的值，則將跳過 post-order 功能。
def vectorize8(op):
“”" Split can vectorize the loops found in find_width8. “”"
if op in loops:
extent = op.extent.value
name = op.loop_var.name
lo, li = te.var(name + “.outer”), te.var(name + “.inner”)
body = tvm.tir.stmt_functor.substitute(op.body, {op.loop_var: lo * 8 + li})
body = tvm.tir.For(li, 0, 8, tvm.tir.ForKind.VECTORIZED, body)
body = tvm.tir.For(lo, 0, extent // 8, tvm.tir.ForKind.SERIAL, body)
return body
return None

@tvm.tir.transform.prim_func_pass(opt_level=0)
def vectorize(f, mod, ctx):
global loops

tvm.tir.stmt_functor.post_order_visit(f.body, find_width8)if not loops:return sf# The last list arugment indicates what kinds of nodes will be transformed.
# Thus, in this case only `For` nodes will call `vectorize8`
return f.with_body(tvm.tir.stmt_functor.ir_transform(f.body, None, vectorize8, ["tir.For"]))

降低膠水Glue to Lowering
到目前為止，已經完成了編寫此IR轉換通道的操作。接下來，需要將該pass粘貼到TVM的較低pass上。
在這種情況下，通過將元組列表作為參數(shù)提供給TVM標準降低passtir.add_lower_pass。“元組”表示降低的不同階段。在TVM中，降級分為四個階段，每個階段完成后將調用用戶自定義的階段。
筆記
以下是每個階段完成的基本轉換：
? 階段0生成原始IR和環(huán)路電平。
? 第1階段將陣列存儲平坦化。
? 階段2轉換循環(huán)，例如展開，向量化和線程綁定。
? 第三階段進行一些清理工作。
因此，放置此轉換過程的好地方就在階段1之后。
with tvm.transform.PassContext(config={“tir.add_lower_pass”: [(1, vectorize)]}):
print(tvm.lower(sch, [a, b, c]))
輸出：
primfn(a_1: handle, b_1: handle, c_1: handle) -> ()
attr = {“global_symbol”: “main”, “tir.noalias”: True}
buffers = {b: Buffer(b_2: Pointer(float32), float32, [128], []),
c: Buffer(c_2: Pointer(float32), float32, [128], []),
a: Buffer(a_2: Pointer(float32), float32, [128], [])}
buffer_map = {a_1: a, b_1: b, c_1: c} {
for (i.outer: int32, 0, 16) {
c_2[ramp((i.outer8), 1, 8)] = ((float32x8)a_2[ramp((i.outer8), 1, 8)] + (float32x8)b_2[ramp((i.outer*8), 1, 8)])
}
}
快速瀏覽
本文提供了編寫自定義IR轉換過程的快速視圖：-tvm.tir.stmt_functor.post_order_visit用于收集每個IR節(jié)點上的信息。-tvm.tir.stmt_functor.ir_transform用于轉換IR節(jié)點。-總結以上兩個內容，編寫一個IR轉換功能。-用tvm.transform.PassContext將此功能用于TVM降準

總結

以上是生活随笔為你收集整理的自定义pass编写的全部內容，希望文章能夠幫你解決所遇到的問題。

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