relay的类型系统

我们在详细介绍Relay的表达语言时简要介绍了类型,但尚未描述其类型系统。relay是一种静态类型化和类型推断的语言,可以在仅需要几个显式类型注释的情况下对程序进行完全类型化。
静态类型在执行编译器优化时很有用,因为它们可以传达有关程序处理的数据的属性,例如运行时形状、数据布局和存储,而无需运行程序。Relay的代数数据类型允许轻松灵活地构成类型,以便构建可以归纳推理并用于编写递归函数的数据结构。
relay的类型系统具有形状的依赖类型(dependent typing)形式。也就是说,其类型系统会在Relay程序中跟踪张量的形状。将张量形状视为类型可以使Relay在编译时执行更强大的推理;特别是,Relay可以以复杂的方式静态地推断出某个操作,该操作的输出形状因输入形状而异 。将形状推断转换为类型推断问题后,Relay可以在编译时推断所有张量的形状,包括使用分支和函数调用的程序。
以这种方式对形状进行静态推理可以使Relay提前进行编译,并提供有关张量的更多信息,以便进一步在编译pipeline中进行优化。这样的优化可以实现为pass,这是relay到relay AST转换,并且可以使用推断的类型(例如,形状信息)来做出有关程序转换的决策。例如,src/relay/pass/fuse_ops.cc给出了一个pass的实现,该pass使用推断的张量形状来替换算子调用(带有融合算子实现的Relay程序中)。
Relay中的关于张量类型的推理是使用类型关系编码的,这意味着Relay中的大量类型检查是约束求解(确保在调用位置满足所有类型关系)。类型关系提供了一种灵活且相对简单的方式,使relay中的依赖类型的功能可用,而不会大大增加其类型系统的复杂性。
下面我们详细介绍Relay中类型的语言以及它们如何分配给Relay表达式。

Type

Type是所有relay类型的基本类型,所有relay类型都是此基本类型的子类。
有关其定义和文档,请参见Type

张量类型TensorType

张量的类型是由数据类型和形状决定的。当前,它们使用TVM的数据类型和形状,但是在将来,Relay可能会为形状包含单独的AST。具体地,数据类型包括boolfloat32int8和各种其它位宽和通道数。形状以不同规模的元组(TVM IndexExpr)的形式给出,例如(5, 5)表示5x5的二维形状,标量也构成元组类型,形状为()。
但是请注意,TVM形状还包含变量以及包括变量的算术表达式,因此Relay的约束求解阶段将尝试找到所有形状变量的分配,以确保在运行程序之前所有形状都是具体的。
例如,这是一个简单的具体张量类型,对应于32位浮点数的10×10张量:
Tensor[(10, 10), float32]
有关其定义和文档,请参见TensorType

元组类型**TupleType**

正如元组只是已知长度的值的序列一样,元组类型也包括与元组的每个成员相对应的类型序列。
因为元组类型的大小是静态已知的,所以元组投影的类型就是元组类型中的相应索引。
例如,在下面的代码,%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

类型参数

类型参数表示用于函数中多态的占位符类型。类型参数是根据kind指定的,与允许这些参数替换的类型相对应:

  • Type,对应于顶级Relay类型,例如张量类型,元组类型和函数类型
  • BaseType,对应于基础类的张量(例如,float32bool
  • Shape,对应于张量形状
  • ShapeVar,对应于张量形状内的变量

relay的类型系统强制规定类型参数只能出现在其类型允许它们出现的地方,因此,如果类型变量tType类型,则Tensor[t, float32]不是有效类型(t位置应该是Shape类型的参数)。
像普通参数一样,必须在调用位置为类型参数提供具体参数。
例如,下面的s是Shape类型的类型参数,它将在下面的调用位置被替换为(10, 10)
def @plus(%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(arg_types) -> ret_type where type_constraints
函数类型中的类型参数可能会出现在参数类型或返回类型中。此外,每个类型约束都必须在函数的每个调用位置都存在。类型约束通常采用函数的参数类型和函数的返回类型作为参数,但也可以采用子集。
有关其定义和文档,请参见FuncType

类型关系

类型关系是Relay中最复杂的类型系统功能。它允许用户使用新规则扩展类型推断。我们使用类型关系为那些以复杂方式处理张量形状的算子定义类型,例如广播算子或 flatten,从而允许Relay在这些情况下静态推断形状。
类型关系R描述了一个relay 函数的输入和输出类型之间的关系。即R是一种类型上的函数,如果该关系成立, 则输出true,如果关系不成立,则输出false。提供给该关系的类型可能不完整或包含形状变量,因此类型推断必须为不完整的类型和形状变量分配适当的值,以保持必要的关系(如果这样的值存在的话)。
例如,我们可以将一个identity关系定义为:
Identity(I, I) :- true
通常,通过定义特定于该算子的关系来对Relay中的算子进行类型化,该关系对参数类型和返回类型的所有必要约束进行编码。例如,我们可以为flatten定义关系:
Flatten(Tensor(sh, bt), O) :-
O = Tensor(sh[0], prod(sh[1:]))
如果我们有一个类似Broadcast的关系,就可以像这样键入算子add
add : fn(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

不完整类型

不完整类型是尚不知道的类型或类型的一部分,仅在类型推断期间使用。任何省略的类型注释都将由不完整的类型替换,稍后将由另一种类型替换。
不完整的类型在编程语言文献中被称为“类型变量”或“类型孔”。为了更清楚地将它们与类型参数区分开,我们使用名称“不完整类型”:类型参数必须绑定到函数,并在调用位置用具体的类型参数(实例化)替换,而不完整类型则可能出现在程序中的任何位置并在类型推断期间填写。
有关其定义和文档,请参见IncompleteType

代数数据类型(ADT)

注意:文本格式当前不支持ADT。
代数数据类型(ADT)在其概述中有更详细的描述。本节介绍它们在类型系统中的实现。
ADT由一组命名构造函数定义,每个命名构造函数都接受某些类型的参数。ADT的实例是一个容器,它存储用于产生该值的构造函数参数的值以及构造函数的名称;可以通过根据实例的构造函数进行匹配来解构实例,从而检索值。因此,ADT有时被称为“标记联合”:就像C样式的联合一样,给定ADT的实例的内容在某些情况下可能具有不同的类型,但是构造函数用作指示如何解释内容的标记。
从类型系统的角度来看,最相关的是ADT可以接受类型参数(构造函数参数可以是类型参数,尽管具有不同类型参数的ADT实例必须视为不同的类型)并且是递归的(ADT的构造函数可以采用该ADT的实例,因此可以归纳地建立一个ADT,例如树或列表)。类型系统中ADT的表示必须能够适应这些事实,以下各节将详细介绍。

全局类型变量

为了紧凑而 轻松地表示ADT,可以使用递归ADT定义,将以唯一标识它的全局类型变量的形式为ADT定义一个句柄。每个ADT定义都有一个新的全局类型变量作为句柄,因此指针相等可用于区分不同的ADT名称。
出于relay类型系统的目的,ADT按名称进行区分。这意味着,如果两个ADT具有不同的句柄,则即使它们的所有构造函数在结构上相同,它们也将被视为不同的类型。
因此,ADT定义中的递归就像全局函数的递归一样:构造函数可以在其定义中简单地引用ADT句柄(全局类型变量)。
有关[GlobalTypeVar](https://docs.tvm.ai/api/python/relay/ty.html#tvm.relay.ty.GlobalTypeVar)其定义和文档,请参见。

定义(类型数据)

除名称外,ADT还需要存储用于定义名称的构造函数以及其中使用的任何类型参数。它们类似于全局函数定义存储在模块中。
在使用ADT进行类型检查时,类型系统有时必须使用ADT名称索引到模块中以查找有关构造函数的信息。例如,如果在匹配表达式子句中对构造函数进行模式匹配,则类型检查器必须检查构造函数的签名,以确保为所有绑定变量分配了正确的类型。
有关**TypeData**其定义和文档,请参见。

类型调用

因为ADT定义可以接受类型参数,所以Relay的类型系统将ADT定义视为类型级别的函数 (因为该定义接受类型参数并返回具有这些类型参数的ADT实例的类型)。因此,可以使用类型调用来键入ADT的任何实例,该调用明确列出提供给ADT定义的类型参数。
列出ADT实例的类型参数很重要,因为使用不同的构造函数构建的两个ADT实例,但相同的类型参数具有相同的类型,而具有不同类型参数的两个ADT实例不应视为相同的类型(例如,整数列表不应与浮点张量对列表具有相同的类型)。
类型调用中的“函数”是ADT句柄,并且ADT定义中的每个类型参数必须有一个参数。(不带参数的ADT定义意味着任何实例都不会将类型参数传递给类型调用)。
有关其定义和文档,请参见TypeCall

示例:列表ADT

本小节使用简单列表ADT(作为relay中的默认ADT包括在内)来说明前面各节中描述的构造。其定义如下:
data List {
Nil : () -> List
Cons : (a, List[a]) -> List
}
因此,全局类型变量List是ADT的句柄。模块注释中的列表ADT的类型数据 List采用一个类型参数,并具有两个构造函数, Nil(具有签名 fn
() -> List[a] )和Cons(具有签名fn(a, List[a]) -> List[a])。 通过在构造函数定义中使用全局类型变量List来实现在Cons构造函数中对List的递归调用。
下面是两个使用类型调用给出类型的列表实例:
Cons(1, Cons(2, Nil())) # List[Tensor[(), int32]]
Cons((1, 1), Cons((2, 2), Nil())) # List[(Tensor[(), int32], Tensor[(), int32])]
请注意,该它Nil()可以是任何列表的实例,因为它不使用任何使用类型参数的参数。(不过,对于的任何特定实例Nil(),必须指定type参数。)
这是类型系统拒绝的两个列表,因为类型参数不匹配:
# 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()))