[Bug已解决] torch.compile 不区分同名但 mutates_args 不同的自定义算子导致 FX 图错误解决方案
一、现象长什么样
你用torch.library自定义了两个算子，名字相同、但一个会原地修改输入（mutates_args不同），另一个不修改：
import torch from torch.library import custom_op # 算子 A：不修改输入 @custom_op("mylib.foo", mutates_args=()) def foo(x: torch.Tensor) -> torch.Tensor: return x + 1 # 算子 B：同名但会原地修改第 0 个输入 @custom_op("mylib.foo", mutates_args={"x"}) # 注意同名为 foo，但 mutates_args 不同 def foo(x: torch.Tensor) -> torch.Tensor: x.add_(1) return x然后你torch.compile一个同时用到这两种「同名不同语义」算子的模型，结果编译出的 FX 图把它们的副作用搞混了——比如该原地修改的没修改、或不该修改的被改了，或者 Dynamo 的别名分析（aliasing analysis）判断错，生成错误的图。 本 issue（pytorch/pytorch#187319）说的是：torch.compile的 FX 图在区分「同名但mutates_args不同」的自定义算子时存在缺陷，导致图的别名/副作用分析出错。 本文聚焦：自定义算子的mutates_args是什么、为什么同名不同副作用会让编译器困惑、怎么正确注册以避免被混淆。
二、背景：custom_op 与 mutates_args
torch.library.custom_op让你注册一个「PyTorch 不认识、但你能提供 forward/backward」的算子。关键参数mutates_args告诉编译器：这个算子会原地修改哪些输入参数。
@custom_op("mylib.foo", mutates_args={"x"}) # 声明：会改 x def foo(x): x.add_(1) return xmutates_args对torch.compile至关重要：编译器做**别名分析（aliasing / mutation analysis）**时，需要知道「调用 foo 后，x 的内容变了」，才能正确安排后续算子（比如不能把 x 当「未变」的值复用，不能错误地做算子融合把 x 的读取合并掉）。 问题：如果两个同名算子mutates_args不同（一个改、一个不改），编译器在构建 FX 图时可能只按「名字」缓存了某一种副作用假设，没按mutates_args区分，于是把两种语义混为一谈 → 图错误。
三、为什么同名不同 mutates_args 会让编译器困惑
Dynamo 在 trace 自定义算子时，会从算子的schema / 注册信息读取mutates_args，据此更新别名分析状态。缺陷在于：
- 按名字做图节点缓存/特化：FX 图里节点以算子名（如
mylib::foo）标识，编译器可能用「名字」作为别名分析的 key，而不是「名字 + mutates_args 组合」； - 两种注册共存时状态冲突：当同一个
mylib.foo注册了两次（不同mutates_args），后注册的覆盖前者的副作用信息，或编译器取了错误的那份； - 生成错误的图：例如把「会改 x」误判为「不改 x」，于是后续
x的读取被错误复用旧值，或错误地与别处融合 → 数值错 / 崩溃。 本质：编译器的「算子副作用」状态 key 没包含mutates_args，导致同名但副作用不同的算子被当成同一个。
四、最小可运行复现（带守卫，结构示意）
下面演示「同名不同 mutates_args 的 custom_op」注册与 compile（实际图错误需特定版本，用守卫+说明）：
import torch from torch.library import custom_op # 注册一个「不修改输入」的同名算子 @custom_op("mylib.foo", mutates_args=()) def foo_pure(x: torch.Tensor) -> torch.Tensor: return x + 1 # 注册「同名但会原地修改输入」的算子（不同 mutates_args） # 注意：真实场景应给不同名字；这里演示同名冲突为何危险 @custom_op("mylib.foo", mutates_args={"x"}) def foo_mut(x: torch.Tensor) -> torch.Tensor: x.add_(1) return x # 编译器若只看名字，会混淆两者的副作用假设 def demo_compile(): if not torch.cuda.is_available(): print("[skip] 无 GPU，仅说明 custom_op 注册") def f(x): y = foo_pure(x) # 编译期可能误用 foo_mut 的副作用假设 return y try: g = torch.compile(f) print("compile 成功（但图可能含副作用误判）") except Exception as e: print("compile 报错:", e) if __name__ == "__main__": demo_compile()要点：同名不同mutates_args是反模式；正确做法是用不同名字。下面给正解。
五、解决方案一：同名算子用不同名字（最关键）
永远不要为「语义/副作用不同」的算子注册同名。给它们不同的算子名，torch.compile就能正确区分：
import torch from torch.library import custom_op # 不修改输入 → 名字 foo_pure @custom_op("mylib.foo_pure", mutates_args=()) def foo_pure(x: torch.Tensor) -> torch.Tensor: return x + 1 # 原地修改 → 名字 foo_mut @custom_op("mylib.foo_mut", mutates_args={"x"}) def foo_mut(x: torch.Tensor) -> torch.Tensor: x.add_(1) return x # 编译时，两个名字的副作用信息互不干扰 def f(x): y = foo_pure(x.clone()) # 不修改 z = foo_mut(x) # 修改 x return y + z compiled = torch.compile(f) print("用不同名字，编译器正确区分副作用")这是从根上规避——不触发「同名不同 mutates_args」的混淆路径。
六、解决方案二：正确声明 mutates_args + setup_context / backward
若你确实需要「会改输入」的算子，按官方模板完整声明，让编译器拿到准确副作用信息：
from torch.library import custom_op, register_autograd @custom_op("mylib.foo_mut", mutates_args={"x"}) def foo_mut(x: torch.Tensor) -> torch.Tensor: x.add_(1) return x # 提供 backward（若需训练） @register_autograd(foo_mut, autograd_fn=None, setup_context_fn=None) def foo_mut_autograd(ctx, x, output): # 自定义反向（略） ... # 编译器读到 mutates_args={"x"}，正确标记 x 被改关键点：mutates_args必须与实际代码一致——声明改了就必须真的原地改，声明不改就绝不能改，否则编译器的别名分析基于错误声明，照样出问题（这是你自己的责任，不是编译器 bug）。
七、解决方案三：用 tags / schema 明确告知
对更复杂的自定义算子，用tags和精确 schema 帮助编译器：
import torch from torch.library import custom_op, impl @custom_op("mylib.foo_mut", mutates_args={"x"}, device_types="cuda") def foo_mut(x: torch.Tensor) -> torch.Tensor: x.add_(1) return x @impl(foo_mut, "CompositeExplicitAutograd", device_types="cuda") def foo_mut_impl(x): x.add_(1) return x精确的device_types/impl注册减少编译器在「同名多实现」下的歧义。
八、解决方案四：升级并验证图正确性
#187319 是编译器缺陷，官方会修（让别名分析 key 包含mutates_args）。升级：
pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124判断修复：注册同名不同mutates_args后torch.compile的 FX 图能正确区分副作用。修复前，用不同名字是唯一可靠做法。验证图正确性的小技巧：
def f(x): y = foo_pure(x) return y # 对比 eager 与 compiled 的数值一致性 xc = x.clone() eager_out = f(x) comp_out = torch.compile(f)(xc) assert torch.allclose(eager_out, comp_out), "图与 eager 不一致，副作用分析错了"九、排查清单
- 自定义算子 +
torch.compile后数值错/图怪 → 查是否「同名不同mutates_args」（#187319）。 - 用不同名字：同名不同副作用是反模式，立即改名规避混淆。
mutates_args必须与实际代码一致（改了才声明改，否则声明空）。- 完整注册：
setup_context/register_autograd/impl齐全，编译器拿到准确信息。 - 升级：关注该 bug 修复（别名分析 key 含 mutates_args）。
- 验证：eager 与 compiled 数值
allclose，不一致即副作用分析错。
十、小结
torch.compile FX graph does not distinguish custom ops with the same name but different mutates_args（#187319）的本质是：torch.compile的别名/副作用分析在区分自定义算子时，key 用的是「算子名」而非「名字 + mutates_args」，当两个同名算子mutates_args不同（一个原地改输入、一个不改）时，编译器把它们当成同一副作用，生成错误的 FX 图。 应对：
- 用不同名字：同名不同副作用是反模式，改名让编译器正确区分（最根本）；
- mutates_args 如实声明：改了才声明改，否则声明空，编译器基于声明分析；
- 完整注册：
impl/setup_context/register_autograd齐全； - 升级 + 验证：关注修复；用 eager vs compiled 数值
allclose验证图正确。 记住：自定义算子的mutates_args是编译器做副作用分析的唯一依据。同名不同副作用会让编译器「张冠李戴」，永远用不同名字区分。
