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 ]
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
