关于c++:TorchDynamo-的执行过程

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 的一种解决方案