自主数据类型：在TVM中启用自定义数据类型探索

自主数据类型：在TVM中启用自定义数据类型探索

介绍

在设计加速器时，一个重要的决定是如何在硬件中近似地表示实数。这个问题有一个长期的行业标准解决方案：IEEE 754浮点标准.1。然而，当试图通过构建高度专业化的设计来最大限度地利用硬件时，使用通用IEEE 754浮点有意义吗？如果知道工作负载的数字需求，是否可以构建一个更小、更快或更省电的数据类型？答案是肯定的！研究人员已经开始在学术和工业加速器设计中尝试新的数据类型。例如，Google的张量处理单元（TPU）使用bfloat类型：一个被截断为16位的单精度IEEE浮点。由于许多深度学习工作负载对数值要求不严格，这种截断通常对模型精度没有影响，同时立即将存储成本降低一半。

然而，在研究人员开始为数据类型构建硬件之前，首先需要确定数据类型在工作负载中的数值行为。这通常需要首先构建其数据类型的软件仿真版本（例如Berkeley SoftFloat或libposit），然后将数据类型直接插入工作负载，以查看工作负载如何使用数据类型执行。更好的方法是将数据类型直接集成到编译器本身中，这样就可以编译许多不同的工作负载来使用该数据类型。这两种路径都可能是乏味的，鉴于现代编译器的规模和复杂性，后一种路径往往变得不可管理。一个来自GitHub的例子显示有人将posit数据类型侵入TensorFlow。结果是237次提交，添加了近6000行代码，并在代码库中处理了200多个文件，这仅仅是添加了一个数据类型！对许多研究人员来说，这样的工作量是令人望而却步的。

为了解决这些问题，提出了自带数据类型框架。该框架允许用户将模拟的数据类型插入TVM，从而可以轻松地在深度学习工作负载中探索新的数据类型。与上面的Tensorflow示例中的positions不同，它支持编译器中的单个新数据类型，而Bring Your Own datatype框架支持各种各样的用户定义类型。

自主数据类型

“自主数据类型”框架的目标是使用户能够使用自定义数据类型运行深入学习的工作负载。在Bring Your Own Datatypes框架中，“datatype”表示标量类型：例如float或uint。不处理像Intel Flexpoint这样复杂的数据块格式（floating point or Intel’s Flexpoint）。此外，只声明支持这些标量数据类型的软件仿真版本；不明确支持在自定义数据类型硬件上编译和运行。

TVM中的每个张量都分配了一个类型代码，它定义了张量中标量的数据类型。这些类型代码在TVM中具有硬编码的含义，映射到诸如int和float之类的常见数据类型。然而，绝大多数类型代码都是未使用的。“自带数据类型”框架允许用户声明这些未使用的类型代码，并在运行时添加自己的新数据类型。

该框架是作为注册表来实现的，它位于TVM的常规数据类型架构上。用户与数据类型注册表交互的主要方式有两种：第一种是数据类型注册，第二种是降低函数注册。这些步骤分别类似于数据类型的声明和实现。

本文中引用的所有代码都基于TVM存储库的主分支提交4cad71d。将使用一个示例posit数据类型，该数据类型位于src/target/datatype/posit/posit-wrapper.cc并且可以在TVM中使用USE_BYODT_POSIT标志进行编译。

Datatype Registration

为了注册数据类型，用户为数据类型分配一个名称和一个类型代码，其中类型代码来自自定义数据类型可用的未使用类型代码范围。

tvm.target.datatype.register('posit', 150)

The above code registers the 'posit' datatype with type code 150. This registration step allows TVM to parse programs which use the custom type:

x = relay.var('x', shape=(3, ), dtype='float32')

y = relay.var('y', shape=(3, ), dtype='float32')

x_posit = relay.cast(x, dtype='custom[posit]16')

y_posit = relay.cast(y, dtype='custom[posit]16')

z_posit = x_posit + y_posit

z = relay.cast(z_posit, dtype='float32')

program = relay.Function([x, y], z)

print(program)

# v0.0.4

# fn (%x: Tensor[(3), float32], %y: Tensor[(3), float32]) {

#   %0 = cast(%x, dtype="custom[posit]16");

#   %1 = cast(%y, dtype="custom[posit]16");

#   %2 = add(%0, %1);

#   cast(%2, dtype="float32")

# }

上面的程序将float32输入x和y转换成posits，相加，然后将结果转换回float32。一旦posit类型被注册，TVM就能够解析特殊的dtype语法custom[<typename>]，其中<typename>是为类型注册的名称。此语法还支持通常的<bits>x<lanes>格式；这里，使用16来表示每个posite的宽度是16位。（通道数默认为1。）

Lowering Function Registration

虽然TVM可以解析上述程序，但它还不能编译它，因为TVM还不知道如何在posit类型上编译操作。为了编译这些程序，为自定义数据类型注册了降低函数，这有助于TVM将操作转换成它可以理解和编译的内容。

一般来说，用户不需要直接将操作降低到LLVM或CUDA。相反，大多数使用自定义数据类型的代码可以通过一些简单的技巧被简化为不使用自定义数据类型的代码。然后，可以依赖于本机TVM来理解和编译代码。

图1：用户注册的降低函数的预期结果。

降低函数应该将使用自定义数据类型的程序转换为本机TVM可以理解和编译的程序（在本例中，调用外部库，使用两个uint16）。

图1显示了一个通用模式。假设对探索posit类型感兴趣，并选择通过将posit仿真库（例如Stillwater Universal）通过自带的数据类型框架插入TVM来运行一些工作负载。工作量是一个简单的程序，它添加了两个posite输入。Native-TVM不知道如何实现posit加法，但它不需要，因为有一个实现数据类型的库！这个库包含一个正数加法的实现，以及乘法和平方根等其他运算符。为了实现这个posit加法，只想调用库。因此，Add节点应该成为一个调用节点，调用库中的一个函数（称之为position16es2add）。为了将输入位置的位存储在TVM理解的类型中，使用16位无符号整数。生成的程序是TVM可以理解和编译的程序，它只是调用一个外部库函数，取两个无符号整数。

为了实现上述降低，为posit注册了一个降低功能：

tvm.target.datatype.register_op(

tvm.target.datatype.create_lower_func({16: 'Posit16es2Add'}),

'Add', 'llvm', 'posit')

上面的代码为特定运算符（Add）、编译目标（LLVM）、数据类型（posit）和位长度（16）注册了一个降低函数。第一个参数是lowering函数。这可以是任何获取TVM IR节点并返回新TVM IR节点的函数。在例子中，使用了Bring Your Own Datatypes框架提供的helper函数。tvm.target.datatype.create_lower_func({16:'Posit16es2Add'})为上述公共模式创建一个降低函数。结果函数将给定节点的参数转换为uint16_t，然后将节点本身转换为对给定函数名的调用（在本例中，对于位长度为16的position16es2add）。我们传递一个字典来create_lower_func，以便TVM可以根据数据类型的位长度将其分派到适当的函数名。

要实现自定义数据类型，用户将需要为他们希望运行的工作负载中的每个运算符注册一个降低函数。对于像ResNet这样的网络，大约有10个操作符，包括Add、Div、各种类型转换和Max。在测试中，注册一个数据类型和所有的降低函数需要大约40行Python。一旦注册了所有需要的算子，自定义数据类型工作负载就可以像任何其他TVM程序一样轻松运行！

结束

“自带数据类型”框架将用户定义的数据类型引入TVM。希望这将鼓励数据类型研究人员在研究中使用TVM；同样，希望这将激发深度学习社区中对定制数据类型的兴趣。有关“自带数据类型”框架的更多文档，请访问“将自己的数据类型带到TVM开发人员”。