Relay 表达式#

Relay IR 是纯粹的、面向表达式的语言。下面的部分描述了在 Relay 中不同的表达式,并给出了其语义的细节。

数据流和控制片段#

为了将 Relay 与传统的基于计算图的 IR 相比,在 dataflow 和控制片段方面考虑 Relay 表达式是很有用的。在写入和表达变换时,只影响 dataflow 的 Relay 程序的每一部分都可以被视为传统的计算图。

dataflow 片段覆盖了不涉及控制流的 Relay 表达式集。也就是说,包含以下结构的程序的任何部分都对应于纯计算图:

控制流表达式允许 graph 拓扑根据先前执行的表达式的值进行更改。Relay 的控制片段包括以下结构:

从计算图的角度来看,函数是子图,函数 call 在子图中,将其参数替换为带有相应名称的子图中的自由变量。因此,如果函数体只使用 dataflow 构造器,那么调用该函数就在 dataflow 片段中;相反,如果函数体包含控制流,那么调用该函数就不属于 dataflow 片段。

变量#

由 LLVM 启发,Relay 明显地区分了 AST 和文本格式之间的局部变量和全局变量。在文本格式中,全局和局部变量由前缀或 sigils 区分。全局变量前缀为 @ 和局部变量前缀为 %

这种显式的区分使某些优化更容易实现。例如,在全局定义中不需要分析:简单地替换定义就足够了。

全局变量#

全局标识符是由前缀 @ 符号化,如 “@global”。全局标识符总是引用在全局可见环境中包含的全局可见定义,称为 module。全局标识符必须是唯一的。

参阅 GlobalVar 的实现和文档。

局部变量#

局部标识符由 % 符号化,如 “%local”。局部标识符总是引用函数参数或被 let 表达式绑定的变量,并将作用于它出现的函数或被 let 表达式绑定之处。

在下面的代码段中,请注意 %a 被定义了两次。这是允许的,就像在大多数函数语言中一样;在第二个 let 表达式的作用域中,名称 %a 是 “shadowed”,这意味着在内部作用域中对 %a 的所有引用后面的定义,而在外部作用域中对 %a 的引用继续引用第一个定义。

let %a = 1;
let %b = 2 * %a;  // %b = 2
let %a = %a + %a; // %a = 2. %a is shadowed
%a + %b           // has value 2 + 2 = 4

(请注意,在 Relay 的实现中,局部变量的每个定义都会创建新的 Var,因此,被隐藏的局部变量尽管与外部作用域中的变量具有相同的名称,但将是不同的对象。这允许通过指针标识比较局部变量,同时知道相同的局部变量对象对应于不同的绑定站点。)

阅读 Var 的实现和文档。

函数#

Relay 中的函数的作用类似于其他编程语言中的过程或函数,并用于推广命名子图的概念。

函数是 Relay 中的 first class,这意味着它们是表达式,就像变量、常量和元组一样。此外,Relay 中的函数是高阶的,这意味着函数可以作为参数传递给函数,也可以由函数返回,因为函数表达式计算为闭包(参见 Closures 小节),这是张量和元组等的值。

阅读 Function 了解函数节点的定义和文档。

语法#

定义至少包括关键字 fn,一个空的参数集,以及由花括号包含的 body 表达式(Expr)。

fn() { body }

定义可以包含任意数量的参数。例如,调用 add 算子的简单函数:

fn(%x, %y) { add(%x, %y) }

注意,在函数体中,参数是局部变量,就像 let 表达式中绑定的那些参数一样。

还可以在函数上注解显式类型。例如,可以将上面的函数限制为只能处理特定类型:

fn(%x : Tensor[(10, 10), float32], %y : Tensor[(10, 10), float32])
           -> Tensor[(10, 10), float32] {
    add(%x, %y)
}

上面的函数只接受 Tensor[(10, 10), float32] 类型的参数,并返回 Tensor[(10, 10), float32] 类型的值。函数参数只是局部变量(LocalVar),可选地使用类型进行注解,写为 %x : T

当类型信息被省略时,Relay 尝试为用户推断出最一般的类型。这个属性被称为泛化(generalization):对于没有显式注解的定义,Relay 试图根据函数体和调用站点(sites)为参数和返回类型分配最通用的类型。

递归函数表达式可以使用 let 绑定定义,如下所示:

let %fact = fn(%x : Tensor[(10, 10), float32]) -> Tensor[(10, 10), float32] {
    if (%x == Constant(0, (10, 10), float32)) {
        Constant(1, (10, 10), float32)
    } else {
        %x * %fact(%x - Constant(1, (10, 10), float32))
    }
};
%fact(Constant(10, (10, 10), float32))

闭包#

函数表达式的计算结果为闭包。闭包是表示为一对局部环境(存储在函数体范围之外定义的所有变量的值)和函数本身的值。

例如,在下面的例子中,最终的结果将是值为 0 的张量,因为 %f 的闭包将 %x 的值存储在 %f 定义的指针上。

let %g = fn() {
  let %x = Constant(0, (10, 10), float32);
  // %x is a free variable in the below function
  fn(%y) { %y * %x }
};
// the %x in %g's body is not in scope anymore
// %f is a closure where %x maps to Constant(0, (10, 10), float32)
let %f = %g();
let %x = Constant(1, (10, 10), float32);
%f(%x) // evaluates to Constant(0, (10, 10), float32)

多态性与类型关系#

注意:文本格式中还不支持类型参数语法。

函数也可以给定一组类型参数,这些参数可以在 call sites 替换特定类型。带类型参数的函数是 类型多态的 (type polymorphic);它们的返回类型或它们将接受的参数类型可以根据调用点给出的类型参数而变化。

类型参数是按 kind 分类的,并且只能出现在类型签名中适合其种类的部分(例如,类型参数 Shape 的种类只能出现在张量类型中需要 shape 的地方);要获得完整的讨论,请参阅关于 类型参数 的文档。

例如,可以为任何 Relay 类型定义多态标识函数如下:

fn<t : Type>(%x : t) -> t {
    %x
}

下面的定义也是多态的,但将其参数限制为张量类型:

fn<s : Shape, bt : BaseType>(%x : Tensor[s, bt]) {
    %x
}

注意,返回类型被省略,并且将被推断。

注意:在文本格式中还不支持 “where” 语法。

函数也可以受一种或多种类型关系的支配,例如如下所示:

fn(%x, %y) where Broadcast { add(%x, %y) }

在上面的定义中,%x%y 的类型以及返回类型都服从 Broadcast 关系,这意味着这三个都必须是张量,它们的形状遵循 elementwise Broadcast 关系。与算子一样,关系的定义对 Relay 不是透明的,而是在外部用 C++ 或 Python 实现的。

Broadcast 的例子中,关系被用来表达对类型(尤其是张量形状)的复杂约束。所有函数关系必须在所有 call sites 保持不变;因此,类型检查被视为 constraint-solving 问题。要了解更多关于类型关系及其实现的细节,请参见 Relay 的类型系统 文档中的相关章节。

算子#

算子是原子运算(primitive operation),如 addconv2d,在 Relay 语言中没有定义。在 C++ 中,算子在全局算子注册表(registry)中声明。TVM 的 Tensor Operator Inventory 支持许多常见的算子。

要注册算子,用户必须提供算子的实现、类型和任何其他所需的元数据。算子注册是基于列的存储,其中算子是键,因此任何元数据(可能通过优化 passes 引用)都可以注册为新列(column)。

从 Relay 类型系统的角度来看,算子是函数,所以算子可以像任何其他函数一样调用,并且具有函数类型。特别地,算子类型是使用单个类型关系注册的(请参阅 关于类型关系 的文档),通常是特定到该算子的关系。例如,add 算子被注册为 Broadcast 关系,表明 add 的参数必须是张量,返回类型是张量,其形状取决于其参数的形状。

当精细打印 Relay 程序时,算子不需要 sigil(例如 conv2d, flatten)。算子显式地包含在程序中,并且可以通过指针唯一地识别。

注意,常见的算术算子(如 addmultiply)可以使用文本格式中相应的算术运算符(如 +*)作为语法糖来编写。

请参阅 Op 获取算子节点的定义和文档,演示了注册算子元数据的基础设施。op 中的其他文件提供了句柄,用于生成对各种预注册算子的回调。向 Relay 添加算子的教程 展示了如何向语言中添加更多的算子。

ADT 结构#

Relay 中的代数数据类型(Algebraic data types,简称 ADTs)将在 单独的概述 中详细描述,并在 这里 描述它们与类型系统的集成。

在本节中,将简单地注意到 ADT 构造函数被赋予了函数类型,应该像函数或算子一样在 call 节点内部使用。ADT 构造函数是通过给出它构造的 ADT 的名称(全局类型变量)和构造函数所需参数的类型来定义的。

如果 ADT 定义包含类型变量,那么这些类型变量可能出现在构造函数中。构造函数不能包含任何其他类型变量。

假设 D 是接受类型参数 ab 的 ADT。如果 C1D 的构造函数,需要两个参数,一个类型 a,一个类型 b,那么 C1 有以下类型签名 fun<a, b>(a, b) -> D[a, b]。(关于返回类型中类型调用的解释,请参阅 ADT 概述或 ADT 类型的讨论。)如果 D 的另一个构造函数 C2 不接受参数,那么它有以下类型签名:fun<a, b>() -> D[a, b];类型参数将始终出现在 return 类型中。

一旦被回调,构造函数将生成 ADT 实例,它是容器,将参数的值以及构造函数的名称(”tag”)存储到构造函数中。当 ADT 匹配 时,该 tag 将用于分解实例和检索值。

参阅 Constructor 了解定义和文档。

Call#

在 Relay 中具有函数类型的表达式是 “callable”,这意味着它们可以通过函数调用来回调。它们由任何求值为闭包的表达式(即函数表达式或全局函数)和 Relay 算子组成。

调用的语法遵循类 C 语言的语法,如下例所示:

let %c = 1;
let %f = fn(%x : Tensor[(), float32], %y : Tensor[(), float32]) { %x + %y + %c };
%f(10, 11)

当调用闭包(参见 Closures )时,闭包的主体将在存储环境中求值(即使用存储的值作为自由变量),并为每个参数添加局部变量绑定;通过求值 body 获得的最后的值是回调的返回值。因此,在上面的例子中,调用的计算结果为 22。在算子的情况下,实现对 Relay 是不透明的,因此结果由注册的 TVM 实现决定。

注意:文本格式中还不支持类型参数。

类型多态函数还可以在 call site 包含类型参数。类型参数在类型检查时被替换为类型参数。如果函数是类型多态的,并且没有给出类型参数,则如果可能,类型推断将尝试推断类型参数。下面的代码给出了显式和隐式类型参数的示例:

// %f : fn<a : Type, b : Type, c : Type>(a, b) -> c
let %x1 = %f<Tensor[(), bool], Tensor[(), bool], Tensor[(), bool)]>(True, False);
// %x1 is of type Tensor[(), bool]
let %x2 : () = %f(%x1, %x1)
// the type arguments in the second call are inferred to be <Tensor[(), bool], Tensor[(), bool], ()>

注意,函数类型中的所有类型关系必须在每个 call site 保持。具体来说,这意味着将根据给定 call site 上特定类型的参数检查关系。这也是多态的一种形式,因为只要满足关系,参数类型和返回类型可以有多个有效赋值。

例如,如果有函数 %f,它接受张量参数,并且具有 Broadcast 关系,那么下面调用中的参数可以有许多不同的形状来满足类型注解:

let %x : Tensor[(100, 100, 100), float32] = %f(%a, %b);
%x

参阅 Call 了解定义及其文档。

Module 和全局函数#

Relay 保留称为 “module” 的全局数据结构(在其他函数式编程语言中通常称为 “environment”),以跟踪全局函数的定义。特别地,该模块保持全局变量到它们所表示的函数表达式的全局可访问映射。模块的实用之处在于,它允许全局函数递归地引用它们自己或任何其他全局函数(例如,在 mutual 递归中)。

注意:Relay 的模块类似于用于跟踪基于计算图的 IRs 中的子图的数据结构。

Relay 中的全局函数的行为与 Functions 中定义的函数表达式相同,但是在文本格式中有语法糖来将它们的定义输入到模块中。也就是说,全局函数定义包含全局标识符,并且允许在函数体中递归地引用该标识符,如下面的例子所示:

def @ackermann(%m : Tensor[(), int32], %n : Tensor[(), int32]) -> Tensor[(), int32] {
    if (%m == 0) {
        %n + 1
    } else if (%m > 0 && %n == 0) {
        @ackermann(%m - 1, 1)
    } else {
        @ackermann(%m - 1, @ackermann(%m, %n - 1))
    }
}

这个定义将产生模块(module)条目,将标识符 @ackermann 映射到带有上述参数、返回类型和 body 的函数表达式。代码中其他地方对标识符 @ackermann 的任何引用都可以在模块中查找标识符,并根据需要替换函数定义。

参阅 IRModule 定义和文档了解 module 更多细节。

Constant#

该节点表示常量张量值(详见 Value)。常量被表示为 NDArray,允许 Relay 使用 TVM 算子来计算常量。

这个节点还可以表示标量常数,因为标量是形状为 () 的张量。因此,在文本格式中,数值和布尔字面值是将张量类型编码为零阶形状的常量的语法糖。

阅读 Constant 定义和文档了解细节。

元组#

构造#

元组节点构建有限(即静态已知大小)的异构数据序列。这些元组与 Python 非常匹配,它们的固定长度允许有效地投影其成员。

fn(%a : Tensor[(10, 10), float32], %b : float32, %c : Tensor[(100, 100), float32]) {
    let %tup = (%a, %b);     // type: (Tensor[(10, 10), float32], float32)
    ((%tup.0 + %tup.1), %c)  // type: (Tensor[(10, 10), float32], Tensor[(100, 100), float32])
}

阅读 Tuple 了解细节。

投影#

元组必须用整数常量进行索引,以便提取元组中的特定成员。投影是 0-indexed。

例如,下面的投影计算结果为 %b

(%a, %b, %c).1

阅读 TupleGetItem 了解更多细节。

let 绑定#

let 绑定是不可变的局部变量绑定,允许用户将表达式绑定到名称。

let 绑定包含局部变量、可选类型注解、值和可以引用绑定标识符的 body 表达式。如果省略了绑定变量上的类型注释,Relay 将尝试推断该变量允许的最通用类型。

let 表达式中的绑定变量只作用在其 body 作用域内,除非该变量定义了函数表达式。当 let 表达式创建函数时,该变量的值也在范围内,以允许递归定义函数(请参阅前一小节)。

let 绑定的值是计算它所依赖的绑定后的最后一个表达式的值。例如,在下面的例子中,整个表达式计算为形状为 (10, 10) 的张量,其中所有元素为 2:

let %x : Tensor[(10, 10), float32] = Constant(1, (10, 10), float32);
%x + %x

一系列 let 绑定可以看作是数据流图,其中绑定是由绑定变量连接的一系列子图。由于这些绑定序列是纯的,因此可以安全地重新排序两个都不依赖于另一个的绑定。例如,下面的第一个和第二个 let 绑定可以按任意顺序进行计算,因为它们之间都没有数据流依赖关系:

let %x = %a + %b;
let %y = %c + %d;
%x * %y

阅读 Let 了解更多信息。

Graph 绑定#

let 绑定创建了命名变量(named variable),该变量绑定到给定值并作用域为后续表达式。相比之下,图绑定允许在 Relay 程序中显式地构造数据流图,方法是将表达式(图节点)直接绑定到没有作用域的临时变量。每个对变量的引用对应数据流图中的一条边。它的语义是在变量出现的地方替换表达式,即使图节点只会被编译后的程序计算一次。

这些绑定允许一种与 NNVM 和其他基于数据流图的输入格式所采用的编程风格相对应的编程风格。与 let 绑定相比,变量没有限定作用域的事实在计算顺序上提供了一些灵活性,尽管这也会在程序中引入一些模糊性(开发人员对 Relay IR 的介绍 包含了关于这种细微差别的更详细的讨论)。

注意:Graph 当前未被文本格式解析。

在 Relay 的文本格式中,graph 绑定可以这样写(注意缺少 let 关键字和分号):

%1 = %a + %b
%2 = %1 + %1
%2 * %2

与 let 绑定不同,图绑定在 Relay 中不表示为 AST 节点,而是表示为引用其 AST 节点值的元变量(meta-variable)。例如,可以在 Relay 的 Python 前端通过设置 Python 变量 等于相应的 Relay AST 节点并重复使用变量来构造类似于上面的程序,如下所示(使用相应的 API 绑定的 C++ 程序可以完成相同的事情):

sum1 = relay.add(a, b)
sum2 = relay.add(sum1, sum1)
relay.multiply(sum2, sum2)

为了开发目的和实现某些优化,Relay 包括了在使用图绑定定义的数据流图和使用 let 绑定定义的 A-范式程序之间进行变换的 passes,它被来自函数式编程社区的许多编译器优化所采用(参见 Matt Might 的 “A-Normalization: Why and How” 了解 A-范式形式)。

If-Then-Else#

Relay 有简单的 if-then-else 表达式,允许程序在 bool 类型的单个值上进行分支,即布尔值的零阶张量(Tensor[(), bool])。

if (%t == %u) {
    %t
} else {
    %u
}

由于 if-then-else 分支是表达式,因此它们可以内联出现在任何需要其他表达式的地方,就像在类 C 语言中调用三元算子一样。if-then-else 表达式在条件值为 True 时计算为 “then” 分支的值,在条件值为 False 时计算为 “else” 分支的值。

阅读 If 了解细节。

ADT 匹配#

ADT 概述 中所讨论的,代数数据类型(ADT)的实例是存储传递给用于创建它们的构造函数的参数的容器,由构造函数名称标记(tagged)。

Relay 中的 match 表达式允许根据它们的构造函数 tag 检索存储在 ADT 实例中的值(解构(”deconstructing”)它)。match 表达式的行为类似于 C 风格的 switch 语句,对于被解构的值的类型在不同的可能构造函数上进行分支。正如 ADT 概述的细节所示,match 表达式能够进行更通用的模式匹配,而不仅仅是简单地由构造函数 splitting:任何嵌套在实例内的 ADT 实例(例如,list 的列表)可以与外部实例同时分解,而实例的不同字段可以绑定到变量。(请参阅 ADT 模式匹配 的详细描述。)

match 表达式使用 input 值(表达式)和一列子句定义,每个子句由模式(pattern)和表达式(expression)组成。当执行时,执行第一个模式匹配查询值结构的子句;子句表达式被求值并返回。

例如,假设有用于自然数的 ADT:

data Nat {
  Z : () -> Nat # zero
  S : (Nat) -> Nat # successor (+1) to a nat
}

然后下面的函数从传递(passes)的 nat 中减去 1:

fn(%v: Nat[]) -> Nat[] {
  match(%v) {
    case Z() { Z() }
    case S(%n) { %n } # the variable %n is bound in the scope of this clause
  }
}

下面的函数如果参数至少为 2,则从参数中减去 2,否则使用嵌套构造函数模式返回参数:

fn(%v : Nat[]) -> Nat[] {
  match(%v) {
     case S(S(%n)) { %n }
     # wildcard pattern: matches all cases not matched already
     case _ { %v }
  }
}

如上所述,match 子句的排序是相关的。在下面的例子中,第一个子句将始终匹配,因此它下面的子句永远不能运行:

fn(%v : Nat[]) -> Nat[] {
  match(%v) {
    case _ { %v }
    case S(S(%n)) { S(%n) }
    case S(%n) { %n }
    case Z() { S(Z()) }
  }
}

阅读 Match 了解更多细节。

TempExprs#

Relay 中的程序变换(passes)可能需要在程序 AST 中插入临时状态来指导进一步的变换。为此,将 TempExpr 节点作为实用工具提供给开发人员;继承自 TempExpr 的节点不能直接出现在用户提供的代码中,但可以插入到 pass 中。在 pass 中创建的任何 TempExpr 都应该在 pass 完成之前被删除,因为 TempExpr 只存储内部状态,并且没有自己的语义。

关于在 pass 中使用 TempExpr 的例子,请参见 src/relay/transforms/fold_scale_axis.cc,它使用 TempExpr 节点存储关于 scaling 参数的信息,因为 pass 试图将这些参数折叠到卷积的 weights 中。

参阅 TempExpr 了解更多细节。