【TVM帮助文档学习】Relay表达式

本文翻译自Expressions in Relay — tvm 0.9.dev0 documentation

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

数据流和控制片段

为了比较Relay和传统的基于计算图形的IR,可以从数据流和控制片段的角度看Relay表达式。在编写和表示转换时,如果Relay的某个片段仅仅只包含影响数据流的表达式,那么这个片段可以看作是一个传统的计算图。
数据流片段涵盖了不涉及控制流的Relay表达式集。也就是说,程序的任何只包含以下结构的部分,对应于一个纯计算图:

  • 变量
  • Tuple构造和投影
  • Let Bindings
  • Graph Bindings
  • 算子和ADT构造调用

控制流表达式使graph拓扑结构随之前执行的(条件)表达式值改变。Relay的控制片段包括以下结构:

  • if-then-else表达式
  • ADT匹配表达式
  • 函数中的递归调用

从计算图的角度来看,函数是一个子图,函数调用内联子图,使用子图中对应名字的自由变量替换函数参数。因此,如果函数体只使用数据流结构,那么对该函数的调用就在数据流片段中。相反,如果函数体包含控制流,则对该函数的调用就不是数据流片段的一部分。

变量

受LLVM的启发,Relay在AST和文本格式中明确区分了本地变量和全局变量。在文本格式中,全局变量和局部变量由前缀或符号区分。全局变量以@作为前缀,局部变量以%作为前缀。
这种显式的区别使得某些优化更容易实现。例如,内联一个全局定义将不需要分析,简单地替换定义即可。

全局变量

全局标识符以@符号作为前缀,例如“@global”。全局标识符常见于包含在全局可见环境(通常是module)中的全局可见定义。全局标识符必须唯一

本地变量

本地标识符以%符号作为前缀,例如“%local”。本地标识符常见于函数参数,或者let表达式中绑定的变量,并且作用域限制在定义该变量的函数中,或者绑定的let表达式中。
下面的代码段中,注意%a定义了两次。在大多数函数式语言中这是允许的。在第三个let表达式作用域中,第一个let表达式定义的%a被隐藏了,也就是在内部作用域中对%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,所以对一个隐藏的局部变量而言,尽管与外部作用域中的变量有相同的名称,仍将是一个不同的对象。这允许通过指针标识来比较不同绑定点的相同局部变量。)

函数

Relay中的函数类似于其他编程语言中的过程或函数,用于扩大命名子图的概念。
函数是Relay中的重要类,它与变量、常量和元组一样都是表达式。此外,Relay中的函数是高阶的,这意味着函数可以作为参数传递给函数,也可以由函数返回,这是因为函数表达式计算结果是闭包,值为张量和元组之类。

语法

一个最简单的函数定义由关键字fn、花括号和作为函数体的一条表达式组成,函数参数为空。

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 binding定义一个递归函数表达式,例如:

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的类型参数只能出现在张量类型中用到shape的地方);有关完整的讨论,请参阅有关类型参数的文档。
例如,可以将多态函数定义为支持任何Relay类型,如下所示:

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

下面的定义也是多态的,但是限定了参数必须是tensor类型

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的情况一样,关系用于表示类型(特别是张量形状)的复杂约束。所有函数关系必须在所有调用点保持;因此,类型检查被视为一个约束解决问题

算子

算子表示一种基础的运算,如add或conv2d,它们不是在Relay语言中定义的。算子在C++的全局操作符注册表中声明。许多常用的算子都被TVM的张量算子库支持。
用户注册算子时,必须提供该算子的实现、类型和任何其他所需的元数据。算子注册表采用基于列的方式存储,其中算子是键,因此任何元数据(在执行优化pass时可能被引用)都可以注册为一个新列。
从Relay的类型系统的角度来看,算子就是函数,因此可以像调用任何其他函数一样调用算子,并具有函数类型特性。特别是,算子类型是使用单一类型关系注册的(请参阅有关类型关系的文档),通常(在算子注册式)一个特定的relation会被指定给算子。例如,add算子注册时指定了Broadcast规则,表明add的参数必须是张量,并且返回类型是一个张量,其shape取决于其参数的shape。
在按照pretty-printing原则(即遵循语言排版格式)打印Relay程序时,算子是没有特定标记符号的(例如conv2d、flatten)。算子显式地包含在程序中,并使用指针惟一地标识。
注意,常见的算术算子,如加和乘,可以直接使用文本格式中的相应算术操作符(例如+或*)表示

ADT构造函数

ADT构造函数被赋予了函数类型,应该像函数或算子一样被调用。定义ADT构造函数时要给出它构造的ADT的名称(一个全局类型变量)和构造函数预期参数的类型。
如果ADT定义包含类型变量,那么这些类型变量可能会出现在构造函数中。构造函数不能包含任何其他类型变量。
假设D是一个ADT,它接受类型参数a和b。如果C1是D的构造函数,并且需要两个参数,一个类型为a,一个类型为b,那么C1具有以下类型签名: fun<a,b>(a,b)->D[a,b]。(关于返回类型中类型调用的解释,请参阅ADT概述或ADT类型的讨论。)如果D有另外一个构造函数C2,并且C2没有参数,那么它具有以下类型签名: fun<a, b>() -> D[a, b];类型参数将始终出现在返回类型中。
调用构造函数后,构造函数生成一个ADT实例,该实例是一个容器,存储构造函数的参数值以及构造函数的名称(“tag”)。标记将用于实例的析构,以及在ADT匹配时值的检索。

Call

在Relay中具有函数类型的表达式是“可调用的”,这意味着可以通过函数调用调用它们。这类表达式包括任何计算结果为闭包的表达式(即函数表达式或全局函数)和Relay操作符。
调用的语法遵循类c语言中使用的语法,如下例所示:

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

 当一个闭包被调用时(请参阅闭包),闭包的主体会在存储环境中计算(例如,使用自由变量存储的值),并为每个参数添加局部变量绑定(使用实参赋值给形参?);通过计算函数体得到的最终值是调用的返回值。因此,在上面的示例中,调用的计算结果为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关系,那么下面调用中的参数可以有多种不同的shape满足要求的类型(这些shape经过Broadcast后满足要求):

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

 Module and Global Functions

Relay有一种全局的数据结构称为“module”(在其他函数式编程语言中通常称为“环境”),用来记录全局函数的定义。特别地,module中有全局变量到它们所表示的函数表达式的全局可访问的映射。module的实用性在于,它允许全局函数递归地调用它们自己或任何其他全局函数(比如在相互递归中)。
注意:Relay的Module类似于一种用于跟踪计算图IR中子图的数据结构
Relay中全局函数的行为与Functions中定义的函数表达式一致,但是在文本格式中使用语法糖来将它们的定义输入到module中。也就是说,全局函数定义包含一个全局标识符,并允许在函数体中递归地引用该标识符,如下例所示:

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标识符的引用都可以在模块中查找标识符,并根据需要替换成函数定义。

常量

常量节点表示一个常量张量值(更多细节请参阅value)。常数使用NDArray表示,允许Relay利用TVM算子对常量求值。
常量节点也可以表示标量常数,因为标量是形状为()的张量。因此,在文本格式中,数字和布尔值是shape rank为0的张量。

Tuple

构造

元组节点构建一个有限的(即静态已知大小的)异构数据序列。这些元组与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

Let Bindings

let binding是一种不可变的局部变量绑定,允许用户将表达式绑定到名称。
let绑定包含一个局部变量、一个可选的类型说明、一个值和一个表达式主体,表达式主体可能引用绑定的标识符。如果忽略了绑定变量的类型说明,Relay将尝试推断出该变量允许的最通用类型。
let表达式中的绑定变量作用域的仅限于其主体,除非该变量定义了函数表达式。当let表达式创建一个函数时,变量的作用域将包括let表达式的value部分,以支持递归函数。
let绑定的值是计算它所依赖的绑定后的最终表达式的值。例如,在下面的例子中,整个表达式的结果是一个形状为(10,10)张量,其中所有元素都是2:

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

一组let bindings表达式在一起可以被看作是一个数据流图,各绑定表达式可以看作子图,子图之间由绑定的变量连接。由于这些绑定序列是清晰明确的,因此两个彼此不依赖的绑定可以安全地重新排序。例如,下面的第一个和第二个let绑定可以按任意顺序求值,因为这两个绑定之间都没有数据流依赖关系:

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

Graph Bindings

let binding创建一个命名变量,该变量被绑定到给定的值并被限定在后续表达式的范围。相比之下,graph binding允许在Relay程序中显式地构造数据流图,方法是将表达式(图节点)直接绑定到没有限定作用域的临时变量。每个对变量的引用都对应于数据流图中的一条边。其语义是在变量出现的任何地方替换表达式,即使已编译的程序只对该图形节点求值一次。
这些绑定与NNVM以及其他基于数据流图的输入格式所使用的编程风格一样。因为变量没有限定作用域,所以相比let binding,在求值顺序上更具灵活性,当然这也会在程序中引入一些歧义(开发人员对Relay IR的介绍中包含了对这种细微差别的更详细的讨论)。
注意:当前文本格式不支持图形绑定的解析。
在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包含了在使用图形绑定定义的数据流图和使用let绑定的a范式程序之间转换的pass,该范式被函数式编程社区的许多编译器优化所采用

If-Then-Else

Relay有一个简单的if-then-else表达式,它允许程序在bool类型的单个值上进行分支

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

因为if-then-else分支是表达式,所以它们可以内联地出现在任何需要其他表达式的地方,类似c语言中调用三元操作符。如果条件值为True,则if-then-else表达式计算为" then "分支的值;如果条件值为False,则计算为" else "分支的值。

ADT Matching

如ADT概述中所讨论的,代数数据类型(ADT)的实例是存储传递给ADT构造函数的参数的容器,由构造函数名称标记。
Relay中的匹配表达式允许根据构造函数标记检索存储在ADT实例中的值。match表达式的行为类似于c语言的switch语句,针对被解构的值的类型,在各可能构造函数上进行分支。如前所述,匹配表达式能够进行更一般的模式匹配,而不仅仅是通过构造函数进行拆分:嵌套在实例中的任何ADT实例(例如列表列表)都可以与外部实例同时解构,而实例的不同字段可以绑定到变量
匹配表达式是使用输入值(一个表达式)和一组子句来定义的,每个子句包含一个模式和一个表达式。执行时,第一个能匹配查询值结构的子句被;计算子句表达式并返回值。
例如,假设我们有一个自然数的ADT:

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

然后下面的函数从传递的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()) }
  }
}

TempExprs

Relay中的程序转换(passes)可能需要将临时状态插入到程序AST中,以指导进一步的转换。TempExpr节点是作为一个实用工具提供给开发人员的;从TempExpr继承的节点不能直接出现在用户提供的代码中,但可以在pass中插入。在pass中创建的任何TempExpr都应该在pass完成之前消除,因为TempExpr只存储内部状态,没有自己的语义。
关于在pass中使用TempExpr的例子,请参见src/relay/transforms/fold_scale_axis.cc,在pass试图将这些参数折叠成卷积的权重时,它使用TempExpr节点来存储关于缩放参数的信息。

  • 2
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值