如何在 Windows CPU/XPU 上使用 torch.compile

作者: Zhaoqiong Zheng, 徐涵

引言

TorchInductor 是新的编译器后端，它将 TorchDynamo 生成的 FX Graphs 编译成优化的 C++/Triton 内核。

本教程介绍了如何在 Windows CPU/XPU 上通过 torch.compile 使用 TorchInductor 的步骤。

软件安装

现在，我们将逐步向您介绍如何在 Windows CPU/XPU 上使用 torch.compile 的教程。

安装编译器

TorchInductor 优化需要 C++ 编译器，我们以 Microsoft Visual C++ (MSVC) 为例。

1. 下载并安装 MSVC。

1. 在安装过程中，选择 工作负载，然后选择 桌面和移动。 1. 勾选 使用 C++ 的桌面开发 并安装。

注意

Windows CPU inductor 也支持 C++ 编译器 LLVM Compiler 和 Intel Compiler 以获得更好的性能。请查看 备用编译器以获得更好的 CPU 性能。

设置环境

接下来，我们配置环境。

1. 通过 cmd.exe 打开命令行环境。

2. 通过以下命令激活 MSVC

   "C:/Program Files/Microsoft Visual Studio/2022/Community/VC/Auxiliary/Build/vcvars64.bat"
   

3. 创建并激活虚拟环境

4. 安装 PyTorch 2.5 或更高版本用于 CPU。安装 PyTorch 2.7 或更高版本用于 XPU，请参考 Intel GPU 入门指南。

5. 以下是如何在 Windows 上使用 TorchInductor 的示例

   import torch
   device="cpu" # or "xpu" for XPU
   def foo(x, y):
       a = torch.sin(x)
       b = torch.cos(x)
       return a + b
   opt_foo1 = torch.compile(foo)
   print(opt_foo1(torch.randn(10, 10).to(device), torch.randn(10, 10).to(device)))
   

6. 以下是上述示例的输出

   tensor([[-3.9074e-02,  1.3994e+00,  1.3894e+00,  3.2630e-01,  8.3060e-01,
           1.1833e+00,  1.4016e+00,  7.1905e-01,  9.0637e-01, -1.3648e+00],
           [ 1.3728e+00,  7.2863e-01,  8.6888e-01, -6.5442e-01,  5.6790e-01,
           5.2025e-01, -1.2647e+00,  1.2684e+00, -1.2483e+00, -7.2845e-01],
           [-6.7747e-01,  1.2028e+00,  1.1431e+00,  2.7196e-02,  5.5304e-01,
           6.1945e-01,  4.6654e-01, -3.7376e-01,  9.3644e-01,  1.3600e+00],
           [-1.0157e-01,  7.7200e-02,  1.0146e+00,  8.8175e-02, -1.4057e+00,
           8.8119e-01,  6.2853e-01,  3.2773e-01,  8.5082e-01,  8.4615e-01],
           [ 1.4140e+00,  1.2130e+00, -2.0762e-01,  3.3914e-01,  4.1122e-01,
           8.6895e-01,  5.8852e-01,  9.3310e-01,  1.4101e+00,  9.8318e-01],
           [ 1.2355e+00,  7.9290e-02,  1.3707e+00,  1.3754e+00,  1.3768e+00,
           9.8970e-01,  1.1171e+00, -5.9944e-01,  1.2553e+00,  1.3394e+00],
           [-1.3428e+00,  1.8400e-01,  1.1756e+00, -3.0654e-01,  9.7973e-01,
           1.4019e+00,  1.1886e+00, -1.9194e-01,  1.3632e+00,  1.1811e+00],
           [-7.1615e-01,  4.6622e-01,  1.2089e+00,  9.2011e-01,  1.0659e+00,
           9.0892e-01,  1.1932e+00,  1.3888e+00,  1.3898e+00,  1.3218e+00],
           [ 1.4139e+00, -1.4000e-01,  9.1192e-01,  3.0175e-01, -9.6432e-01,
           -1.0498e+00,  1.4115e+00, -9.3212e-01, -9.0964e-01,  1.0127e+00],
           [ 5.7244e-04,  1.2799e+00,  1.3595e+00,  1.0907e+00,  3.7191e-01,
           1.4062e+00,  1.3672e+00,  6.8502e-02,  8.5216e-01,  8.6046e-01]])
   

备用编译器以获得更好的 CPU 性能

为了提升 inductor 在 Windows CPU 上的性能，可以使用 Intel Compiler 或 LLVM Compiler。然而，它们依赖于 Microsoft Visual C++ (MSVC) 的运行时库。因此，第一步应该安装 MSVC。

Intel 编译器

1. 下载并安装 Windows 版本的 Intel 编译器。

2. 通过环境变量 set CXX=icx-cl 设置 Windows Inductor 编译器。

LLVM 编译器

1. 下载并安装 LLVM 编译器，并选择 win64 版本。

2. 通过环境变量 set CXX=clang-cl 设置 Windows Inductor 编译器。

结论

在本教程中，我们介绍了如何在 Windows CPU 上使用 PyTorch 2.5 或更高版本的 Inductor，以及如何在 Windows XPU 上使用 PyTorch 2.7 或更高版本。我们还可以使用 Intel 编译器 或 LLVM 编译器来获得更好的 CPU 性能。