TVM自定义数据类型

TVM自定义数据类型

本文将介绍“自定义数据类型”框架，该框架可在TVM中使用自定义数据类型。

介绍

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

在研究人员开始为其数据类型构建硬件之前，需要确定其数据类型在关心的工作负载中如何以数字方式表现。这通常涉及建立其数据类型的软件仿真版本（例如Berkeley SoftFloat或libposit），将数据类型直接入侵工作负载中，以查看工作负载如何使用该数据类型执行工作。更好的是将数据类型直接集成到编译器本身中，以便可以编译许多不同的工作负载以使用该数据类型。两种路由都可能很乏味，考虑到现代编译器的大小和复杂性，后一种路由通常变得难以管理。取自GitHub的一个示例显示有人入侵了将数据类型存入TensorFlow。结果是237次提交，添加了将近6000行代码，并在整个代码库中触摸了200多个文件，而这仅仅是添加一种数据类型！对于许多研究人员来说，这项工作量是令人望而却步的。

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

自定义数据类型

自定义数据类型框架的目标，使用户能够使用自定义数据类型运行深度学习工作负载。在“自定义数据类型”框架中，“数据类型”表示标量类型： 例如，float 或uint。不处理更复杂的数据格式，例如块浮点数 或Intel的Flexpoint。此外，仅声称支持 这些标量数据类型的软件仿真版本；不明确支持在自定义数据类型硬件上进行编译和运行。

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

该框架被实现为一个注册表，与TVM的常规数据类型设施并排放置。用户与数据类型注册表进行交互的主要方式有两种：第一，数据类型注册， 第二，降低功能注册。

这些步骤分别类似于数据类型的声明和实现。

请注意，本文中所有引用的代码均基于TVM存储库的master分支commit 4cad71d。将使用一个示例posit数据类型，该数据类型可以src/target/datatype/posit/posit-wrapper.cc在TVM下找到，并可以在带有USE_BYODT_POSIT标志的TVM中进行编译。⁴

数据类型注册

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

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

上面的代码'posit'使用类型代码150注册数据类型。此注册步骤允许TVM解析使用自定义类型的程序：

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用来表示每个posit都是16位宽。（车道数默认为1。）

降低功能注册

尽管TVM可以解析上述程序，但它尚不能编译，TVM尚不了解如何在该posit类型上编译操作。为了编译这些程序，为自定义数据类型注册了降级函数，这有助于TVM将操作转换为它可以理解和编译的内容。

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

图1：用户注册的降低功能的预期结果。降低功能应将使用自定义数据类型的程序转换为本机TVM可以理解和编译的程序（在这种情况下，需要使用两个uint16_t来调用外部库）。

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

为了实现上述降低，为以下对象注册了降低功能posit：

tvm.target.datatype.register_op(

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

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

上面的代码为特定的运算符（Add），编译目标（LLVM），数据类型（posit）和位长（16）注册了一个降低函数。第一个参数是降低功能。这可以是采用TVM IR节点并返回新的TVM IR节点的任何功能。在案例中，使用Bring Your Own Datatypes框架提供的帮助程序功能。 tvm.target.datatype.create_lower_func({16:'Posit16es2Add'}) 为上述通用模式创建降低功能。结果函数将给定节点的参数转换为uint16_t，将节点本身转换为对给定函数名称的调用（在这种情况下，位长度'Posit16es2Add'为posits）。将一个字典传递给create_lower_func，以便TVM可以根据数据类型的位长，将其分配给适当的函数名称。

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

包起来wrapping up

自定义数据类型框架将用户定义的数据类型引入TVM。鼓励数据类型研究人员在研究中使用TVM；同样，引起深度学习社区中对自定义数据类型的兴趣。有关“携带自己的数据类型”框架的更多文档。