TorchDynamo 的执行过程

1 执行流程

设置全局变量 extra_index。这是 PEP 523 设置 PyCodeObject.co_extra 需要的。
初始化 TSS Key。这是 PEP 539 需要的。Tss 的 key 是固定的，这里定义的这个 key 只用于存储 callback，tss 是每个线程私有的（Each thread has a distinct mapping of the key to a void* value）。
- Dynamo callback 都是一样的？

set_eval_frame_py 输入参数中的 args[0] 应该就是 Python 传入的 callback。

在 Python _TorchDynamoContext 中调用这个函数，Python 的方法名是 set_eval_frame。

这是 torch._dynamo.optimize(...) 对应的 context manager。callback 是 Python 代码。

__enter__ 中通过 self.prior = set_eval_frame(self.callback)保存之前的回调。
- set_eval_frame 就是上面提到的 set_eval_frame_py。
__exit__ 中通过 set_eval_frame(self.prior) 恢复之前的 eval frame。
__call__: 当我们调用 optimizer('inductor')(fn) 时，Dynamo 会将 fn 的帧评估函数替换成 Dynamo 自定义的，并且传入回调函数。传入的回调函数会被自定义的帧评估函数调用。回调函数会解析重构 frame 中原有的字节码，并在过程中追踪模型执行图结构。当然了，帧评估时也不是每次都会调用回调函数，例如某个 frame 已经被解析重构过了（cached），此时就会直接执行缓存里已经重构好的代码。

在这里获取 callback，并用于 eval：

_custom_eval_frame_shim

callback = eval_frame_callback_get(): 获取 _TorchDynamoContext 中设置的 callback
- PyThread_tss_get
_custom_eval_frame
- CacheEntry* extra = get_extra(frame->f_code): 获取当前 frame 的 PyCodeObject 的 extra 字段。
  - 该字段用于放置自定义的编译结果。
  - 第一次调用时，extra 为空、还没有设置。下面调用 set_extra 后才会有值。
- result = call_callback(...): 调用 Python 逻辑，分析原始的字节码、捕获静态图、图编译优化、转为字节码。
- extra = create_cache_entry(extra, result): 构造真实的 extra 类型
  - extra_info 的内容是 CacheEntry，其中 code 字段是修改后的字节码。
- set_extra(frame->f_code, extra): 把修改后的 byte code 存到 frame->f_code
  - 后续用到时可以直接取。
- eval_frame_callback_set(callback): 将 callback 保存到 tss
  - PyThread_tss_set
- return eval_custom_code(...)
  - PyFrame_New: 使用 custom_code 创建一个自定义的 frame
  - result = eval_frame_default(tstate, shadow, throw_flag): 调用 Python 默认的 frame eval 执行自定义 frame