Pytorch/XLA 中的重编译来源¶

让我们先来介绍一些事实/约束:¶

XLA 中的图编译代价很高。
XLA 仅处理静态形状。换句话说，即使对于相同的 IR 图，当输入形状改变时，XLA 也会重新编译。
重编译发生时会严重影响 torch_xla 的性能，并且从普通 Python 用户的角度来看，很难理解和调试。

通常，当发生重编译时，我们说我们只需要动态形状支持，然后就可以放心了，当未来支持动态形状时，所有重编译都会神奇地消失。但事实并非如此，XLA 现在已经有了很好的有界动态形状覆盖，但我们仍然看到重编译，而且这是意料之中的。

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

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

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

Tensorflow 图模式用户更习惯于进行填充/分桶（tf.pad）来将输入形状归一化为一个或少数几个桶。但这对于 PyTorch 动态图前端用户（也是 lazy tensor 前端试图瞄准的用户）来说是一种反模式，因为对于动态图 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 是一个动态形状张量时，此 op 基本上会强制 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 作为输入的运算符也需要一个张量重载。这是一个很大的要求，因为这很容易导致我们的运算符集爆炸。
- 如果我们能使标量到张量的转换非常便宜，那么我们只需要关注张量重载，这样会更容易。
- 实际上，并非所有 op 都接受来自先前计算的标量，因此我们一直在通过临时请求添加张量变体。
- 我认为这也是基于跟踪的方法经常提出的要求。

好的，现在我们假设 PyTorch 中的每个 op 都具有我们需要的张量版本，我们是否完成了？

#3. 来自控制流¶

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

看下面的例子

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

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

目前我们还没有办法解决这个问题。要解决它，我们需要将控制流从 Python 降低到图！在不深入思考实现的情况下，我们可以通过两种方式做到这一点：

要求用户显式使用控制流 op 而不是 Python 的 if/else/while/for。这目前在 torch_xla 中的自定义 API 中得到支持，但并未在用户的代码中广泛采用。（Python 用户习惯于 if/else/for，除非有巨大的性能提升，否则很难让他们切换到更丑陋的 API）。
解析 Python 源代码。代码以自动获取控制流语句。这就像 Torchscript 一样，并以某种方式将 torchscripted 图与懒惰跟踪图正确合并（包括形状信息等）。我还没有真正思考过实现这些步骤：P

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

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

是的！现在您的整个计算都表示在 Tensor 运算的图中，包括控制流，因此编译器现在可以消耗并执行其智能技巧！但说实话，此时您的程序不再那么 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 有助于告诉用户“嘿，您的程序将无法超快运行”。但对于这些数据依赖的 op，除非用户更改模型中的逻辑，否则没有简单的重写方法。（另一个例子是 torch.unique()）

像 nonzero 这样的 op 如何影响我们对 sizes() 进行动态绑定的能力？如果我们想对 sizes() 进行动态绑定，我们将需要能够为每个 op 优先计算大小——这是否意味着每次遇到像 nonzero 这样的 op 时，我们都必须评估图？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 可能会在多个分支中被更新，这仍然会产生不同组合的大图。