共计 2282 个字符，预计需要花费 6 分钟才能阅读完成。
1 执行流程
1.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 都是一样的？
1.2 设置 eval frame 以代替默认的 eval frame 函数
set_eval_frame_py 输出参数中的 args[0] 应该就是 Python 传入的 callback。
在 Python _TorchDynamoContext 中调用这个函数，Python 的办法名是 set_eval_frame。
1.2.1 函数的次要作用
- 将传入的 callback 保留到 thread specific storage
- 如果需要的话，设置 eval_frame 为 custom_eval_frame_shim
- 返回之前保留在 tss 的、旧的 callback
1.2.2 set_eval_frame_py 执行流程
- set_eval_frame
- increment_working_threads，如果需要的话
- enable_eval_frame_shim
- tstate->interp->eval_frame = &custom_eval_frame_shim
1.3 _TorchDynamoContext: 用于设置转换字节码的 callback 函数
这是 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），此时就会间接执行缓存里曾经重构好的代码。
1.4 执行 Python 字节码时，调用 custom_eval_frame_shim 进行 eval frame
在这里获取 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
2 TODO
- torch.fx
- 解析字节码、生成新的字节码、设置 co_extra 的过程。
3 参考资料
3.1 cpython
- Python behind the scenes
- Python behind the scenes 中文翻译
- PEP 523 – Adding a frame evaluation API to CPython
- PEP 539 – A New C-API for Thread-Local Storage in CPython
- Initialization, Finalization, and Threads: PyThread_tss_create, PyThread_tss_set, PyThread_tss_get
3.2 Dynamo
- PyTorch 2.0 之 Dynamo: eager 模式的救星，减速背地的假相
- TorchDynamo 初探：Python ByteCode 的动静批改
3.3 torch.fx
- 官网文档
- PyTorch 新技能解锁：torch.fx
- 用沐神的办法浏览 PyTorch FX 论文
- 深度学习框架量化感知训练的思考及 OneFlow 的一种解决方案
