torch_xla 中重新编译的来源¶

我们先从一些事实/约束开始：¶

XLA 中的图编译非常昂贵。
XLA 只支持静态形状。换句话说，即使是相同的 IR 图，当输入形状改变时，XLA 也会重新编译。
重新编译在发生时会严重影响 torch_xla 的性能，并且从普通 Python 用户的角度来看，很难理解和调试。

通常，当发生重新编译时，我们会说我们只需要动态形状支持，然后就可以放心，当动态形状在未来得到支持时，所有重新编译都会神奇地消失。但事实并非如此，XLA 现在已经具有相当不错的有界动态形状覆盖，但我们仍然会看到重新编译，并且这是可以预期的。

**本文档旨在详细解释几种常见的重新编译来源，以及我们需要做些什么才能消除它们。它将主要面向没有任何上下文的初学者解释问题。为了便于理解，“解决方案”可能依赖于不切实际的假设。 **

#1. 来自输入数据集。¶

是的，输入数据集包含不同形状的示例非常普遍，例如，句子长度不同或图像大小不同。如果不进行归一化，每次输入形状新颖时都会导致重新编译。

Tensorflow 图模式用户更习惯于进行填充/分桶（tf.pad）来将输入形状归一化为一个或几个桶。但这对 PyTorch 动态图前端用户来说是一种反模式（这与惰性张量前端试图针对的同一用户相同），因为对于 CPU/CUDA 后端，不同的输入形状根本无关紧要。

建议的解决方法： 好的，现在我们假设可以通过教用户进行填充/分桶来解决这个问题（这在实践中很难 :P）。接下来呢？

#2. 来自算子输出¶

有些算子在语义上是数据依赖的，并产生动态形状的输出：例如，torch.nonzero 返回其输入张量中非零元素的索引。因此，即使此算子的输入张量形状始终相同，其输出形状也可能不同并导致重新编译。

2.1 有界动态形状可以解决您将具有动态形状的张量用作张量，而无需查询其真实维度的情况。¶

建议的解决方法： 假设现在 XLA 为所有算子支持有界动态形状，这样足够了吗？

有界动态形状意味着我们可以将张量填充到理论上的最大值，用更多的内存使用量来换取更少的重新编译/更快的速度。

嗯，算是一种。让我们看下面的例子

a = torch.tensor([1, 2, 0, 1, 3], device='xla')
b = torch.nonzero(a)
c = b * 2
d = c + 1
print(torch_xla._XLAC._get_xla_tensors_text([d]))

在上面的例子中，图中 b 下方的每个节点（即 c, d 以及依赖于它们的所有内容）都将具有动态形状，很明显 b 在第 0 维上具有动态形状，如下所示

  %9 = (s64[<=5,1]{1,0}, s64[]) aten::nonzero(%8), num_outputs=2 # b
  %10 = s64[5,1]{1,0} aten::mul(%9.0, %3) # c
  %11 = s64[5,1]{1,0} aten::add(%10, %2), ROOT=0 # d

尽管图中未直接显示，但 c & d 实际上也具有动态形状（换句话说，[5, 1] 只是填充后的形状并且被屏蔽了）。

print(torch_xla._XLAC._get_xla_tensor_dimension_size(d, 0)) # prints 4 instead of 5

您可以看到，在这种情况下，只要输入张量 a 的形状是 [5]，我们就只编译图一次。有界动态形状支持起到了作用！

2.2 如果在具有动态形状的张量上查询真实维度会怎样？¶

这实际上非常常用，因为并非所有 PyTorch 计算都以张量形式完成。

例如，PyTorch 中的 tensor.size() 返回一个整数元组而不是一个 dtype=int 的张量。当 tensor 是一个动态形状的张量时，这个操作基本上会迫使 XLA 切割图并进行评估，以便我们可以返回正确的标量（否则它只会返回填充后的形状，这是错误的）。

更糟糕的是，许多 PyTorch 也接受标量输入。在执行 s = tensor.size(0) 并将 s 用于其他算子后，它也成为一个动态源。在这种情况下，我们可能知道如何填充它以及它的上限，但我们无法这样做，因为它甚至不是一个张量！

 a = torch.tensor([1, 2, 0, 1, 3], device='xla')
 b = torch.nonzero(a)
 s = a.size(0) # evaluation happens! nit: we use size() for simplicity, the actual API is _get_xla_tensor_dimension_size.
 c = torch.rand(s, device='xla') # c can be of any shape between [0, 5] which causes more recompilations!
 d = c + 1

因此，如果没有 PyTorch 前端的帮助，这实际上很难解决。我们需要什么？

简而言之，我们需要一个张量世界！

例如，

tensor.size() 应该返回一个张量，这样它就可以成为一个具有动态形状的张量并在图中保持，而无需提前评估。
张量访问器，例如，对于 2D 张量，tensor[0][0] 现在返回一个值，但这需要也返回一个张量。
这隐含着所有目前接受 int/float/double 作为输入的算子也需要一个张量重载。这是一个巨大的请求，因为它可以轻易地使我们的算子集爆炸。
- 如果我们能让标量到张量的转换变得非常便宜，那么我们只需要关心张量重载，这样会更容易。
- 实际上，并非所有算子都接受来自先前计算的标量，因此我们一直在根据临时请求添加张量变体。
- 我认为这也是基于跟踪的方法中一个常见的请求。

好的，现在我们假设 PyTorch 中的每个算子都有我们需要的张量版本，我们完成了吗？

#3. 来自控制流¶

不！我们实际上只解决了没有数据依赖控制流的问题……

请看下面的例子

if x[0][0] == 3:
  bla
else:
  blabla

即使 x[0][0] 是一个张量，我们也需要执行/实例化它的值才能让 Python 解释器继续执行。多个控制流中的不同分支选择组合意味着我们需要编译大量图！

目前我们没有办法解决这个问题。要解决它，我们需要将控制流从 Python 降低到图！无需过多考虑实现，我们可以通过两种方式做到这一点：

要求用户显式使用控制流算子而不是 Python 的 if/else/while/for。这目前作为 torch_xla 中的自定义 API 受支持，但并未在用户代码中广泛采用。（Python 用户习惯于 if/else/for，除非有巨大的性能提升，否则很难让他们切换到更难用的 API）。
解析 Python 源代码。代码以自动获取控制流语句。这类似于 Torchscript，并将 torchscripted 图与惰性跟踪图正确合并（包括形状信息等）。我还没有仔细考虑过如何实现这一点 :P

但以上任何一种解决方案都需要大量的工作，无论是在用户端还是在框架端。这就是为什么考虑到我们目前的带宽，我们目前只是接受早期评估和多次编译的损失作为短期解决方案。

好的，现在我们假设控制流也已自动降低到图中，我们是否就成功了？

是的！现在您所有的计算都表示为张量运算的图，包括控制流，以便编译器可以消耗它们并执行其智能技巧！但老实说，此时您的程序不再非常“PyTorch 式”了。

结论：¶

实际上存在多种重新编译来源，有界动态形状支持无法解决所有问题。本文档中提出的解决方法有时肯定是不切实际的，并且可能有更好的方法来正确解决每个来源，而我完全不知道。但我希望随着我们在这篇文档中不断努力实现理想的惰性张量栈，您现在更容易理解我们面前仍然存在的障碍。

附录：¶

NNC 使用符号形状，这有帮助吗？

是的，但部分有效。通过使用符号形状，您的编译优化不再需要具体的形状值。换句话说，您生成的内核比 XLA 的静态形状内核更通用。

它具体帮助了哪些问题？

它有助于解决 #1 和 #2.1 的情况。

shape [3, 5] -> add -> transpose -> ... -> mul
shape [6, 2] -> add -> transpose -> ... -> mul

# with symbolic shape
shape [x, y] -> add -> transpose -> ... -> mul

使用符号形状，您的生成的内核不会像 XLA 使用静态形状那样重新编译。

XLA 通过填充/分桶（针对 #1）和有界动态形状（针对 #2.1）来解决这个问题。

Brian Hirsh (@bdhirsh) 在评论中提出了一些非常好的问题，移至此处以使其更加突出

在产生数据依赖输出形状的 XLA 内核中加入 TORCH_WARN 是否值得？

是的，torch_warn 对于告知用户“嘿，您的程序不会飞速运行”很有用。但对于这些数据依赖的算子，除非用户更改其模型中的逻辑，否则没有简单的重写方法。（另一个例子是 torch.unique()）

像 nonzero 这样的算子如何影响我们对 sizes() 进行虚函数消除的能力？如果我们想对 sizes() 进行虚函数消除，我们将需要能够为每个算子进行及时计算 sizes - 这是否意味着我们每次遇到像 nonzero 这样的算子时都会被迫评估图？Vs. 目前，听起来当用户调用 nonzero() 时，我们实际上并没有强制评估？

是的，这是一个很好的问题！所以以目前的形式，它不是一个硬性障碍，因为 XLA 张量上的 size() 不携带真实来源的大小信息。正如示例所示，真实来源存在于 IRValue 中，只能通过 _get_xla_tensor_dimension_size 检索。所以如果我们决定对 size() 进行虚函数消除，它只会强化这种差异。

作为后续，如果我们让 size() 返回张量而不是值，如上面建议的解决方法中所述。在这种情况下，size() 将无法进行虚函数消除，因为它变成了一个算子（接收张量并产生张量，对不同的后端有不同的实现）。

如果我在循环中调用 torch.add(input, 1)，其中 input 的大小从 1-1000 不等，通常我们需要编译 1000 个不同的图——但有了动态形状，XLA 似乎会在内部生成一个图，其中说“如果输入大小小于等于 1000，则使用此图”。我的问题是：“动态形状”只是图的属性吗？还是图和输入的属性？也就是说，如果我的代码在循环中调用 x = torch.add(input, 1); x.sizes()，此时 x 是否具有动态形状，这意味着我们需要运行图来获取大小？或者即使存在具有动态形状的图，我们也能将其作为及时计算的属性。

是的，在这种情况下，您将编译 1000 个不同的图。动态形状意味着其输入在其内部具有动态维度。因此，当您查询 x.sizes()（目前需要使用 get_dimention_size 来获取正确的大小）时，它将触发执行（由于大小没有改变，因此不会触发重新编译）。如果没有访问大小的行，则当输入具有动态维度时，它不会触发任何重新编译/执行。

将控制流提供到图中的一种替代方法是否是想办法确保 XLA 图不包含控制流？也就是说，如果我们有一个模型，中间有一个条件判断，然后让 XLA 生成 3 个图：1 个用于条件判断之前的所有内容，1 个用于 if 分支，1 个用于 else 分支。这意味着您不会为每个路径组合产生新的图的指数级增长，但是（a）图会更小，提供的优化机会更少，并且（b）让 XLA 识别条件路径的发生位置可能相当困难。

说得好！所以如果我们能将其分解成更小的图，这是可行的。但实际上这种模式很烦人

y = <some computation>
x = y + 2
if x[0] == 2 :
  z = y +1
else:
  z = y - 1

请注意，当您遇到控制流时，您将使用子图来评估 x，但分支计算中可能还包含之前的变量（例如 y 比 x 小一个节点，但在评估 x 时并未实例化）。所以您实际上正在为这个示例评估 1 个小图和 2 个大图。并且随着更多控制流的参与，y 可能会在多个分支中更新，这仍然会产生不同的大图组合。