[Bug已解决] torch.compile 自定义算子 DCE 在动态形状下失效（CUDA，DISABLED test）解决方案

[Bug已解决] torch.compile 自定义算子 DCE 在动态形状下失效CUDADISABLED test解决方案一、现象长什么样PyTorch CI 里有一条DISABLED test_custom_op_dce_dynamic_shapes_cuda (__main__.DynamicShapesCodegenGPUTests)它测试的是当你用torch.library.custom_op注册了自定义算子并在torch.compile(dynamicTrue)动态形状下编译时Inductor 的**死代码消除DCE, Dead Code Elimination**能否正确工作。被DISABLED说明在 CUDA 动态形状下自定义算子的 DCE 有 bug——可能该删的没删性能问题或删错了正确性问题比如把有副作用的自定义算子误删。本文讲清楚自定义算子怎么注册、DCE 是什么、为什么动态形状下会踩坑以及你写自定义算子时怎么规避。二、自定义算子custom_op是什么有时 PyTorch 内置算子不够用你需要注册自己的算子并让它能被torch.compile识别、编译。现代做法是torch.library.custom_opimport torch # 注册一个无副作用的自定义算子对输入每个元素 1 torch.library.custom_op(mylib::add_one, mutates_args()) def add_one(x: torch.Tensor) - torch.Tensor: return x 1 # 声明它在 CUDA 上的实现fake/meta 用于编译期形状推导 add_one.register_fake def _add_one_fake(x): return torch.empty_like(x)mutates_args()告诉编译器「这个算子不修改任何输入」于是编译器可以放心做 DCE / 重排。如果这里声明错了DCE 就可能误删或误保留。三、DCE死代码消除是什么编译优化里DCE 指「删掉对最终结果没有贡献的计算」。比如y add_one(x) # y 后面没被用到 z x * 2 # 只有 z 被返回 return z编译器可以删掉add_one(x)这次调用因为 y 没用省一次 kernel。DCE 是 Inductor 的重要优化。但 DCE 的前提是编译器能确定「这个算子没有副作用、删了不影响结果」。对自定义算子它依赖你声明的mutates_args/is_pure等元信息。如果元信息在动态形状下被错误推导DCE 就可能把「其实是纯函数、该删」的保留下来浪费性能但不算 bug或把「有副作用、不该删」的删掉正确性问题严重。测试DISABLED说明在 CUDA 动态形状下出现了后者或前者导致的断言失败。四、可运行注册自定义算子并在 compile 下用下面脚本演示一个正确声明的自定义算子在动态形状下编译无 GPU 时演示 CPU 逻辑GPU 部分会优雅跳过import torch torch.library.custom_op(mylib::scale, mutates_args()) def scale(x: torch.Tensor, factor: float) - torch.Tensor: return x * factor scale.register_fake def _scale_fake(x, factor): return torch.empty_like(x) def model(x): # unused 不会被返回 - 理想情况下被 DCE 删掉 unused scale(x, 2.0) return x * 3.0 if __name__ __main__: device cuda if torch.cuda.is_available() else cpu x torch.randn(4, 8, devicedevice) compiled torch.compile(model, dynamicTrue) out compiled(x) print(编译输出形状, out.shape) print(提示确认 scale 的 unused 调用已被 DCE 删掉性能 且没有误删有副作用的算子正确性)如果你怀疑 DCE 误删了有副作用的算子可以加一个「标记副作用」的版本对比见下节。五、解决方案一正确声明副作用避免被误删如果你的自定义算子有副作用比如修改了全局状态、写入外部 buffer、或mutates_args非空必须如实声明否则 DCE 可能删掉它import torch # 这个算子会修改输入有副作用 torch.library.custom_op(mylib::inplace_add, mutates_args(x,)) def inplace_add(x: torch.Tensor, v: float) - None: x.add_(v) # 原地修改 x # 有副作用 - 编译器知道不能 DCE 删掉它关键点mutates_args要如实填写所有被原地修改的参数。填错该填没填就会让 DCE 误以为它无害而删掉导致逻辑错误。如果需要让编译器知道「它有副作用、必须保留」还可标记为非纯torch.library.custom_op(mylib::side_effect_op, mutates_args(), device_typescuda) def side_effect_op(x: torch.Tensor) - torch.Tensor: # 即使不改输入但有外部副作用 return x # 通过 register_kernel / 标注让它不被 DCE具体 API 见版本文档六、解决方案二动态形状下用 fullgraphFalse 容忍如果 DCE 在动态形状下不可靠可以让编译图在自定义算子处「断裂」退回 eager避免 DCE 错误处理它compiled torch.compile(model, dynamicTrue, fullgraphFalse)代价是图被切开性能略降但自定义算子走 eagerDCE 不再碰它正确性有保证。七、解决方案三给自定义算子注册 autograd 显式 fake自定义算子要在torch.compile下完全可靠需要补齐register_fake编译期形状推导必须有autograd 反向如果参与训练正确的mutates_args/ 纯度声明。import torch torch.library.custom_op(mylib::square, mutates_args()) def square(x: torch.Tensor) - torch.Tensor: return x * x square.register_fake def _square_fake(x): return torch.empty_like(x) # 反向用于训练 square.register_autograd def _square_autograd(ctx, x, needs_grad): def backward(grad_output): return grad_output * 2 * x return backward # 这样 square 在 compile 训练下都可靠缺失任何一环尤其register_fake动态形状下编译器就无法正确推导DCE / 重排都可能出错。八、解决方案四升级 PyTorchtest_custom_op_dce_dynamic_shapes_cuda是 Inductor 对自定义算子 DCE 支持持续完善中的 Known Issue。新版本可能已修复动态形状下的 DCE 行为。查看版本import torch print(PyTorch, torch.__version__)九、如何判断你踩的是同一条你用了torch.library.custom_op注册的算子用torch.compile(..., dynamicTrue)现象是「该删的没删变慢」或「有副作用的被删了结果错」关掉 compileeager或用fullgraphFalse后恢复正常。命中即说明踩中该 disabled 测试覆盖的问题。十、小结DISABLED test_custom_op_dce_dynamic_shapes_cuda揭示自定义算子在 CUDA 动态形状下Inductor 的死代码消除不可靠。应对如实声明副作用mutates_args写全所有被原地修改的参数避免被误删动态形状不可靠时fullgraphFalse让自定义算子走 eager补齐register_fake autograd让编译器在动态形状下能正确推导升级到修复该 DCE bug 的 PyTorch 版本。自定义算子是扩展 PyTorch 的利器但编译器对它的所有判断都建立在你提供的元信息上。元信息填错优化器就会「好心办坏事」——DCE 删掉你以为会执行的逻辑。诚实声明副作用是自定义算子进torch.compile的第一原则。