重编译故障排除¶

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

XLA 中的图编译成本相当高。
XLA 仅支持静态形状。换句话说，即使对于相同的 IR 图，当输入形状发生变化时，XLA 也会重新编译。
重编译会极大地损害 torch_xla 的性能，而且对于普通的 Python 用户来说，很难理解和调试。

通常，当发生重编译时，我们会说我们只需要动态形状支持，然后认为以后有了动态形状支持，所有重编译都会神奇地消失。但这并非如此，XLA 目前已经有了相当好的有界动态形状覆盖，但我们仍然会看到重编译，而且这些是符合预期的。

本文档旨在详细解释一些常见的重编译源，以及我们需要做什么来消除它们。它将主要面向没有背景知识的初学者进行解释。为了易于理解，“解决方案”可能依赖于不切实际的假设。

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

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

Tensorflow 图模式用户更习惯于进行填充/分桶（tf.pad）来将输入形状归一化到一个或几个桶中。但这对于 PyTorch eager 前端用户（也是 lazy tensor 前端尝试针对的用户）来说是一种反模式，因为不同的输入形状对于 eager 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() 进行符号展开，我们就需要能够为每个算子预先计算大小 - 这是否意味着我们每次遇到像 nonzero 这样的算子时都会被迫评估图？还是说，正如您所说，我们目前并没有强制评估当用户调用 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 可能会在多个分支中更新，从而仍然产生不同组合的大图。