【TVM 教程】在 TVM 中使用 Bring Your Own Datatypes 原創
Apache TVM 是一個深度的深度學習編譯框架，適用于 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →https://tvm.hyper.ai/
作者：Gus Smith,?Andrew Liu
本教程將展示如何利用 Bring Your Own Datatypes 框架在 TVM 中使用自定義數據類型。注意，Bring Your Own Datatypes 框架目前僅處理數據類型的軟件模擬版本。該框架不支持開箱即用地編譯自定義加速器數據類型。
數據類型庫?
Bring Your Own Datatypes 允許用戶在 TVM 的原生數據類型（例如?float）旁邊注冊自己的數據類型實現。這些數據類型實現通常以庫的形式出現。例如：
- libposit，一個位置庫
- Stillwater Universal，一個包含位置、定點數和其他類型的庫
- SoftFloat，伯克利的 IEEE 754 浮點軟件實現
Bring Your Own Datatypes 使用戶能夠將這些數據類型實現插入 TVM！
本節中我們將用到一個已經實現的示例庫（位于?3rdparty/byodt/myfloat.cc）。這種稱之為「myfloat」的數據類型實際上只是一個 IEE-754 浮點數，但它提供了一個有用的示例，表明任何數據類型都可以在 BYODT 框架中使用。
設置?
由于不使用任何 3rdparty 庫，因此無需設置。
若要用自己的數據類型庫嘗試，首先用?CDLL?把庫的函數引入進程空間：
ctypes.CDLL('my-datatype-lib.so', ctypes.RTLD_GLOBAL)
一個簡單的 TVM 程序?
從在 TVM 中編寫一個簡單的程序開始，之后進行重寫，從而使用自定義數據類型。
import tvm
from tvm import relay
# 基本程序：Z = X + Y
x = relay.var("x", shape=(3,), dtype="float32")
y = relay.var("y", shape=(3,), dtype="float32")
z = x + y
program = relay.Function([x, y], z)
module = tvm.IRModule.from_expr(program)
現使用 numpy 為程序創建隨機輸入：
import numpy as np
np.random.seed(23) # 可重復性
x_input = np.random.rand(3).astype("float32")
y_input = np.random.rand(3).astype("float32")
print("x: {}".format(x_input))
print("y: {}".format(y_input))
輸出結果：
x: [0.51729786 0.9469626 0.7654598 ]
y: [0.28239584 0.22104536 0.6862221 ]
最后，準備運行程序：
z_output = relay.create_executor(mod=module).evaluate()(x_input, y_input)
print("z: {}".format(z_output))
輸出結果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
z: [0.7996937 1.168008 1.4516819]
添加自定義數據類型?
接下來使用自定義數據類型進行中間計算。
使用與上面相同的輸入變量?x?和?y，但在添加?x + y?之前，首先通過調用?relay.cast(...)?將?x?和?y?轉換為自定義數據類型。
注意如何指定自定義數據類型：使用特殊的?custom[...]?語法來表示。此外，注意數據類型后面的「32」：這是自定義數據類型的位寬，告訴 TVM?myfloat?的每個實例都是 32 位寬。
try:
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
x_myfloat = relay.cast(x, dtype="custom[myfloat]32")
y_myfloat = relay.cast(y, dtype="custom[myfloat]32")
z_myfloat = x_myfloat + y_myfloat
z = relay.cast(z_myfloat, dtype="float32")
except tvm.TVMError as e:
# 打印最后一行錯誤
print(str(e).split("\n")[-1])
嘗試生成此程序會從 TVM 引發錯誤。TVM 不知道如何創造性地處理所有自定義數據類型！因此首先要從 TVM 注冊自定義類型，給它一個名稱和一個類型代碼：
tvm.target.datatype.register("myfloat", 150)
注意，類型代碼 150 目前由用戶手動選擇。參閱?include/tvm/runtime/c_runtime_api.h?中的?TVMTypeCode::kCustomBegin。下面再次生成程序：
x_myfloat = relay.cast(x, dtype="custom[myfloat]32")
y_myfloat = relay.cast(y, dtype="custom[myfloat]32")
z_myfloat = x_myfloat + y_myfloat
z = relay.cast(z_myfloat, dtype="float32")
program = relay.Function([x, y], z)
module = tvm.IRModule.from_expr(program)
module = relay.transform.InferType()(module)
現在有了一個使用 myfloat 的Relay 程序！
print(program)
輸出結果：
fn (%x: Tensor[(3), float32], %y: Tensor[(3), float32]) {
%0 = cast(%x, dtype="custom[myfloat]32");
%1 = cast(%y, dtype="custom[myfloat]32");
%2 = add(%0, %1);
cast(%2, dtype="float32")
}
現在可以準確無誤地表達程序，嘗試運行！
try:
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
z_output_myfloat = relay.create_executor("graph", mod=module).evaluate()(x_input, y_input)
print("z: {}".format(y_myfloat))
except tvm.TVMError as e:
# 打印最后一行錯誤
print(str(e).split("\n")[-1])
輸出結果：
Check failed: (lower) is false: Cast lowering function for target llvm destination type 150 source type 2 not found
編譯該程序會引發錯誤，下面來剖析這個報錯。
該報錯發生在代碼降級的過程中，即將自定義數據類型代碼，降級為 TVM 可以編譯和運行的代碼。TVM 顯示，當從源類型 2（float，在 TVM 中）轉換到目標類型 150（自定義數據類型）時，它無法找到?Cast?操作的降級函數。
當對自定義數據類型進行降級時，若 TVM 遇到對自定義數據類型的操作，它會查找用戶注冊的降級函數，這個函數告訴 TVM 如何將操作降級為 TVM 理解的數據類型的操作。由于我們還沒有告訴 TVM 如何降級自定義數據類型的?Cast?操作，因此會報錯。
要修復這個錯誤，只需要指定一個降級函數：
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func(
{
(32, 32): "FloatToCustom32", # cast from float32 to myfloat32 # 從 float32 轉換為 myfloat32
}
),
"Cast",
"llvm",
"float",
"myfloat",
)
register_op(...)?調用接受一個降級函數和一些參數，這些參數準確地指定了應該使用提供的降級函數降級的操作。在這種情況下，傳遞的參數指定此降級函數用于將 target?“llvm”?的?Cast?從?float?降級到?myfloat。
傳遞給此調用的降級函數非常通用：它應該采用指定類型的操作（在本例中為?Cast）并返回另一個僅使用 TVM 理解的數據類型的操作。
通常，我們希望用戶借助對外部庫的調用，來對其自定義數據類型進行操作。在示例中，myfloat?庫在函數?FloatToCustom32?中實現了從?float?到 32 位?myfloat?的轉換。一般情況下，創建一個輔助函數?create_lower_func(...)，它的作用是：給定一個字典，它將給定的?Call的操作，替換為基于操作和位寬的適當函數名稱。它還通過將自定義數據類型存儲在適當寬度的不透明?uint?中，從而刪除自定義數據類型的使用；在我們的例子中，如?uint32_t。有關更多信息，參閱?源代碼。
# 現在重新嘗試運行程序：
try:
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
z_output_myfloat = relay.create_executor("graph", mod=module).evaluate()(x_input, y_input)
print("z: {}".format(z_output_myfloat))
except tvm.TVMError as e:
# 打印最后一行錯誤
print(str(e).split("\n")[-1])
輸出結果：
Check failed: (lower) is false: Add lowering function for target llvm type 150 not found
新報錯提示無法找到?Add?降級函數，這并不是壞事兒，這表明錯誤與?Cast無關！接下來只需要在程序中為其他操作注冊降級函數。
注意，對于?Add，create_lower_func?接受一個鍵（key）是整數的字典。對于?Cast?操作，需要一個 2 元組來指定?src_bit_length?和?dest_bit_length，對于其他操作，操作數之間的位長度相同，因此只需要一個整數來指定?bit_length。
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Add"}),
"Add",
"llvm",
"myfloat",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({(32, 32): "Custom32ToFloat"}),
"Cast",
"llvm",
"myfloat",
"float",
)
# 現在，可以正常運行程序了。
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
z_output_myfloat = relay.create_executor(mod=module).evaluate()(x_input, y_input)
print("z: {}".format(z_output_myfloat))
print("x:\t\t{}".format(x_input))
print("y:\t\t{}".format(y_input))
print("z (float32):\t{}".format(z_output))
print("z (myfloat32):\t{}".format(z_output_myfloat))
# 或許正如預期的那樣，``myfloat32`` 結果和 ``float32`` 是完全一樣的！
輸出結果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
z: [0.7996937 1.168008 1.4516819]
x: [0.51729786 0.9469626 0.7654598 ]
y: [0.28239584 0.22104536 0.6862221 ]
z (float32): [0.7996937 1.168008 1.4516819]
z (myfloat32): [0.7996937 1.168008 1.4516819]
使用自定義數據類型運行模型?
首先選擇要使用 myfloat 運行的模型，本示例中，我們使用的是?Mobilenet。選擇 Mobilenet 是因為它足夠小。在 Bring Your Own Datatypes 框架的這個 alpha 狀態下，還沒有為運行自定義數據類型的軟件仿真實現任何軟件優化；由于多次調用數據類型仿真庫，導致性能不佳。
首先定義兩個輔助函數，獲取 mobilenet 模型和貓圖像。
def get_mobilenet():
dshape = (1, 3, 224, 224)
from mxnet.gluon.model_zoo.vision import get_model
block = get_model("mobilenet0.25", pretrained=True)
shape_dict = {"data": dshape}
return relay.frontend.from_mxnet(block, shape_dict)
def get_cat_image():
from tvm.contrib.download import download_testdata
from PIL import Image
url = "https://gist.githubusercontent.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/fa7ef0e9c9a5daea686d6473a62aacd1a5885849/cat.png"
dst = "cat.png"
real_dst = download_testdata(url, dst, module="data")
img = Image.open(real_dst).resize((224, 224))
# CoreML's standard model image format is BGR
img_bgr = np.array(img)[:, :, ::-1]
img = np.transpose(img_bgr, (2, 0, 1))[np.newaxis, :]
return np.asarray(img, dtype="float32")
module, params = get_mobilenet()
輸出結果：
Downloading /workspace/.mxnet/models/mobilenet0.25-9f83e440.zipe0e3327d-26bc-4c47-aed4-734a16b0a3f8 from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/mobilenet0.25-9f83e440.zip...
用原生 TVM 很容易執行 MobileNet：
ex = tvm.relay.create_executor("graph", mod=module, params=params)
input = get_cat_image()
result = ex.evaluate()(input).numpy()
# 打印前 10 個元素
print(result.flatten()[:10])
輸出結果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
[ -7.5350165 2.0368009 -12.706646 -5.63786 -12.684058 4.0723605
2.618876 3.4049501 -9.867913 -24.53311 ]
若要更改模型在內部使用 myfloat，需要轉換網絡。為此首先定義一個函數來幫助轉換張量：
def convert_ndarray(dst_dtype, array):
"""Converts an NDArray into the specified datatype"""
x = relay.var("x", shape=array.shape, dtype=str(array.dtype))
cast = relay.Function([x], x.astype(dst_dtype))
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
return relay.create_executor("graph").evaluate(cast)(array)
為了實際轉換整個網絡，我們在 Relay 中編寫了?一個 pass，它簡單地將模型中的所有節點轉換為使用新的數據類型。
from tvm.relay.frontend.change_datatype import ChangeDatatype
src_dtype = "float32"
dst_dtype = "custom[myfloat]32"
module = relay.transform.InferType()(module)
# 目前，自定義數據類型僅在預先運行 simple_inference 時才有效
module = tvm.relay.transform.SimplifyInference()(module)
# 在更改數據類型之前運行類型推斷
module = tvm.relay.transform.InferType()(module)
# 將數據類型從 float 更改為 myfloat 并重新推斷類型
cdtype = ChangeDatatype(src_dtype, dst_dtype)
expr = cdtype.visit(module["main"])
module = tvm.relay.transform.InferType()(module)
# 轉換參數：
params = {k: convert_ndarray(dst_dtype, v) for k, v in params.items()}
# 還需要轉換輸入：
input = convert_ndarray(dst_dtype, input)
# 最后，可以嘗試運行轉換后的模型：
try:
# 向量化不是用自定義數據類型實現的。
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
result_myfloat = tvm.relay.create_executor("graph", mod=module).evaluate(expr)(
input, **params
)
except tvm.TVMError as e:
print(str(e).split("\n")[-1])
輸出結果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
Check failed: (lower) is false: Intrinsic lowering function for target llvm, intrinsic name tir.sqrt, type 150 not found
嘗試運行模型時，會收到一個熟悉的報錯，提示需要為 myfloat 注冊更多函數。
因為這是一個神經網絡，所以需要更多的操作。下面注冊所有需要的函數：
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "FloatToCustom32"}),
"FloatImm",
"llvm",
"myfloat",
)
tvm.target.datatype.register_op(
tvm.target.datatype.lower_ite, "Call", "llvm", "myfloat", intrinsic_name="tir.if_then_else"
)
tvm.target.datatype.register_op(
tvm.target.datatype.lower_call_pure_extern,
"Call",
"llvm",
"myfloat",
intrinsic_name="tir.call_pure_extern",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Mul"}),
"Mul",
"llvm",
"myfloat",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Div"}),
"Div",
"llvm",
"myfloat",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Sqrt"}),
"Call",
"llvm",
"myfloat",
intrinsic_name="tir.sqrt",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Sub"}),
"Sub",
"llvm",
"myfloat",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Exp"}),
"Call",
"llvm",
"myfloat",
intrinsic_name="tir.exp",
)
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({32: "Custom32Max"}),
"Max",
"llvm",
"myfloat",
)
tvm.target.datatype.register_min_func(
tvm.target.datatype.create_min_lower_func({32: "MinCustom32"}, "myfloat"),
"myfloat",
)
注意，我們使用的是：register_min_func?和?create_min_lower_func。
register_min_func?接收一個整數?num_bits?作為位長，然后返回一個表示最小有限可表示值的操作，這個值是具有指定位長的自定義數據類型。
與?register_op?和?create_lower_func?類似，create_min_lower_func?處理通過調用一個外部庫，實現最小可表示的自定義數據類型值的一般情況。
接下來運行模型：
# 向量化不是用自定義數據類型實現的。
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
result_myfloat = relay.create_executor(mod=module).evaluate(expr)(input, **params)
result_myfloat = convert_ndarray(src_dtype, result_myfloat).numpy()
# 打印前 10 個元素
print(result_myfloat.flatten()[:10])
# 再次注意，使用 32 位 myfloat 的輸出與 32 位浮點數完全相同，
# 因為 myfloat 就是一個浮點數！
np.testing.assert_array_equal(result, result_myfloat)
輸出結果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
[ -7.5350165 2.0368009 -12.706646 -5.63786 -12.684058 4.0723605
2.618876 3.4049501 -9.867913 -24.53311 ]
