Apache TVM 是一个深度的深度学习编译框架，适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 →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 tvmfrom tvm import relay# 基本程序：Z = X + Yx = relay.var("x", shape=(3,), dtype="float32")y = relay.var("y", shape=(3,), dtype="float32")z = x + yprogram = relay.Function([x, y], z)module = tvm.IRModule.from_expr(program)现使用 numpy 为程序创建随机输入：
import numpy as npnp.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_myfloatz = 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 ChangeDatatypesrc_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 ]