j6-plugin-dynamic-block

为 Horizon QAT 自动添加 dynamic_block 标注（Scope 约束版）

目标

在使用 horizon_plugin_pytorch 的基于图（如 PrepareMethod.JIT_STRIP / PrepareMethod.JIT）的 prepare 流程适配模型时，针对动态循环/动态控制流里会触发 function 算子替换 或 算子融合 的代码段，自动补齐 dynamic_block 标注，把该段逻辑定义为独立的 Scope，避免：

• 同一行 function 多次调用在不同 trace 次数/路径下导致的 scale/scope 错位
• 动态部分与静态部分意外融合，造成 forward 报错
• 算子替换/融合在动态控制流中发生时的 量化信息错乱

适用范围

• 你正在为 horizon_plugin_pytorch 做量化适配，并使用 prepare(..., method=PrepareMethod.JIT_STRIP/JIT) 这类基于图的模式。
• 模型 forward 存在运行次数/执行路径不稳定的逻辑，例如：
  • 循环次数由输入/随机数/外部状态决定（numpy.random/数据相关分支等）
• 且该动态逻辑内部包含可能被 prepare 替换/融合的算子（例如某些 torch function、可融合的 Conv+BN(+Add) 等 pattern）。

重要限制：

• 如果你的逻辑是纯静态的（没有数据依赖的 for/while、没有根据输入/随机数决定执行次数），即使其中有可替换或可融合算子，也不应该额外包 dynamic_block。

强约束（本 Skill 的“必须做到”）

• 只标注需要算子替换/融合的逻辑块，不要把整个 for/while 循环都包进去：标注的是“块内的算子替换/融合逻辑”，不是控制流本身。
• Tracer.dynamic_block 的第一个参数必须是当前 nn.Module 实例（通常是 self），不是 function / lambda / 普通 callable：正确用法是 Tracer.dynamic_block(self, "BlockName")，不要写成 Tracer.dynamic_block(some_func, "BlockName")。
• with Tracer.dynamic_block(...) 必须放在循环/分支内部，只包住需要替换/融合的语句，不要把 for/while 控制流包进 with：控制流负责“执行次数/路径”，dynamic_block 负责给其中那段非 module scope 的动态算子逻辑建立稳定 Scope。
• 不要对“已经处于 Module Scope 内”的逻辑额外加 dynamic_block：
  • nn.Module.forward 本身就是一个 Scope。
  • 因此如果动态控制流里只有“子 module 调用”（例如循环里只做 x = self.layer1(x)），通常不需要也不应该再包 dynamic_block。
  • dynamic_block 的主要目标是：动态控制流中出现的 非 module scope 的 function/Tensor 逻辑（例如 x = x + 1、x = torch.relu(x)、x = sub_one(x)）导致的 function 替换/融合错误的问题。
• 动态块必须是稳定可复现的 Scope：
  • 同一段逻辑每次执行应使用同一个 dynamic_block 名称（如果用命名 API）。
  • 不要在同一 forward 里复用同一个名字去标不同语义的块。
• 优先使用 Tracer.dynamic_block(self, "<Name>")（当代码中已引入/可引入 Tracer 且希望显式命名）。
• 如果已有 dynamic_block 标注，不重复嵌套/不改语义：只在缺失且确有必要的位置补齐。

标准改法（优先推荐）

1) 引入 Tracer（推荐写法）

from horizon_plugin_pytorch.fx.jit_scheme import Tracer

2) 在动态循环/动态分支中，包住需要替换/融合的算子段

for _ in range(n):  # n 可能是动态的
    # 只标注需要算子替换/融合的逻辑块（不是整个循环）
    with Tracer.dynamic_block(self, "ConvBnAdd"):
        x = self.conv(x)
        x = self.bn(x)
        x = x + y