Relay 表达式

Relay 表达式#

Relay IR 是一种纯粹的、面向表达式的语言。以下部分描述了 Relay 中的不同表达式，并详细说明了它们的语义。

数据流和控制片段#

为了将 Relay 与传统基于计算图的中间表示（IRs）进行比较，从数据流和控制片段的角度来考虑 Relay 表达式是有益的。当编写和表达变换时，Relay 程序中仅影响数据流的每个部分都可以被视为传统的计算图。

数据流片段涵盖了不涉及控制流的 Relay 表达式集合。也就是说，程序中仅包含以下结构的任何部分都对应于纯计算图：

变量
Tuple Construction 和 Projection
Let Bindings
Graph Bindings
调用 Operators 和 ADT Constructors

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

If-Then-Else 表达式
ADT Matching 表达式
函数中的递归调用

从计算图的角度来看，函数是一个子图，而函数调用则内联该子图，将其参数替换为子图中具有相应名称的自由变量。因此，如果函数体仅使用数据流结构，则对该函数的调用属于数据流片段；相反，如果函数体包含控制流，则对该函数的调用不属于数据流片段。

变量#

受 LLVM 启发，Relay 在 AST 和文本格式中明确区分了局部变量和全局变量。在文本格式中，全局变量和局部变量通过前缀或符号来区分。全局变量以 @ 为前缀，局部变量以 % 为前缀。

这种明确的区分使得某些优化更容易实现。例如，内联全局定义不需要进行分析：简单地替换定义就足够了。

全局变量#

全局标识符以 @ 符号为前缀，例如 "@global"。全局标识符始终引用包含在全局可见环境中的全局可见定义，称为模块。全局标识符必须是唯一的。

请参阅 GlobalVar 了解其实现和文档。

局部变量#

局部标识符以 % 符号为前缀，例如 "%local"。局部标识符始终引用函数参数或在 let 表达式中绑定的变量，并且分别限定在它出现的函数或绑定的 let 表达式的作用域内。

在下面的代码段中，请注意 %a 被定义了两次。这是允许的，就像在大多数函数式语言中一样；在第二个 let 表达式的作用域内，名称 %a 被“遮蔽”了，这意味着内部作用域中对 %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 中是一等公民，这意味着它们像变量、常量和元组一样是表达式。此外，Relay 中的函数是高阶的，这意味着函数可以作为参数传递给另一个函数，或者由函数返回，因为函数表达式会求值为闭包（参见 Closures 小节），闭包就像张量和元组一样是值。

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

语法#

定义至少包含关键字 fn、空的参数集以及由花括号括起来的体表达式（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 会尝试为用户推断最通用的类型。这一特性被称为泛化：对于没有显式注解的定义，Relay 会尝试根据函数体和调用位置为参数和返回类型分配最通用的类型。

可以使用 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))

闭包#

函数表达式求值为闭包。闭包是表示为局部环境（存储函数体外部定义的所有变量的值）和函数本身这对组合的值。

例如，在下面的示例中，最终结果将是零值张量，因为 %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)

多态性与类型关系#

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

函数还可以被赋予一组类型参数，这些参数可以在调用位置替换为特定类型。具有类型参数的函数是 类型多态 的；它们的返回类型或接受的参数类型可以根据调用位置提供的类型参数而变化。

类型参数按种类分类，并且只能出现在类型签名中其种类适用的部分（例如，种类为 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 关系的约束，这意味着三者都必须是张量，并且它们的形状遵循逐元素广播关系。与算子一样，关系的定义对 Relay 是不透明的，它们是在外部用 C++ 或 Python 实现的。

与 Broadcast 的情况一样，关系用于表达对类型（尤其是张量形状）的复杂约束。所有函数关系必须在所有调用位置都成立；因此，类型检查被视为一个约束求解问题。有关类型关系及其实现的更多详细信息，请参阅 Relay 类型系统文档中的相关部分。

算子#

算子是一种元运算（primitive operation），例如 add 或 conv2d，它们不是在 Relay 语言中定义的。算子在 C++ 的全局算子注册表中声明。许多常见的算子由 TVM 的张量算子库支持。

要注册算子，用户必须提供该算子的实现、其类型以及任何其他所需的元数据。算子注册表是基于列的存储，其中算子是键，因此任何元数据（可能被优化过程引用）都可以注册为新列。

从 Relay 类型系统的角度来看，算子是函数，因此算子可以像任何其他函数一样被调用，并且具有函数类型。特别是，算子类型使用单一类型关系注册（参见关于类型关系的文档），通常是专门为该算子定制的关系。例如，add 算子使用 Broadcast 关系注册，表明 add 的参数必须是张量，并且返回类型的形状取决于其参数的张量。

在漂亮打印 Relay 程序时，算子没有符号前缀（例如 conv2d，flatten）。算子显式地包含在程序中，并且可以通过指针唯一标识。

请注意，常见的算术算子（如 add 和 multiply）可以使用文本格式中相应的算术算子（例如 + 或 *）作为语法糖来编写。

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

ADT 构造器#

Relay 中的代数数据类型（Algebraic data types，简称 ADTs）在单独的概述中有详细描述，并且它们如何集成到类型系统中的内容在这里进行了说明。

在本节中，将简要说明，ADT 构造器被赋予了函数类型，并且应该像函数或算子一样在调用节点中使用。ADT 构造器的定义包括指定它所要构造的 ADT 的名称（全局类型变量）以及构造器期望的参数类型。

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

假设 D 是接受类型参数 a 和 b 的 ADT（代数数据类型）。如果 C1 是 D 的构造函数，并且期望接收两个参数，一个类型为 a，另一个类型为 b，那么 C1 将具有以下类型签名：fun<a, b>(a, b) -> D[a, b]。（关于返回类型中的类型调用的解释，请参见 ADT 概述或ADT类型讨论部分。）如果 D 的另一个构造函数 C2 不接收任何参数，那么它将具有以下类型签名：fun<a, b>() -> D[a, b]；类型参数将始终出现在返回类型中。

一旦被调用，构造函数就会生成 ADT 实例，这是容器，它不仅存储了传递给构造函数的参数值，还存储了构造函数的名称（"tag"）。该标签将用于在 ADT Matching 时解构实例并检索其中的值。

请参阅 Constructor 以获取定义和相关文档。

调用#

在 Relay 中，具有函数类型的表达式是“可调用的”，这意味着它们可以通过函数调用来执行。这些表达式包括任何可以求值为闭包的表达式（例如函数表达式或全局函数）以及 Relay 算子。

调用的语法遵循类 C 语言中使用的语法，如下例所示：

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

当闭包被调用时（参见 Closures），闭包的主体会在存储的环境中（即使用自由变量的存储值）进行评估，并为每个参数添加局部变量绑定；通过评估主体获得的最终值即为调用的返回值。因此，在上面的例子中，调用评估的结果是 22。对于算子的情况，其实现对 Relay 是不透明的，因此结果由注册的 TVM 实现决定。

注意：文本格式目前尚不支持类型参数。

类型多态函数在调用时也可以包含类型参数。在类型检查时，类型参数会被替换为类型变量。如果函数是类型多态的且未提供类型参数，类型推断将尝试在可能的情况下推断出类型参数。以下代码展示了显式和推断类型参数的示例：

// %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], ()>

需要注意的是，函数类型中的所有类型关系在每个调用点都必须成立。具体来说，这意味着关系将根据给定调用点的参数具体类型进行检查。这也是一种多态形式，因为只要满足关系，可能存在多种参数类型和返回类型的有效组合。

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

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

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

模块与全局函数#

Relay 维护了称为“模块”(module)的全局数据结构（在其他函数式编程语言中通常称为“环境”），用于跟踪全局函数的定义。具体来说，模块维护了全局可访问的映射，将全局变量与其所表示的函数表达式关联起来。模块的实用性在于它允许全局函数递归地引用自身或任何其他全局函数（例如，在相互递归的情况下）。

需要注意的是，Relay 的模块类似于在基于计算图的中间表示（IR）中用于跟踪子图的数据结构。

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))
    }
}

这个定义将导致模块中创建条目，将标识符 @ackermann 映射到具有上述参数、返回类型和主体的函数表达式。代码中其他任何对标识符 @ackermann 的引用都可以在模块中查找该标识符，并根据需要替换函数定义。

请参阅 IRModule 以获取模块的定义和文档。

常量#

该节点表示常量张量值（更多详情请参阅 Value）。常量以 NDArray 的形式表示，这使得 Relay 能够利用 TVM 算子进行常量求值。

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

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

元组#

Construction#

元组节点构建了有限（即静态已知大小）的异构数据序列。这些元组与 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 开始索引的。

例如，下面的投影评估结果为 %b：

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

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

Let 绑定#

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

let 绑定包含局部变量、可选的类型注解、值以及可以引用绑定标识符的主体表达式。如果省略了绑定变量的类型注解，Relay 会尝试推断该变量允许的最通用类型。

在 let 表达式中，绑定的变量仅在其主体范围内有效，除非该变量定义了函数表达式。当 let 表达式创建函数时，该变量在其值范围内也有效，以允许递归定义的函数（参见上一小节）。

let 绑定的值是在评估其所依赖的绑定后最终表达式的值。例如，在以下示例中，整个表达式评估为形状为 (10, 10) 且所有元素均为 1 的张量：

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 以获取其定义和文档。

图绑定#

let 绑定创建命名变量，该变量绑定到给定值并限定在后续表达式的范围内。相比之下，图绑定允许通过将表达式（图节点）直接绑定到不限范围的临时变量来显式构建 Relay 程序中的数据流图。每次对变量的引用都对应于数据流图中的一条边。这具有在变量出现的任何地方替换表达式的语义，尽管编译后的程序只会对图节点进行一次评估。

这些绑定允许一种编程风格，与 NNVM 和其他基于数据流图的输入格式已经采用的风格相对应。与 let 绑定相比，变量不限范围的事实提供了一些评估顺序的灵活性，尽管这也可能在程序中引入一些歧义。

注意：图绑定目前尚未被文本格式解析。

在 Relay 的文本格式中，图绑定可以写成如下形式（注意缺少 let 关键字和分号）：

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

与 let 绑定不同，图绑定在 Relay 中不作为 AST 节点表示，而是作为引用其 AST 节点值的元变量。例如，像上面这样的程序可以通过在 Relay 的 Python 前端中将 Python变量 设置为相应的 Relay AST 节点并重复使用这些变量来构建，如下所示（使用相应 API 绑定的 C++ 程序可以实现相同的功能）：

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

出于开发目的和实现某些优化，Relay 包含了在数据流图和 A-normal 形式的 let 绑定程序之间进行转换的步骤，这种形式被许多来自函数式编程社区的编译器优化所采用（有关 A-normal 形式的介绍，请参见 Matt Might的 "A-Normalization: Why and How"）。

If-Then-Else#

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

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

由于 if-then-else 分支是表达式，它们可以内联出现在任何其他表达式可能出现的地方，就像类 C 语言中的三元算子调用一样。如果条件值评估为 True，则 if-then-else 表达式评估为 “then” 分支的值；如果条件值评估为 False，则评估为 “else” 分支的值。

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

ADT 匹配#

代数数据类型（ADT）的实例，如 ADT概述中所述，是存储传递给用于创建它们的构造函数的参数的容器，并由构造函数名称标记。

Relay 中的匹配表达式允许根据构造函数标签检索存储在 ADT 实例中的值（“解构”它）。匹配表达式的行为类似于 C 风格的 switch 语句，根据被解构值的类型的不同可能构造函数进行分支。正如 ADT 概述中详细说明的那样，匹配表达式能够进行比简单按构造函数拆分更一般的模式匹配：任何嵌套在实例内部的ADT实例（例如，列表的列表）都可以与外部实例同时解构，而实例的不同字段可以绑定到变量。（有关ADT模式匹配的详细描述，请参见此部分。）

匹配表达式使用输入值（表达式）和子句列表定义，每个子句由模式和表达式组成。执行时，第一个 模式与查询值结构匹配的子句将被执行；子句表达式被评估并返回。

例如，假设有用于自然数的 ADT：

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

那么以下函数将从传入的 nat 中减去一：

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

以下函数如果其参数至少为二，则从中减去二，否则返回参数，使用嵌套构造函数模式：

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

如前所述，匹配子句的顺序是相关的。在下面的示例中，第一个子句将始终匹配，因此其下面的子句永远不会运行：

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

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

临时表达式#

在 Relay 中，程序转换（passes）可能需要向程序 AST 中插入临时状态以指导进一步的转换。为此，TempExpr 节点作为工具提供给开发者使用；继承自 TempExpr 的节点不能直接出现在用户提供的代码中，但可以在 pass 中插入。理想情况下，在 pass 完成之前应消除任何创建的 TempExpr，因为 TempExpr 仅存储内部状态，没有自身的语义。

有关在 pass 中使用 TempExpr 的示例，请参见 src/relay/transforms/fold_scale_axis.cc，该文件使用 TempExpr 节点在 pass 尝试将这些参数折叠到卷积的权重中时存储有关缩放参数的信息。

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