TorchDynamo 基础原理
综述
TorchDynamo 的核心目标是：尽可能提高 JIT 对动态 Python 程序的覆盖率，使更多 PyTorch 程序能够在不修改原始代码的情况下获得编译加速。
对于一个待执行的 Python 函数，TorchDynamo 的基本思路是：
- 尽可能捕获可编译的计算片段。 即使无法将整个函数编译成一个 FX Graph，也可以尝试从中提取部分可编译的计算。
- 尽可能生成少量、较大的 FX Graph。 Dynamo 会贪婪地尝试捕获连续的计算操作。当遇到无法处理的逻辑时，通过 Graph Break 将当前计算图截断，并尽可能从后续位置继续捕获。
- 通过原生 Python 执行逻辑将这些计算图串联起来。 最终，程序由编译后的 FX Graph 和无法捕获的 Python 代码共同组成，从而兼顾编译加速和 Python 的灵活性。
例如：
def foo(x):
x = x.sin()
print("hello")
x = x.cos()
return x假设 print 无法被 Dynamo 捕获，那么最终生成的执行逻辑可以近似理解为：
def optimized_foo(x):
x = compiled_graph_0(x)
print("hello")
x = compiled_graph_1(x)
return x这里的关键是：Graph Break 不意味着整个函数编译失败，而是允许部分代码继续以原生 Python 方式执行。
对于 Python 控制流，例如：
def foo(x, n):
...
for i in range(n):
...
...Dynamo 也会尝试在追踪过程中处理这些控制流。由于JIT时 x,n的值都是可知的，它通常可以通过展开循环来生成更大的 FX Graph，而不一定需要保留原始的 Python for 循环。
因此，Graph Break 的边界并不直接等同于 Python 控制流的边界，而是取决于 Dynamo 能否理解和捕获对应的执行逻辑。
技术背景
CPython 的虚拟机负责执行 Python Bytecode。从概念上，可以将函数的执行过程简化为：
run(frame)其中，frame 包含待执行的 Bytecode、局部变量和运行时状态等信息。
CPython 通过 PEP 523 提供了 Frame Evaluation Hook，允许开发者介入 Python Frame 的执行过程，使用自定义逻辑替代默认的 Frame Evaluation 流程。
TorchDynamo 的核心思路就是利用这个机制：
- 在 Python Frame 开始执行时，拦截其执行过程。
- 基于实际运行时输入，对原始 Bytecode 进行符号执行（Symbolic Execution），尽可能将其中的 PyTorch 计算捕获为 FX Graph。
- 将捕获到的 FX Graph 交给编译 Backend，并生成新的 Bytecode，用对编译结果的调用替代原始计算逻辑，同时保留必要的 Python 执行逻辑。
- 后续执行时，在满足相应 Guards 的情况下复用已经编译的结果。
因此，TorchDynamo 本质上是一个基于 Python Bytecode 的动态 Graph Capture 与 Bytecode Transformation 系统。
从概念上看，这类似于在运行时对 Python 的执行逻辑进行 Monkey Patch。 理论上，也可以考虑从 AST 层面进行程序转换。但 Bytecode 层面的方案能够直接接入 CPython 的 Frame 执行机制，不依赖原始源码并可以结合运行时信息进行动态分析与代码改写，整个流程在Python的现有框架下都可以比较清晰明了的实现。
技术细节
1. torch.compile 的入口
从概念上看，torch.compile(foo) 主要做的是创建一个 Wrapper，在执行原始函数时启用 TorchDynamo 的 Frame Evaluation Hook。
可以近似理解为：
def compiled_foo(*args, **kwargs):
previous = set_eval_frame_cb(dynamo_cb)
try:
return foo(*args, **kwargs)
finally:
set_eval_frame_cb(previous)以上只是示意代码，并非真实 API。
需要注意：torch.compile(foo) 不会立即完成函数的编译，而是在后续使用实际输入调用 Wrapper 时触发 Dynamo 的分析与编译流程。
2. Bytecode 分析与 FX Graph 构建
当 Python VM 准备执行某个 Frame 时，Dynamo 的 Hook 会尝试介入执行流程。
Dynamo 并不是简单地静态扫描 Bytecode，而是实现了一套自定义的 Bytecode Interpreter，通过符号执行模拟 Python 程序的运行过程，并逐步构建 FX Graph。
整个过程可以简化为：
- 创建一个空的 FX Graph，并初始化符号执行所需的状态。
- 依次解释 Python Bytecode，维护操作数栈、局部变量等符号状态。
- 对于可以捕获的 PyTorch 计算，将对应操作加入当前 FX Graph。
- 对于无法处理的操作，尝试触发 Graph Break：
- 将已经捕获的计算构建为独立的 FX Graph，并交给 Backend 编译。
- 生成新的 Bytecode，用于执行编译后的 Graph 和必要的原生 Python 逻辑。
- 尽可能在 Graph Break 之后恢复追踪，继续生成后续 FX Graph。
在无法安全恢复的情况下，Dynamo 也可能放弃后续部分的捕获，直接回退到原生 Python 执行。
3. 函数调用与 Inlining
对于 Python Bytecode 中的 CALL 指令，Dynamo 通常会尝试对被调用函数（Callee）进行 Inlining，以便将跨函数的计算捕获到同一个 FX Graph 中。
其基本逻辑是：
- 遇到
CALL指令时，尝试进入 Callee 的 Bytecode，继续进行符号执行。
- 如果 Inlining 成功，就将 Callee 中可捕获的计算合并到当前 FX Graph，并继续分析 Caller。
- 如果 Inlining 失败，且无法通过其他方式处理该调用，则通常需要在调用处触发 Graph Break：
- 将 CALL 之前捕获的计算编译为 FX Graph。
- 保留对 Callee 的原生调用。
- 在 Callee 返回后，尽可能恢复 Caller 剩余部分的 Graph Capture。
对于回退到原生 Python 执行的 Callee，其 Frame 也可能再次进入 Dynamo 的 Hook，从而获得独立的编译机会。
因此，Dynamo 不仅能够在单个函数内部捕获计算图，还能够通过 Inlining 尝试跨越 Python 函数边界，生成更大的 FX Graph。
4. Guards 与编译结果复用
由于 Dynamo 基于运行时信息进行符号执行，因此捕获到的 FX Graph 可能依赖一些特定假设，例如：
- Tensor 的 Shape、Dtype、Device。
- 某些 Python 变量的具体取值。
- 特定的对象类型或程序状态。
Dynamo 会为这些假设生成 Guards。
后续再次执行函数时：
- 如果 Guards 检查通过，则可以直接复用缓存的编译结果。
- 如果 Guards 检查失败，则可能需要重新进行追踪和编译。
Guards 使得 Dynamo 可以大胆地对运行时信息进行特化，同时保证编译结果只在适用的条件下复用。
核心总结
TorchDynamo 可以概括为：
Frame Evaluation Hook + Bytecode Symbolic Execution + FX Graph Capture + Graph Break / Resume + Guards。
其最重要的设计思想不是要求整个 Python 函数都能够被编译，而是尽可能从动态 Python 程序中提取较大的、可独立编译的计算图，再通过 Python 原生执行机制把它们连接起来。
Dynamo 负责提取计算图，真正的计算图优化与代码生成则由后续的编译 Backend 完成。