Torch 导出中间表示规范

调用函数

call_function 节点表示对操作符的调用。

定义

• 功能性：如果一个可调用对象满足以下所有要求，则称其为“功能性”的：

  • 非修改操作：该操作不更改其输入的值（对于张量，这包括元数据和实际数据）。

  • 无副作用：操作符不会改变外部可见的状态，例如不修改模块参数的值。

• 运算符：是一个具有预定义模式的函数式调用。这类运算符的例子包括功能性的ATen运算符。

外汇表示

%name = call_function[target = operator](args = (%x, %y, …), kwargs = {})

与标准 FX call_function 的区别

1. 在FX图中，call_function可以引用任何可调用对象。但在导出IR时，我们只允许它使用ATen操作符、自定义操作符和控制流操作符的特定子集。

2. 在导出IR时，常量参数将会被嵌入到图形中。

3. 在FX图中，get_attr节点可以表示读取图模块中存储的任意属性。但在导出IR时，只能读取子模块中的属性，因为所有的参数和缓冲区都会被作为输入传递到图模块中。

占位符

占位符代表图的输入，其语义与FX中的完全一致。占位符节点必须是图的节点列表中的前N个节点，其中N可以为零。

外汇表示

%name = placeholder[target = name](args = ())

目标字段是一个字符串，表示输入的名称。

如果 args 不为空，则其大小应为 1，并表示该输入的默认值。

元数据

占位符节点也有像 call_function 节点一样的 meta[‘val’]。在这种情况下，val 字段表示图期望接收的此输入参数的形状和数据类型。

输出

输出调用代表函数中的返回语句，因此会终止当前的图。图中只有一个输出节点，并且它总是位于图的最后。

外汇表示

output[](args = (%something, …))

这与torch.fx 中的语义完全一致。其中，args 表示需要返回的节点。

元数据

输出节点和call_function节点具有相同的元数据。

获取属性

get_attr 节点表示从封装的 torch.fx.GraphModule 中读取子模块。与来自 torch.fx.symbolic_trace() 的普通 FX 图不同，后者使用 get_attr 节点从顶级 torch.fx.GraphModule 读取属性（如参数和缓冲区），在导出的程序中，这些参数和缓冲区作为图模块的输入传递，并存储在顶级 torch.export.ExportedProgram 中。

外汇表示

%name = get_attr[target = name](args = ())

示例

考虑以下模型：

from functorch.experimental.control_flow import cond

def true_fn(x):
    return x.sin()

def false_fn(x):
    return x.cos()

def f(x, y):
    return cond(y, true_fn, false_fn, [x])

图表:

graph():
    %x_1 : [num_users=1] = placeholder[target=x_1]
    %y_1 : [num_users=1] = placeholder[target=y_1]
    %true_graph_0 : [num_users=1] = get_attr[target=true_graph_0]
    %false_graph_0 : [num_users=1] = get_attr[target=false_graph_0]
    %conditional : [num_users=1] = call_function[target=torch.ops.higher_order.cond](args = (%y_1, %true_graph_0, %false_graph_0, [%x_1]), kwargs = {})
    return conditional

该行代码 %true_graph_0 : [num_users=1] = get_attr[target=true_graph_0] 读取子模块 true_graph_0，其中包含 sin 操作符。

SymInt

A SymInt 是一个对象，它可以是一个具体的整数值或代表整数的符号（在 Python 中由 sympy.Symbol 类表示）。当 SymInt 作为符号时，它描述了一个类型为整数但在编译阶段图中未知的变量，也就是说，它的具体值仅在运行时才能确定。

FakeTensor

FakeTensor 是一个包含张量元数据的对象，可以视为具有以下属性：

class FakeTensor:
  size: List[SymInt]
  dtype: torch.dtype
  device: torch.device
  dim_order: List[int]  # This doesn't exist yet

FakeTensor 的 size 字段是一个由整数或 SymInt 组成的列表。如果存在 SymInt，说明张量具有动态形状；若只有整数，则表示张量具有固定的静态形状。TensorMeta 的秩始终是固定的。dtype 字段表示该节点输出的数据类型。Edge IR 中没有隐式的类型提升机制。FakeTensor 不包含 strides。

换句话说：

• 如果节点的目标操作返回一个张量，那么 node.meta['val'] 就是描述这个张量的一个 FakeTensor。

• 如果节点中的操作符返回一个包含 n 个张量的元组，那么 node.meta['val'] 就会是一个包含每个张量的 FakeTensor 元组。

• 如果节点中的操作符返回一个编译时已知的整数、浮点数或标量值，则 node.meta['val'] 为 None。

• 如果节点中的操作符返回一个在编译时未知的整数、浮点数或标量值，则 node.meta['val'] 的类型为 SymInt。

例如：

• aten::add 返回一个张量，因此它的规格将是一个假张量，具有此操作返回的张量的数据类型和大小。

• aten::sym_size 返回一个整数，因此它的值会是一个 SymInt，因为在运行时才能确定其具体值。

• max_pool2d_with_indexes 返回一个包含两个张量的元组 (Tensor, Tensor)；因此，规格将是一个包含两个 FakeTensor 对象的 2 元组，其中第一个 TensorMeta 描述返回值的第一个元素等。

Python代码：

def add_one(x):
  return torch.ops.aten(x, 1)

图表:

graph():
  %ph_0 : [#users=1] = placeholder[target=ph_0]
  %add_tensor : [#users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%ph_0, 1), kwargs = {})
  return [add_tensor]

FakeTensor:

FakeTensor(dtype=torch.int, size=[2,], device=CPU)

可以转换为Pytree的类型

我们将一种类型定义为“可构成Pytree的类型”，如果它要么是叶子类型，要么是包含其他可构成Pytree类型的容器类型。

注意：

pytree 的概念与 JAX 文档中的定义相同：

以下类型被定义为叶类型：

以下类型被定义为容器类型：