Pytorch/XLA 中重新编译的来源¶

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

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，并将 torchscript 图以某种方式正确合并到延迟跟踪图中（包括形状信息等）。我还没有想清楚如何实现这一点的步骤 :P

但是，上述任何一种解决方案都需要付出相当大的努力，无论是在用户方面还是在框架方面。这就是为什么考虑到我们拥有的带宽，我们目前只是接受早期评估和多次编译的打击作为短期解决方案。

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

是的！现在您的整个计算都以张量操作图表示，包括控制流，以便编译器现在可以使用并执行其智能技巧！但说实话，在这一点上，您的程序不再是非常 PyTorch-y 的了。

结论：¶

实际上存在多个重新编译来源，并且有界动态形状支持无法解决所有这些问题。本文档中提出的建议的解决方法有时肯定是不切实际的，并且可能存在更好的方法来正确解决每个来源，而我完全没有意识到。但我希望，当我们不断突破障碍，朝着本文档中理想的延迟张量堆栈前进时，您现在可以更容易地理解我们面前剩下的障碍是什么。

附录：¶

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() 将无法取消虚拟化，因为它变成了一个运算符（输入张量并生成张量，对于不同的后端具有不同的实现。）

如果我在循环中调用 torch.add(input, 1)，其中输入大小从 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 个小图和两个大图。并且随着更多控制流的参与，y 可能会在多个分支中更新，这仍然会产生不同的大图组合。