Relay 类型系统

Relay 类型系统#

在详细描述 Relay 的表达式语言时简要介绍了类型，但尚未描述其类型系统。Relay 是一种静态类型和类型推断的语言，允许程序完全类型化，同时只需要少量显式类型注解。

静态类型在执行编译器优化时非常有用，因为它们传达了程序操作的数据的属性，例如运行时形状、数据布局和存储，而无需运行程序。Relay的 代数数据类型 允许轻松灵活地组合类型，以便构建可以归纳推理并用于编写递归函数的数据结构。

Relay 的类型系统具有一种形状的 依赖类型 形式。也就是说，它的类型系统跟踪 Relay 程序中张量的形状。将张量形状视为类型允许 Relay 在编译时执行更强大的推理；特别是，Relay 可以静态地推理那些输出形状以复杂方式基于输入形状变化的操作。将形状推断视为类型推断问题允许 Relay 在编译时推断所有张量的形状，包括使用分支和函数调用的程序。

以这种方式静态推理形状允许 Relay 进行提前编译，并为编译管道中的优化提供更多关于张量的信息。这些优化可以作为 pass 实现，即 Relay 到 Relay 的 AST 转换，并且可以使用推断的类型（例如，形状信息）来做出程序转换的决策。例如，src/relay/transforms/fuse_ops.cc 提供了 pass 的实现，该 pass 使用推断的张量形状来替换 Relay 程序中算子的调用为融合算子实现。

在 Relay 中推理张量类型使用 类型关系 编码，这意味着 Relay 中的大部分类型检查是约束求解（确保所有类型关系在调用点都得到满足）。类型关系提供了一种灵活且相对简单的方式，使依赖类型的力量在 Relay 中可用，而不会大大增加其类型系统的复杂性。

下面详细介绍了 Relay 中的类型语言以及它们如何分配给 Relay 表达式。

类型#

所有 Relay 类型的基础类型。所有 Relay 类型都是此基础类型的子类。

请参阅 Type 以获取其定义和文档。

张量类型#

Relay 中的具体张量类型。

张量根据数据类型和形状进行类型化。目前，这些使用 TVM 的数据类型和形状，但将来 Relay 可能包括单独的 AST 用于形状。特别是，数据类型包括 bool、float32、int8 以及各种其他位宽和通道数。形状以维度元组（TVM IndexExpr）给出，例如 (5, 5)；标量也给出元组类型，并具有形状 ()。

但请注意，TVM 形状也可以包括变量和包括变量的算术表达式，因此 Relay 的约束求解阶段将尝试为所有形状变量找到赋值，以确保在运行程序之前所有形状都是具体的。

例如，这里简单的具体张量类型，对应于 10x10 的 32 位浮点数张量：

Tensor[(10, 10), float32]

请参阅 TensorType 以获取其定义和文档。

元组类型#

Relay 中元组的类型。

正如元组只是静态已知长度的值序列，元组的类型由对应于元组每个成员的类型的序列组成。

由于元组类型是静态已知大小的，元组投影的类型只是元组类型中的相应索引。

例如，在下面的代码中，%t 的类型为 (Tensor[(), bool], Tensor[(10, 10), float32])，而 %c 的类型为 Tensor[(10, 10), float32]。

let %t = (False, Constant(1, (10, 10), float32));
let %c = %t.1;
%c

请参阅 TupleType 以获取其定义和文档。

类型参数#

类型参数表示用于函数多态性的占位符类型。类型参数根据 种类 指定，对应于允许这些参数替换的类型：

  • Type，对应于顶级 Relay 类型，如张量类型、元组类型和函数类型。

  • BaseType，对应于张量的基础类型（例如，float32、bool）。

  • Shape，对应于张量形状。

  • ShapeVar，对应于张量形状中的变量。

Relay 的类型系统强制类型参数只允许出现在其种类允许的地方，因此如果类型变量 t 的种类为 Type，则 Tensor[t, float32] 不是有效类型。

与普通参数一样，必须在 call sites 为类型参数提供具体参数。

例如，下面的 s 是种类为 Shape 的类型参数，它将在下面的调用点被替换为 (10, 10)：

def @plus<s : Shape>(%t1 : Tensor[s, float32], %t2 : Tensor[s, float32]) {
     add(%t1, %t2)
}
plus<(10, 10)>(%a, %b)

请参阅 TypeVar 以获取其定义和文档。

类型约束#

这是表示类型约束的抽象类，将在进一步的版本中详细说明。目前，类型关系是唯一提供的类型约束；它们将在下面讨论。

请参阅 TypeConstraint 以获取其定义和文档。

函数类型#

Relay 中的函数类型，更多详情请参见 tvm/relay/type.h。

这是分配给 Relay 中函数的类型。函数类型由类型参数列表、类型约束集合、参数类型序列和返回类型组成。

非正式地将函数类型写为：fn<type_params>(arg_types) -> ret_type where type_constraints。

函数类型中的类型参数可能出现在参数类型或返回类型中。此外，每个类型约束必须在函数的每个调用点都成立。类型约束通常将函数的参数类型和返回类型作为参数，但也可能只取子集。

请参阅 FuncType 以获取其定义和文档。

类型关系#

类型关系是 Relay 中最复杂的类型系统特性。它允许用户使用新规则扩展类型推断。使用类型关系来定义以复杂方式处理张量形状的运算符的类型，例如广播运算符或 flatten，允许 Relay 在这些情况下静态推理形状。

类型关系 R 描述了 Relay 函数的输入和输出类型之间的关系。即，R 是类型上的函数，如果关系成立则输出 true，如果关系不成立则输出 false。提供给关系的类型可能是不完整的或包括形状变量，因此类型推断必须为不完整类型和形状变量分配适当的值以使必要的关系成立，如果存在这样的值。

例如，可以定义恒等关系为：

Identity(I, I) :- true

通常，通过定义特定于该算子的关系来为 Relay 中的算子类型化是很方便的，该关系编码了参数类型和返回类型的所有必要约束。例如，可以为 flatten 定义关系：

Flatten(Tensor(sh, bt), O) :-
  O = Tensor(sh[0], prod(sh[1:]))

如果有像 Broadcast 这样的关系，就可以为像 add 这样的算子类型化：

add : fn<t1 : Type, t2 : Type, t3 : Type>(t1, t2) -> t3
            where Broadcast

上面包含的 Broadcast 表示参数类型和返回类型必须是张量，其中 t3 的形状是 t1 和 t2 形状的广播结果。只要参数类型和返回类型满足 Broadcast，类型系统就会接受它们。”

请注意，上述示例关系是用类似 Prolog 的语法编写的，但目前这些关系必须由用户在 C++ 或 Python 中实现。更具体地说，Relay 的类型系统使用了一种 临时 求解器来处理类型关系，其中类型关系实际上是作为 C++ 或 Python 函数实现的，这些函数检查关系是否成立，并强制更新任何形状变量或不完整类型。在当前实现中，如果关系不成立，实现关系的函数应返回 False；如果关系成立或没有足够的信息来确定关系是否成立，则应返回 True。

所有关系的函数会在需要时（例如当输入更新时）运行，直到满足以下条件之一为止：

  1. 所有关系均成立且不存在未完成的类型（类型检查成功）。

  2. 某个关系未能成立（出现类型错误）。

  3. 达到固定点，其中形状变量或不完整类型仍然存在（可能需要类型错误或更多类型注释）。

目前，Relay 中使用的所有关系都是用 C++ 实现的。有关用 C++ 实现的关系的示例，请参见 src/relay/op 目录中的文件。

请参阅 TypeRelation 以了解其定义和文档说明。

不完整类型#

不完整类型”是指尚未确定的类型或类型的一部分。它仅在类型推断过程中使用。任何省略的类型注解都会被替换为不完整类型，随后会在某个时刻被其他类型所替代。

在编程语言文献中，不完整类型被称为“类型变量”("type variables")或“类型孔”("type holes")。使用“不完整类型”这一名称，以便更清晰地将它们与类型参数区分开来：类型参数必须绑定到函数上，并在调用点被具体类型参数替换（实例化），而不完整类型可能出现在程序的任何位置，并在类型推断过程中被填充。

请参阅 IncompleteType 以了解其定义和文档说明。

代数数据类型#

注意：文本格式目前不支持ADT（代数数据类型）。

代数数据类型（Algebraic data types，简称 ADTs）在 其概述 中有更详细的描述；本节将介绍它们在类型系统中的实现。

代数数据类型（ADT）由一组命名的构造函数定义，每个构造函数接受特定类型的参数。ADT 的实例是容器，它存储用于生成该实例的构造函数参数的值以及构造函数的名称；这些值可以通过基于其构造函数进行匹配来解构实例并检索。因此，ADT 有时被称为“带标签的联合体”：与 C 风格的联合体类似，在某些情况下，给定 ADT 的实例内容可能具有不同的类型，但构造函数充当标签以指示如何解释这些内容。

从类型系统的角度来看，最关键的是 ADT 可以接受类型参数（构造函数的参数可以是类型参数，尽管具有不同类型参数的 ADT 实例必须被视为不同的类型），并且可以是递归的（ADT 的构造函数可以接受该 ADT 的实例，因此可以归纳地构建像树或列表这样的 ADT）。类型系统中 ADT 的表示必须能够适应这些特性，以下部分将详细说明。

全局类型变量#

为了简洁地表示 ADT 并方便递归 ADT 的定义，ADT 定义会被赋予全局类型变量作为句柄，以唯一标识它。每个 ADT 定义都会被赋予新的全局类型变量作为句柄，因此可以通过指针相等性来区分不同的 ADT 名称。

在 Relay 类型系统中，ADT 通过名称进行区分；这意味着如果两个 ADT 具有不同的句柄，即使它们的所有构造函数在结构上完全相同，它们也会被视为不同的类型。

因此，ADT 定义中的递归与全局函数的递归类似：构造函数只需在其定义中引用 ADT 句柄（全局类型变量）即可。

请参阅 GlobalTypeVar 以了解其定义和文档说明。

定义（类型数据）#

除了名称之外，ADT 还需要存储用于定义它的构造函数以及其中使用的任何类型参数。这些内容存储在模块中，类似于全局函数定义。

在类型检查 ADT 的使用时，类型系统有时必须使用 ADT 名称索引到模块中，以查找有关构造函数的信息。例如，如果在匹配表达式子句中对构造函数进行模式匹配，类型检查器必须检查构造函数的签名，以确保为任何绑定变量分配了正确的类型。

请参阅 TypeData 以了解其定义和文档说明。

类型调用#

由于 ADT 定义可以接受类型参数，Relay 的类型系统将 ADT 定义视为一种 类型级函数 （即该定义接受类型参数并返回具有这些类型参数的 ADT 实例的类型）。因此，任何 ADT 实例的类型都通过类型调用来表示，其中明确列出了提供给 ADT 定义的类型参数。

列出 ADT 实例的类型参数非常重要，因为使用不同构造函数但相同类型参数构建的两个 ADT 实例属于 同一类型，而具有不同类型参数的两个 ADT 实例不应被视为同一类型（例如，整数列表不应与浮点张量对列表具有相同的类型）。

类型调用中的“函数”是 ADT 句柄，并且必须为 ADT 定义中的每个类型参数提供参数。（没有参数的 ADT 定义意味着任何实例都不会有类型参数传递给类型调用）。

请参阅 TypeCall 以了解其定义和文档说明。

示例：列表 ADT#

本小节使用简单的列表 ADT（作为 Relay 中的默认 ADT 包含）来说明前面部分描述的结构。其定义如下：

data List<a> {
  Nil : () -> List
  Cons : (a, List[a]) -> List
}

因此，全局类型变量 List 是 ADT 的句柄。模块中列表 ADT 的类型数据表明 List 接受类型参数并具有两个构造函数：Nil`（签名为 :code:`fn<a>() -> List[a]）和 Cons`（签名为 :code:`fn<a>(a, List[a]) -> List[a]）。在 Cons 构造函数中对 List 的递归引用是通过在构造函数定义中使用全局类型变量 List 来实现的。

以下是两个列表实例及其类型，使用类型调用表示：

Cons(1, Cons(2, Nil())) # List[Tensor[(), int32]]
Cons((1, 1), Cons((2, 2), Nil())) # List[(Tensor[(), int32], Tensor[(), int32])]

请注意，Nil() 可以是任何列表的实例，因为它不接受任何使用类型参数的参数。（然而，对于 Nil() 的任何*特定*实例，必须指定类型参数。）

以下是两个被类型系统拒绝的列表，因为类型参数不匹配：

# attempting to put an integer on a list of int * int tuples
Cons(1, Cons((1, 1), Nil()))
# attempting to put a list of ints on a list of lists of int * int tuples
Cons(Cons(1, Cons(2, Nil())), Cons(Cons((1, 1), Cons((2, 2), Nil())), Nil()))