Dynamo概览
请在阅读torch.compiler之后再阅读本节内容。
TorchDynamo（或简称 Dynamo）是一个在 Python 级别的即时（JIT）编译器，旨在使未修改的 PyTorch 程序运行得更快。它通过 CPython 中的帧评估 API（PEP 523）在执行前动态地修改 Python 字节码。Dynamo 重写字节码，提取 PyTorch 操作序列，并将其转换为一个 FX 图形，然后使用可自定义的后端进行编译。它是通过字节码分析创建这个 FX 图形，并设计用于将 Python 执行与编译后的后端混合起来，以获得最佳效果——易用性和性能。
Dynamo 使用单行装饰器 torch._dynamo.optimize() 让试验不同的编译后端以加速 PyTorch 代码变得非常容易，该装饰器由 torch.compile() 方便地封装。
下图展示了 PyTorch 在使用 torch.compile 和不使用它的情况下分别如何工作：
TorchInductor 是 Dynamo 图形 支持的后端之一，用于将代码转换为针对 GPU 的 Triton 或针对 CPU 的 C++/OpenMP。我们有一个训练性能仪表板，提供了不同训练后端的性能比较。您可以在 PyTorch dev-discuss 上的 TorchInductor 帖子 中了解更多详情。
请详细阅读下方各部分内容，观看深入解析视频，并查阅dev-discuss相关讨论。
TorchDynamo dev-discuss主题
Dynamo 内部
作者: Jason Ansel 和 Kaichao You
本节将介绍Dynamo的一些内部机制，并展示其背后的工作原理。
什么是守卫?
Dynamo 实时运行，并根据动态属性对图进行专门化处理。下面是一个使用 Dynamo 的基本示例：可以通过使用 torchdynamo.optimize 装饰器来装饰一个函数或方法，以启用 Dynamo 优化。
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() called with FX graph:") gm.graph.print_tabular() return gm.forward # return a python callable @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 类（如张量子类）
-
数据类型（dtype）
-
设备
-
requires_grad
-
dispatch_key（已应用线程局部的包含和排除规则）
-
维度数
-
尺寸
-
步伐*
全专业化模式允许后端编译器假设整个图为静态的。然而，大多数后端都要求这样做。如果不在动态形状模式下，那些返回动态形状的操作将会导致图中断。
Dynamo 在做些什么？
如果你想更清楚地了解Dynamo的工作原理，可以使用以下方法来运行你的代码：
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),) {} ORIGINAL BYTECODE toy_example example.py line 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 MODIFIED BYTECODE toy_example example.py line 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 possible source 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) If you find the decompiled code is wrong,please submit an issue at https://github.com/youkaichao/depyf/issues.
在顶部你可以看到FX图。接着是函数的原始字节码，然后是Dynamo生成的修改后字节码，以及供参考的反编译源代码。最后，你会看到我们之前提到的guards。
在修改后的字节码中，__compiled_fn_0 是 my_compiler() （编译的图）的返回值。__resume_at_30_1 和 __resume_at_38_2 都是在图中断后（字节码偏移量 30 和 38 处）生成的续行函数，用于继续执行。这些函数的形式如下：
__resume_at_<offset>: ... restore stack state if needed ... JUMP_ABSOLUTE <offset> into toy_example ... original bytecode of toy_example ...
通过生成这个 resume_at 函数，我们强制剩余的部分在新的 Python 帧中执行。这样会递归地触发 Dynamo，在第一次执行到这一点时重新开始捕获。
如何查看Dynamo生成的资源
要查看Dynamo生成的工件，可以使用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 # the guard takes the local variables of an input frame, and tells whether a re-compilation should be triggered. 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)
代码中提到的一些名称有：
-
编译后的函数存储在包含原始函数
toy_example的模块的全局命名空间中。这些函数包括像__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}的字典。因此，你可以看到代码使用L['a']来引用输入变量a。
图中断出现在编译后的 toy_example 代码中，我们需要使用 Python 解释器来选择要执行的下一个图。
注意，我们将一个简单的my_compiler函数作为后端编译器传递，因此子图代码__resume_at_38_2、__resume_at_30_1和__compiled_fn_0仍然是Python代码。这些代码也可以被检查（请忽略函数名，仅关注函数签名和函数体）：
print("source code of __compiled_fn_0:") print(innermost_fn(__compiled_fn_0).__self__.code) print("=" * 60) print("source code of __resume_at_30_1:") print(decompile(__resume_at_30_1)) print("=" * 60) print("source code of __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) # To see more debug info, please use ``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如何转换和优化用户编写的代码：它首先从用户编写的函数中提取计算图，然后将这些图编译为优化后的函数，并将其组装成一个新的函数。这个新函数在功能上与用户编写的代码等同，但经过了优化以实现更快的计算速度。
了解这些功能在内部是如何实现的，请参阅Dynamo 深入解析。