https://pytorch.org/blog/understanding-lazytensor-system-performance-with-pytorch-xla-on-cloud-tpu/

概述

PyTorch/XLA是PyTorch的一个扩展，与Google的XLA（Accelerated Linear Algebra）编译器相结合，旨在解决PyTorch在训练过程中的性能瓶颈。它利用XLA的图执行（graph execution）功能对计算图进行性能优化，提高模型在硬件加速器（如GPU/TPU）上的计算效率。

PyTorch默认采用动态图计算的即时执行（eager execution）方式，即每个操作都会立即执行并返回结果。这种实时交互对于实验、调试和开发非常有用，也是PyTorch能够迅速赢得用户心智并取得成功的原因之一。Eager模式确实是个很好的特性，这一点从TensorFlow 2也跟随默认采用eager execution可以看出。

然而，即时执行也存在明显的缺点：没有图就没有优化。为了弥补这一点，TensorFlow 2提供了tf.function API，使用户能够方便地将计算封装成函数，并通过这些API生成、优化和编译计算图，实现eager和graph的混合执行。

虽然相较于TensorFlow 1，TensorFlow 2的图执行门槛已经降低很多，但使用体验仍然不够灵活和pythonic。不仅需要把代码封装进函数，还需要确保其中没有不能编译的算子，比如自定义算子。

PyTorch采用了LazyTensor 的设计使得在即时执行中混合图执行变得更加灵活、自然，无论是在性能优化还是编程体验方面都带来了显著的改进。

LazyTensor and XLA

LazyTensor是一种延迟计算的策略，它将即时执行和图执行的优势结合起来，实现了动态图和静态图的合一。

在运行模型或函数时，我们通常只关心最终的输出，而不关心中间结果。LazyTensor利用这一点，通过跟踪计算过程中的操作，生成XLA计算图。当需要获取计算结果时，可以通过即时编译（JIT）的方式，对计算图进行优化和编译，用于后续反复执行，从而获得比动态图执行更快的性能。

通过LazyTensor，用户可以享受到动态图的易用性和灵活性，同时又能获得静态图执行的高效性和优化能力，为深度学习的研究和应用带来了更大的便利和效益。

Lazy IR

PyTorch/XLA对源码的侵入性改动很小，只需要用户将张量放在XLA设备上（.to(dev)），LazyTensor就能发挥作用：在即时执行的过程中构建计算图，并在适当的时机执行图。

通过torch_xla._XLAC._get_xla_tensors_text([y1])，可以获取以y1为根节点的计算子图IR。这些IR是一段op序列，看起来类似于TorchScript，但也带有HLO（High Level Optimizer）的特点。

Device Data

x1和x2是通过torch.randn创建的，但在IR中没有看到aten::randn，相反却有两个xla::device_data操作，这些操作的目的是什么？如果你对异构计算有所了解，就会知道每个输入/输出张量都需要在主机端（host）和设备端（device）创建相应的张量，将输入推送到设备张量并在计算后将结果复制回主机张量。

xla::device_data用于创建设备张量并将张量数据推送到设备或从设备复制回主机。PyTorch内部维护一个设备张量缓存，可以通过主机张量可以找到它相应的设备张量。IR中的xla::device_data操作就是为主机张量在设备端创建相应的设备张量，并将数据从主机端推送给它。

Dynamic Graph

LazyTensor的另一个令人振奋的特性是其出色的动态图支持（例如动态形状和控制流）。正如之前所介绍的，PyTorch的静态图在支持控制流（如if、while）方面相当费力，更不用说支持动态输入形状了（ref: 从PyTorch到TorchScript: 打通深度学习模型的生产和应用）。

LazyTensor通过追踪（tracing）的方式构建计算图。如果输入形状是动态的或者存在if分支判断，它会将捕获到的新计算图通过JIT编译并保存到缓存中，可以动态地选择执行不同的计算图。（ref: JIT: Tensorflow如何实现即时编译）。

上述示例代码展示了不同分支产生的不同计算图。

IR Builder

Lazy IR的构建过程与torch.jit.trace()构建TorchScript的过程相似。当输入张量（x1、x2）被赋予XLA设备属性后，PyTorch的调度器（dispatcher）会将操作首先分发给与XLA调度键（dispatch key）。而torch_xla为XLA调度键注册了一个回退函数（fallback function），所有的XLA操作都会首先分派到该函数来构建图节点（参考：理解PyTorch分发机制的内部工作原理）。

Graph Execution

图执行的触发时机发生在设备端的张量需要得到计算结果的时候。例如，print(y1) 时，需要将设备端上的y1传输到主机端：dev_y1 -> y1，那这前提自然就是编译和执行计算图以获得dev_y1。 这个过程对于用户来说是隐性、透明的。

除了 print tensor 之外，当出现 IR 生成障碍的时候也会触发 graph execution。比如说有一个 custom op，XLA 并不支持让它运行在加速器上，那它就会 fallback 到 CPU上运行。同样的，需要将之前的计算结果传给 CPU 上的 fallback op。

除了打印张量（print tensor）之外，当生成 IR 遇到障碍时，也会触发图执行。举例来说，如果存在一个自定义操作（custom op），XLA 不支持在加速器上运行该操作，那么它会回退（fallback）到在 CPU 上运行。这个过程也会触发图执行，将设备上的计算结果作为输入传递给在 CPU 上回退运行的自定义操作。

除了被动触发图执行，还可以通过主动设置障碍 xm.mark_step() 的方式来触发图执行。当调用 mark_step() 时，会将当前的计算步骤标记为一个障碍点，表示在此处触发图执行：找出活跃 tensor -- loss，用它的IR 构建 HLO（XLA IR）的序列化表示 -- XLA Computation，然后交给 XLA 编译器编译执行，得到这些tensor的结果。

Conclusion

PyTorch/XLA是PyTorch的扩展，与Google的XLA编译器结合，旨在通过静态图执行优化提升模型在加速器上的计算效率。LazyTensor采用延迟计算策略，兼具即时执行和图执行的优势，实现了动静合一。它通过追踪操作生成XLA计算图，并即时编译优化以提高重复计算性能。LazyTensor不仅支持动态图和控制流，还能通过tracing捕获新的计算图并进行编译保存。图执行可以通过获取主机张量值或设置障碍来触发。PyTorch/XLA的设计实现了高效的加速器执行，并实现了动态图与图执行之间的平滑切换。

END