如何记录torch.compile编译日志、验证效果及提前触发编译？

1. 查看编译后的代码或确认编译生效的方法 #

查看编译产物与日志 #

• 开启编译日志：运行代码前设置环境变量 TORCH_LOGS="compiler,graph"，会输出从Python代码转换到C内核的全流程日志。如果test2无日志，大概率是函数逻辑过于简单，被PyTorch的优化器直接内联或跳过了编译流程，开启graph日志可查看是否被优化。
• 查看中间代码与优化图：使用torch._dynamo.explain()工具，传入编译后的函数和匹配的输入，能获取原始计算图、优化后的计算图以及编译细节，示例代码：

  import torch
  
  def test2(x):
      return x + x
  
  compiled_test2 = torch.compile(test2)
  # 传入符合输入类型的张量触发解释
  orig_graph, opt_graph, explanation = torch._dynamo.explain(compiled_test2)(torch.randn(2))
  print("优化后的计算图:", opt_graph)
  

• 查看生成的C内核：默认用inductor后端时，编译后的C代码会存在/tmp/torchinductor_<你的用户名>目录下，文件名对应编译的函数，可直接打开.cpp文件查看生成的内核代码。

确认编译是否生效 #

• 对比执行耗时：用timeit测试原函数和编译后函数的运行时间，编译后的函数首次执行会有编译开销，后续执行速度会明显提升（CPU场景下简单函数提升可能有限，但能观察到差异）。
• 检查函数类型：编译后的函数会是torch._dynamo.eval_frame.OptimizedModule类型，执行type(compiled_test2)，如果输出和原函数类型不同，说明编译成功。
• 禁用Eager模式验证：设置环境变量TORCHDYNAMO_EAGER=0，如果编译失败会直接报错，反之则说明编译正常生效。

2. 强制提前完成编译（绕过懒加载） #

torch.compile默认是懒加载机制，只有首次执行时才触发编译，要提前完成编译可以这样做：

• 手动触发一次函数调用：传入和实际业务场景匹配的输入张量，比如：

  compiled_test2 = torch.compile(test2)
  # 传入一个形状、类型都和实际使用一致的张量，触发编译
  compiled_test2(torch.randn(3, 3))  # 这一步会完成编译，后续调用直接用编译后的内核
  

• 针对多输入形状场景：如果函数需要处理多种输入形状，可以多次传入不同形状的张量，让PyTorch提前编译好对应形状的内核。
• 模块类提前编译：如果是编译模型模块，可以先初始化模块，编译后调用一次forward方法：

  class MyModel(torch.nn.Module):
      def forward(self, x):
          return x + x
  
  model = torch.compile(MyModel())
  # 提前触发编译
  model(torch.randn(2, 2))
  

内容的提问来源于stack exchange，提问作者Max Ostrowski