作者：Evgeni Burovski, Ralf Gommers 和 Mario Lezcano

Quansight 的工程师在 PyTorch 2.1 中实现了通过 torch.compile 跟踪 NumPy 代码的支持。此功能利用 PyTorch 的编译器生成高效的融合向量化代码，而无需修改您原始的 NumPy 代码。更重要的是，只需在 torch.device("cuda") 下通过 torch.compile 运行 NumPy 代码，它就可以在 CUDA 上执行！

在这篇文章中，我们将介绍如何使用此功能，并提供一些充分利用它的提示和技巧。

将 NumPy 代码编译为并行 C++

我们以 K-Means 算法中的一步作为运行示例。这部分代码引自这本 NumPy 书籍

import numpy as np

def kmeans(X, means):
    return np.argmin(np.linalg.norm(X - means[:, None], axis=2), axis=0)

我们创建一个包含 2000 万个随机 2D 点的合成数据集。我们可以看到，如果均值选择得当，该函数会为所有点返回正确的聚类。

npts = 10_000_000
X = np.repeat([[5, 5], [10, 10]], [npts, npts], axis=0)
X = X + np.random.randn(*X.shape)  # 2 distinct "blobs"
means = np.array([[5, 5], [10, 10]])
np_pred = kmeans(X, means)

在 AMD 3970X CPU 上对该函数进行基准测试，基线时间为 1.26 秒。

现在，编译此函数就像用 torch.compile 包装它并使用示例输入执行它一样简单。

import torch

compiled_fn = torch.compile(kmeans)
compiled_pred = compiled_fn(X, means)
assert np.allclose(np_pred, compiled_pred)

编译后的函数在 1 个核心上运行时可获得 9 倍加速。更好的是，与 NumPy 不同，我们生成的代码确实利用了处理器中的所有核心。因此，当我们在 32 个核心上运行时，可以获得 57 倍加速。请注意，PyTorch 始终使用所有可用核心，除非明确限制，因此这是使用 torch.compile 时获得的默认行为。

我们可以通过设置环境变量 TORCH_LOGS=output_code 运行脚本来检查生成的 C++ 代码。这样做时，我们可以看到 torch.compile 能够将广播和两个归约编译成一个 for 循环，并使用 OpenMP 对其进行并行化。

extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) {
    #pragma omp parallel num_threads(32)
    #pragma omp for
    for(long i0=0L; i0<20000000L; i0+=1L) {
        auto tmp0 = in_ptr0[2L*i0];
        auto tmp1 = in_ptr1[0L];
        auto tmp5 = in_ptr0[1L + (2L*i0)];
        auto tmp6 = in_ptr1[1L];
        // Rest of the kernel omitted for brevity

将 NumPy 代码编译为 CUDA

将我们的代码编译使其在 CUDA 上运行，只需将默认设备设置为 CUDA 即可。

with torch.device("cuda"):
    cuda_pred = compiled_fn(X, means)
assert np.allclose(np_pred, cuda_pred)

通过 TORCH_LOGS=output_code 检查生成的代码，我们看到 torch.compile 生成的是相当易读的 triton 代码，而不是直接生成 CUDA 代码。

def triton_(in_ptr0, in_ptr1, out_ptr0, XBLOCK : tl.constexpr):
    xnumel = 20000000
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (2*x0), xmask)
    tmp1 = tl.load(in_ptr1 + (0))
    // Rest of the kernel omitted for brevity

在 RTX 2060 上运行这个小代码片段，比原始 NumPy 代码有了 8 倍加速。这确实有提升，但考虑到我们在 CPU 上看到的加速效果，并不特别令人印象深刻。让我们看看如何通过一些小的改动来最大限度地利用我们的 GPU。

float64 对比 float32。许多 GPU，特别是消费级 GPU，在运行 float64 上的操作时相当缓慢。因此，将数据生成更改为 float32 后，原始 NumPy 代码仅略有加速，大约 9%，但我们的 CUDA 代码却提速 40%，比纯 NumPy 代码快了 11 倍。

torch.compile 默认遵循 NumPy 语义，因此在其所有创建操作中都使用 np.float64 作为默认 dtype。如前所述，这可能会影响性能，因此可以通过设置以下内容来更改此默认值

from torch._dynamo import config
config.numpy_default_float = "float32"

CPU <> CUDA 拷贝。11 倍加速固然不错，但与 CPU 上的性能相比仍然有很大差距。这是由 torch.compile 在幕后进行的小转换引起的。上面的代码接受 NumPy 数组并返回 NumPy 数组。所有这些数组都在 CPU 上，但计算是在 GPU 上执行的。这意味着每次调用函数时，torch.compile 都必须将所有这些数组从 CPU 复制到 GPU，然后再将结果复制回 CPU 以保留原始语义。在 NumPy 中没有这个问题的原生解决方案，因为 NumPy 没有 device 的概念。话虽如此，我们可以通过创建一个包装函数来解决这个问题，使其接受 PyTorch 张量并返回 PyTorch 张量。

@torch.compile
def tensor_fn(X, means):
    X, means = X.numpy(), means.numpy()
    ret = kmeans(X, means)
    return torch.from_numpy(ret)

def cuda_fn(X, means):
    with torch.device("cuda"):
        return tensor_fn(X, means)

这个函数现在接受 CUDA 内存中的张量并返回 CUDA 内存中的张量，但函数本身是用 NumPy 编写的！torch.compile 使用 numpy() 和 from_numpy() 调用作为提示，并将其优化掉，内部只是使用 PyTorch 张量，完全不移动内存。当我们将张量保留在 CUDA 中并在 float32 中执行计算时，相对于初始 NumPy 在 float32 数组上的实现，我们看到了 200 倍加速。

混合使用 NumPy 和 PyTorch。在这个示例中，我们必须编写一个小适配器来将张量转换为 ndarray，然后再转换回张量。在混合使用 PyTorch 和 NumPy 的程序中，将张量转换为 ndarray 通常实现为 x.detach().cpu().numpy()，或者简单地使用 x.numpy(force=True)。由于在 torch.compile 下运行时，我们可以在 CUDA 中运行 NumPy 代码，我们可以将这种转换模式实现为调用 x.numpy()，就像我们上面所做的那样。这样做并在 device("cuda") 下运行生成的代码，将从原始 NumPy 调用中生成高效的 CUDA 代码，而无需将数据从 CUDA 复制到 CPU。请注意，生成的代码不能在没有 torch.compile 的情况下运行。要在 eager 模式下运行，需要回退到 x.numpy(force=True)。

更多加速技巧

一般建议。我们展示的 CUDA 代码已经相当高效，但确实运行示例比较短。在处理更大的程序时，我们可能需要调整其中的部分代码以提高效率。一个好的起点是 torch.compile 的多个教程和常见问题解答 (FAQ)。这些资源展示了检查跟踪过程以及如何识别可能导致速度变慢的代码的多种方法。

编译 NumPy 代码时的建议。NumPy 虽然与 PyTorch 相似，但使用方式通常非常不同。在 NumPy 中进行计算然后根据数组中的值执行 if/else，或者通过布尔掩码执行原地操作，这种情况相当常见。这些结构虽然受到 torch.compile 的支持，但会影响其性能。像以无分支方式编写代码以避免图中断 (graph breaks)，或避免原地操作 (in-place ops) 等更改可以大有帮助。

为了编写快速的 NumPy 代码，最好避免循环，但有时它们是不可避免的。在跟踪循环时，torch.compile 会尝试完全展开循环。这有时是可取的，但有时甚至不可能，例如当我们有动态停止条件时，就像在 while 循环中一样。在这些情况下，最好只编译循环体，可能一次编译几次迭代（循环展开）。

调试 NumPy 代码。当涉及到编译器时，调试相当棘手。要确定遇到的错误是 torch.compile 错误还是程序本身的错误，您可以通过将 NumPy 导入替换为 import torch._numpy as np 来在不使用 torch.compile 的情况下执行您的 NumPy 程序。这仅应用于调试目的，绝不能替代 PyTorch API，因为它慢得多，而且作为私有 API，可能会在没有通知的情况下更改。另请参阅此 FAQ 以获取其他技巧。

NumPy 与 torch.compile 下的 NumPy 之间的差异

NumPy 标量。NumPy 在几乎所有 PyTorch 会返回 0 维张量的情况下（例如来自 np.sum）都返回 NumPy 标量。在 torch.compile 下，NumPy 标量被视为 0 维数组。这在大多数情况下都没有问题。它们的行为唯一不同的是当 NumPy 标量被隐式用作 Python 标量时。例如，

>>> np.asarray(2) * [1, 2, 3]  # 0-D array is an array-like
array([2, 4, 6])
>>> u = np.int32(2)
>>> u * [1, 2, 3]              # scalar decays into a Python int
[1, 2, 3, 1, 2, 3]
>>> torch.compile(lambda: u * [1, 2, 3])()
array([2, 4, 6])               # acts as a 0-D array, not as a scalar ?!?!

如果我们编译前两行，我们会看到 torch.compile 将 u 视为 0 维数组。要恢复 eager 语义，我们只需进行显式类型转换。

>>> torch.compile(lambda: int(u) * [1, 2, 3])()
[1, 2, 3, 1, 2, 3]

类型提升和版本控制。NumPy 的类型提升规则有时可能有点令人惊讶

>>> np.zeros(1, dtype=np.int8) + 127
array([127], dtype=int8)
>>> np.zeros(1, dtype=np.int8) + 128
array([128], dtype=int16)

NumPy 2.0 正在改变这些规则，使其更接近 PyTorch 的规则。相关的技术文档是 NEP 50。torch.compile 已提前实现 NEP 50，而不是即将废弃的规则。

总的来说，torch.compile 中的 NumPy 遵循 NumPy 2.0 预发布版本。

NumPy 之外：SciPy 和 scikit-learn

与这项使 torch.compile 理解 NumPy 代码的努力并行，Quansight 的其他工程师设计并提出了一种在 scikit-learn 和 SciPy 中支持 PyTorch 张量的方法。这受到了这些库的其他维护者的热烈欢迎，因为事实证明，使用 PyTorch 作为后端通常会带来显著的加速。这两个项目现在已经在许多 API 和子模块中合并了对 PyTorch 张量的初步支持。

这为迈向未来 PyTorch 张量可以在 Python 数据生态系统中的其他库中使用奠定了基础。此外，这将使得这些其他库可以在 GPU 上运行，甚至可以编译混合这些库和 PyTorch 的代码，就像我们在本文中讨论的那样。

如果您想了解更多关于这项工作的信息，如何使用它，或者如何帮助推动其进展，请参阅这篇博客文章。

结论

PyTorch 自成立以来就致力于成为一个与 Python 生态系统其他部分兼容的框架。支持编译 NumPy 程序，并建立对其他重要库做同样事情所需的工具，是朝着这个方向迈出的另外两步。Quansight 和 Meta 继续携手合作，改善 PyTorch 与生态系统其他部分之间的兼容性。

Quansight 感谢 Mengwei、Voz 和 Ed 在将我们的工作与 torch.compile 集成方面的宝贵帮助。我们还要感谢 Meta 资助本项目以及之前在改进 PyTorch 内的 NumPy 兼容性方面的工作，以及促成在 scikit-learn 和 SciPy 中支持 PyTorch 的项目。这些是朝着将 PyTorch 打造成开源 Python 数据生态系统中的首选框架迈出的巨大一步。