TVM：Python 创建 Relay 表达式

快速了解 Relay¶

Relay IR 是纯粹的、面向表达式的语言。

从计算图的角度来看，函数（{class}tvm.relay.function.Function ）是计算图的子图，函数调用在子图中，将其参数替换为带有相应名称的子图中的自由变量。

Relay 明显地区分了 AST 和文本格式之间的局部变量（{class}tvm.ir.expr.Var 使用 % 标识）和全局变量（{class}tvm.ir.expr.GlobalVar 使用 @ 标识）。

• 全局标识符总是引用在全局可见环境中包含的全局可见定义，称为 模块 （module）。全局标识符必须是唯一的。
• 局部标识符总是引用函数参数或被 let （{class}tvm.relay.expr.Let） 表达式绑定的变量，并将作用于它出现的函数或被 let 表达式绑定之处。

局部变量¶

局部变量可用于声明函数的输入参数或中间变量。可由 tvm.relay.Var(name_hint, type_annotation=None) 创建。其中

• name_hint 指定了局部变量的名字。
• type_annotation 用于局部变量的类型注解。

 

In [ ]:

from tvm import relay

x = relay.Var("x") # 创建局部变量 x
x

 

Out[ ]:

Var(x)

 

 

可以查看文本表示：

 

In [ ]:

print(x)

 

 

free_var %x;
%x

 

 

如果想要声明给定 dtype 和形状已知的张量的类型，可以指定 type_annotation 参数创建：

 

In [ ]:

type_annotation = relay.TensorType(shape=(5, 5),
                                   dtype="float32")
x = relay.Var("x", type_annotation)
x

 

Out[ ]:

Var(x, ty=TensorType([5, 5], float32))

 

 

查看文本格式：

 

In [ ]:

print(x)

 

 

free_var %x: Tensor[(5, 5), float32];
%x

 

 

还有一种便捷函数：tvm.ir.expr.var(name_hint, type_annotation=None, shape=None, dtype="float32")：

 

 

创建变量的四种等效方式：

 

In [ ]:

x = relay.Var("x", relay.TensorType([1, 2]))
x = relay.var("x", relay.TensorType([1, 2]))
x = relay.var("x", shape=[1, 2])
x = relay.var("x", shape=[1, 2], dtype="float32")

 

 

同样，下面两列也是等效的：

 

In [ ]:

y = relay.var("x", "float32")
y = relay.var("x", shape=(), dtype="float32")

 

 

函数¶

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

可以直接定义函数：

 

In [ ]:

a, b = [relay.var(name) for name in "ab"]
add_op = a + b
add_func = relay.Function([a, b], add_op)

 

 

此函数的文本形式：

 

In [ ]:

add_func

 

Out[ ]:

fn (%a, %b) {
  add(%a, %b)
}

 

 

也可以使用 Python 函数回调的形式：

 

In [ ]:

def add(a, b):
    add_op = a + b
    return relay.Function([a, b], add_op)

add

 

Out[ ]:

<function __main__.add(a, b)>

 

 

想要使用，需要：

 

In [ ]:

a, b = [relay.var(name) for name in "ab"]
add_func = add(a, b)
add_func

 

Out[ ]:

fn (%a, %b) {
  add(%a, %b)
}

 

 

也可以添加变量注解：

 

In [ ]:

type_annotation = relay.TensorType(shape=(5, 5),
                                   dtype="float32")


def add(a, b):
    add_op = a + b
    return relay.Function([a, b],
                          add_op,
                          ret_type=type_annotation)


a, b = [relay.var(name, type_annotation) for name in "ab"]
add_func = add(a, b)
add_func

 

Out[ ]:

fn (%a: Tensor[(5, 5), float32], %b: Tensor[(5, 5), float32]) -> Tensor[(5, 5), float32] {
  add(%a, %b)
}

 

 

模块¶

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

 

In [ ]:

# 定义变量
names = "xy"
x, y = [relay.var(name) for name in names]
# 定义函数
add_op = x + y
add_func = relay.Function([x, y], add_op)

 

 

声明全局变量：

 

In [ ]:

add_gvar = relay.GlobalVar("AddFunc")
print(add_gvar)

 

 

@AddFunc

 

 

定义将 add_func 提升为全局变量：

 

In [ ]:

from tvm import IRModule

mod = IRModule({add_gvar: add_func})
print(mod)

 

 

def @AddFunc(%x, %y) {
  add(%x, %y)
}

 

 

获取模块的全局变量内容：

 

In [ ]:

mod[add_gvar]

 

Out[ ]:

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

 

 

也可以直接借助全局变量的名字获取其内容：

 

In [ ]:

mod["AddFunc"]

 

Out[ ]:

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

 

 

也可以分配新的全局变量给模块：

 

In [ ]:

names = "xy"
x, y = [relay.var(name) for name in names]
# 定义函数
mul_op = x * y
mul_func = relay.Function([x, y], mul_op)
mod["MulFunc"] = mul_func

print(mod)

 

 

def @AddFunc(%x, %y) {
  add(%x, %y)
}

def @MulFunc(%x1, %y1) {
  multiply(%x1, %y1)
}

 

 

也可以通过 Python 字典更新全局变量：

 

In [ ]:

names = "xyz"
x, y, z = [relay.var(name) for name in names]
# 定义函数
v1 = x * y
muladd_op = v1 + z
muladd_func = relay.Function([x, y, z], muladd_op)

mod.update({"MulAddFunc": muladd_func})
print(mod)

 

 

def @AddFunc(%x, %y) {
  add(%x, %y)
}

def @MulAddFunc(%x1, %y1, %z) {
  %0 = multiply(%x1, %y1);
  add(%0, %z)
}

def @MulFunc(%x2, %y2) {
  multiply(%x2, %y2)
}

 

 

查看全局变量：

 

In [ ]:

mod.get_global_vars()

 

Out[ ]:

[GlobalVar(AddFunc), GlobalVar(MulFunc), GlobalVar(MulAddFunc)]

 

 

回调全局变量¶

{class}tvm.relay.expr.Call 可以在模块中回调全局变量。

比如，定义 add 算子：

 

In [ ]:

data = relay.var("data")
bias = relay.var("bias")
add_op = data + bias

 

 

初始化 Relay 模块：

 

In [ ]:

mod = IRModule()

 

 

创建并绑定 add 全局函数到 mod：

 

In [ ]:

mod['AddFunc'] = relay.Function([data, bias], add_op)

 

 

下面定义三个变量用于定义“连加”运算：

 

In [ ]:

a, b, c = [relay.var(name) for name in "abc"]

 

 

获取全局变量 add：

 

In [ ]:

add_gvar = mod.get_global_var('AddFunc')

 

 

定义“连加”运算：

 

In [ ]:

add_01 = relay.Call(add_gvar, [a, b])
add_012 = relay.Call(add_gvar, [c, add_01])

 

 

绑定到 mod：

 

In [ ]:

mod['main'] = relay.Function([a, b, c], add_012)
print(mod)

 

 

def @AddFunc(%data, %bias) {
  add(%data, %bias)
}

def @main(%a, %b, %c) {
  %0 = @AddFunc(%a, %b);
  @AddFunc(%c, %0)
}

 

 

常量¶

{class}tvm.relay.Constant 节点表示常量张量值。常量被表示为 {class}~tvm.runtime.ndarray.NDArray，允许 Relay 使用 TVM 算子来计算常量。

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

 

In [ ]:

from tvm import nd

 

In [ ]:

const_a = nd.array(7)
const_op = relay.Constant(const_a)
str(const_op), repr(const_op)

 

Out[ ]:

('7i64', 'Constant(7)')

 

 

元组¶

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

 

In [ ]:

a = relay.var("a", shape=(10, 10))
b = relay.var("b", shape=(100, 20))
c = relay.var("c", shape=(100, 20))
f_tuple = relay.Tuple([a, b, c])

print(f_tuple)

 

 

free_var %a: Tensor[(10, 10), float32];
free_var %b: Tensor[(100, 20), float32];
free_var %c: Tensor[(100, 20), float32];
(%a, %b, %c)

 

 

支持索引：

 

In [ ]:

f_tuple[0]

 

Out[ ]:

Var(a, ty=TensorType([10, 10], float32))

 

 

If 条件表达¶

 

In [ ]:

a, b = [relay.var(name) for name in "ab"]
c = a + b
d = a * b
cond = relay.const(a == b)
f = relay.If(cond, c, d)
print(f)

 

 

if (False) {
  free_var %a;
  free_var %b;
  add(%a, %b)
} else {
  multiply(%a, %b)
}

 

 

let 绑定¶

其实上述介绍的模块绑定属于 Relay graph 绑定，对应于计算图。

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

• let 绑定包含局部变量、可选类型注解、值和可以引用绑定标识符的 body 表达式。如果省略了绑定变量上的类型注释，Relay 将尝试推断该变量允许的最通用类型。
• let 表达式中的绑定变量只作用在其 body 作用域内，除非该变量定义了函数表达式。当 let 表达式创建函数时，该变量的值也在范围内，以允许递归定义函数（请参阅前一小节）。
• let 绑定的值是计算它所依赖的绑定后的最后一个表达式的值。

 

In [ ]:

x = relay.var("x")
sb = relay.ScopeBuilder()
v1 = sb.let("v1", relay.log(x))
v2 = sb.let("v2", v1 + v1)
sb.ret(v2)
f = relay.Function([x], sb.get())

f

 

Out[ ]:

fn (%x) {
  let %v1 = log(%x);
  let %v2 = add(%v1, %v1);
  %v2
}

 

 

也可以定义 if-else 语句：

 

In [ ]:

sb = relay.ScopeBuilder()
cond = relay.var("cond", 'bool')
x = relay.var("x")
y = relay.var("y")

with sb.if_scope(cond):
    one = relay.const(1, "float32")
    t1 = sb.let("t1", relay.add(x, one))
    sb.ret(t1)
with sb.else_scope():
    sb.ret(y)
f = relay.Function([x, y, cond], sb.get())
f

 

Out[ ]:

fn (%x, %y, %cond: bool) {
  if (%cond) {
    let %t1 = add(%x, 1f);
    %t1
  } else {
    %y
  }
}

 

更多内容见我的 Github：TVM 手册 — tvm-book 0.0.2 文档 (xinetzone.github.io)