TorchDynamo 深入探讨¶
在阅读本节之前，请先阅读torch.compiler。
TorchDynamo 是一个 Python 级别的即时（JIT）编译器，旨在使未修改的 PyTorch 程序运行得更快。TorchDynamo 通过钩入 CPython 中的帧评估 API（PEP 523）来动态修改 Python 字节码，就在其执行之前。它重写 Python 字节码以提取 PyTorch 操作序列到一个 FX Graph 中，然后使用可定制的后端进行编译。它通过字节码分析创建这个 FX Graph，并设计为将 Python 执行与编译后端混合，以获得两者的最佳效果 — 可用性和性能。
TorchDynamo 使得使用不同的编译器后端来加速 PyTorch 代码变得容易，只需使用一行装饰器 torch._dynamo.optimize()，该装饰器由 torch.compile() 方便地包装。
以下图表展示了PyTorch在使用torch.compile和不使用它的情况下的工作方式：
TorchInductor 是 TorchDynamo Graph 支持的后端之一， 用于 Triton 的GPU或 C++/OpenMP 的CPU。我们有一个 训练性能仪表板， 提供了不同训练后端的性能比较。你可以在 PyTorch dev-discuss 上的 TorchInductor 帖子 中阅读更多内容。
如需深入了解，请阅读以下章节，观看深入探讨视频，并查看开发者讨论主题。
TorchDynamo 内部机制¶
作者: Jason Ansel 和 Kaichao You
本节将介绍一些TorchDynamo的内部工作原理，并演示TorchDynamo在幕后是如何工作的。
什么是守卫？¶
TorchDynamo 采用即时操作，并根据动态属性专门化图表。以下是一个使用 TorchDynamo 的基本示例。可以通过使用 torchdynamo.optimize 装饰一个函数或方法来启用 TorchDynamo 优化：
from typing import List
import torch
from torch import _dynamo as torchdynamo
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
print("my_compiler() 被调用，传入的 FX 图为:")
gm.graph.print_tabular()
return gm.forward # 返回一个 Python 可调用对象
@torchdynamo.optimize(my_compiler)
def toy_example(a, b):
x = a / (torch.abs(a) + 1)
if b.sum() < 0:
b = b * -1
return x * b
for _ in range(100):
toy_example(torch.randn(10), torch.randn(10))
例如，上面的第一个图表有以下守卫：
GUARDS:
hasattr(L['a'], '_dynamo_dynamic_indices') == False
hasattr(L['b'], '_dynamo_dynamic_indices') == False
utils_device.CURRENT_DEVICE == None
___skip_backend_check() or ___current_backend() == ___lookup_backend(140355900538256)
check_tensor(L['a'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[10], stride=[1])
check_tensor(L['b'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[10], stride=[1])
如果这些保护中的任何一个失败，图将被重新捕获并重新编译。那里有趣的保护是 check_tensor，它检查以下 torch.Tensor 属性：
张量的Python类（张量子类化等）
数据类型
设备
requires_grad
dispatch_key（应用了线程本地的包含/排除）
维度
尺寸*
步幅*
完全特化模式允许后端编译器假设一个完全静态的图。不幸的是，大多数后端都需要这个。当不在动态形状模式时，返回动态形状的操作符将触发图的中断。
Dynamo在做什么？¶
如果你想更好地理解TorchDynamo在做什么，你可以使用以下代码运行你的代码：
TORCH_LOGS="+dynamo,guards,bytecode"
如果你不熟悉Python字节码，你可以添加一个反编译钩子，将字节码反编译成人类可读的源代码。一个可用的工具是depyf。如果你还没有安装depyf，运行pip install depyf。然后，在你运行任何代码之前，添加以下代码来安装反编译钩子。
import depyf
depyf.install()
此代码会触发有用的（但可能被视为垃圾信息的）打印输出。
例如，toy_example中第一个图的打印输出如下：
__compiled_fn_0 <eval_with_key>.1
opcode name target args kwargs
------------- ------- ------------------------------------------------------ ---------------- --------
placeholder a a () {}
placeholder b b () {}
call_function abs_1 <built-in method abs of type object at 0x7f9ca082f8a0> (a,) {}
call_function add <built-in function add> (abs_1, 1) {}
call_function truediv <built-in function truediv> (a, add) {}
call_method sum_1 sum (b,) {}
call_function lt <built-in function lt> (sum_1, 0) {}
output output output ((truediv, lt),) {}
原始字节码 toy_example example.py 行 12
14 0 LOAD_FAST 0 (a)
2 LOAD_GLOBAL 0 (torch)
4 LOAD_METHOD 1 (abs)
6 LOAD_FAST 0 (a)
8 CALL_METHOD 1
10 LOAD_CONST 1 (1)
12 BINARY_ADD
14 BINARY_TRUE_DIVIDE
16 STORE_FAST 2 (x)
15 18 LOAD_FAST 1 (b)
20 LOAD_METHOD 2 (sum)
22 CALL_METHOD 0
24 LOAD_CONST 2 (0)
26 COMPARE_OP 0 (<)
28 POP_JUMP_IF_FALSE 19 (to 38)
16 30 LOAD_FAST 1 (b)
32 LOAD_CONST 3 (-1)
34 BINARY_MULTIPLY
36 STORE_FAST 1 (b)
17 >> 38 LOAD_FAST 2 (x)
40 LOAD_FAST 1 (b)
42 BINARY_MULTIPLY
44 RETURN_VALUE
修改后的字节码 toy_example example.py 行 12
12 0 LOAD_GLOBAL 3 (__compiled_fn_0)
2 LOAD_FAST 0 (a)
4 LOAD_FAST 1 (b)
6 CALL_FUNCTION 2
8 UNPACK_SEQUENCE 2
10 STORE_FAST 2 (x)
12 POP_JUMP_IF_FALSE 12 (to 24)
14 LOAD_GLOBAL 4 (__resume_at_30_1)
16 LOAD_FAST 1 (b)
18 LOAD_FAST 2 (x)
20 CALL_FUNCTION 2
22 RETURN_VALUE
>> 24 LOAD_GLOBAL 5 (__resume_at_38_2)
26 LOAD_FAST 1 (b)
28 LOAD_FAST 2 (x)
30 CALL_FUNCTION 2
32 RETURN_VALUE
可能的源代码:
def toy_example(a, b):
__temp_1 = __compiled_fn_0(a, b)
x = __temp_1[0]
if __temp_1[1]:
return __resume_at_30_1(b, x)
return __resume_at_38_2(b, x)
如果你发现反编译的代码是错误的，请在 https://github.com/youkaichao/depyf/issues 提交一个issue。
在顶部你可以看到FX图表。 接下来，你会看到函数的原始字节码，然后是由TorchDynamo生成的修改后的字节码，以及用于参考的反编译源代码。最后，你会看到我们在上面讨论过的保护措施。
在修改后的字节码中，__compiled_fn_0 是
my_compiler() 的返回值（编译后的图）。__resume_at_30_1 和
__resume_at_38_2 都是生成的延续函数，它们在图中断后（在字节码偏移量30和38处）恢复执行。这些函数的形式如下：
__resume_at_<offset>:
... 如果需要，恢复堆栈状态 ...
JUMP_ABSOLUTE <offset> 进入 toy_example
... toy_example 的原始字节码 ...
通过生成这个 resume_at 函数，我们强制函数的剩余部分在一个新的 Python 帧中执行，这会递归地触发 TorchDynamo 在第一次执行到达该点时重新启动其捕获。
如何检查由TorchDynamo生成的工件？¶
要检查由 TorchDynamo 生成的工件，有一个 API torch._dynamo.eval_frame._debug_get_cache_entry_list，它可以从函数的 __code__ 对象中检索编译后的代码和保护。一个编译后的函数可以有多个缓存条目，每个缓存条目由一个生成的函数来检查保护条件，以及一个 types.CodeType 对象来保存如果保护条件满足时将要执行的代码。
from torch._dynamo.eval_frame import _debug_get_cache_entry_list, innermost_fn
cache_entries = _debug_get_cache_entry_list(innermost_fn(toy_example))
cache_entry = cache_entries[0]
guard, code = cache_entry.check_fn, cache_entry.code
# 该守卫函数接收输入帧的局部变量，并判断是否需要触发重新编译。
import dis
dis.dis(guard)
dis.dis(code)
如果你了解Python字节码，你可以理解上述输出。
对于守卫函数，无需检查字节码。我们可以直接访问其守卫条件：
for code_part in guard.code_parts:
print(code_part)
输出为：
___guarded_code.valid
___check_global_state()
hasattr(L['a'], '_dynamo_dynamic_indices') == False
hasattr(L['b'], '_dynamo_dynamic_indices') == False
utils_device.CURRENT_DEVICE == None
___skip_backend_check() or ___current_backend() == ___lookup_backend(140215810860528)
___check_tensors(L['a'], L['b'], tensor_check_names=tensor_check_names)
只有当所有条件都满足时，守卫函数返回true，编译后的代码才会执行。
对于编译后的代码，我们无法直接访问其源代码，但必须对其进行反编译。
from depyf import decompile
print(decompile(code))
输出为：
def toy_example(a, b):
__temp_1 = __compiled_fn_0(a, b)
x = __temp_1[0]
if __temp_1[1]:
return __resume_at_30_1(b, x)
return __resume_at_38_2(b, x)
代码中引用的一些名称包括：
编译后的函数，存储在包含原始函数的模块的全局命名空间中。这些包括名称如
__compiled_fn_0/__resume_at_30_1/__resume_at_38_2。用于检查守卫的闭包变量。名称可以从
guard.__code__.co_freevars访问，值存储在guard.__closure__中。这些包括诸如___guarded_code/___is_grad_enabled/___are_deterministic_algorithms_enabled/___is_torch_function_enabled/utils_device/___check_tensors/tensor_check_names之类的名称。函数
guard的参数L。这是一个字典，将toy_example的参数名称映射到其值。这仅在函数被调用时可用，此时帧评估 API 开始发挥作用。简而言之，L是一个具有{'a': value_a, 'b': value_b}结构的dict。因此，您可以看到代码使用L['a']来引用输入变量a。
图中断在编译的toy_example代码中显示，我们需要使用Python解释器来选择以下图进行执行。
请注意，我们将一个简单的 my_compiler 函数作为后端编译器传递，因此子图代码 __resume_at_38_2、__resume_at_30_1 和 __compiled_fn_0 仍然是 Python 代码。这也可以被检查（请忽略函数名称，仅使用函数签名和函数体代码）：
print("__compiled_fn_0 的源代码:")
print(innermost_fn(__compiled_fn_0).__self__.code)
print("=" * 60)
print("__resume_at_30_1 的源代码:")
print(decompile(__resume_at_30_1))
print("=" * 60)
print("__resume_at_38_2 的源代码:")
print(decompile(__resume_at_38_2))
source code of __compiled_fn_0:
def forward(self, L_a_ : torch.Tensor, L_b_ : torch.Tensor):
l_a_ = L_a_
l_b_ = L_b_
abs_1 = torch.abs(l_a_)
add = abs_1 + 1; abs_1 = None
truediv = l_a_ / add; l_a_ = add = None
sum_1 = l_b_.sum(); l_b_ = None
lt = sum_1 < 0; sum_1 = None
return (truediv, lt)
# 要查看更多调试信息，请使用 ``graph_module.print_readable()``
============================================================
source code of __resume_at_30_1:
def <resume in toy_example>(b, x):
b = b * -1
return x * b
============================================================
source code of __resume_at_38_2:
def <resume in toy_example>(b, x):
return x * b
然而，如果我们使用其他后端，如内置的 inductor，子图代码将被编译为用于 GPU 的 CUDA 内核或用于 CPU 的 C++ 代码。
总结来说，编译后的代码在概念上等同于以下代码：
def compiled_example(a, b):
L = {'a': a, 'b': b}
for guard, code in get_cache_entries():
if guard(L):
return code(a, b)
recompile_and_add_another_cache_entry()
以下图表展示了torch.compile如何转换和优化用户编写的代码：它首先从用户编写的函数中提取计算图，并将这些图编译成优化的函数，然后将它们组装成一个新的函数，该函数在功能上等同于用户编写的代码，但经过优化以获得良好的计算速度。
