通过 torch.compile 将 NumPy 代码编译为 C++ 或 CUDA

通过 torch.compile 将 NumPy 代码编译为 C++ 或 CUDA#

来源：Compiling NumPy code into C++ or CUDA via torch.compile

利用 PyTorch 编译器，可以在不修改原始 NumPy 代码的情况下生成高效的融合向量化代码。更重要的是，它还允许在 CUDA 上执行 NumPy 代码，只需将其通过 torch.device("cuda") 下的 torch.compile 运行即可！

将 NumPy 代码编译成并行 C++#

使用 K-Means 算法中的一个步骤作为示例。

import numpy as np

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

创建了包含 \(2000\) 万个随机二维点的合成数据集。可以看到，假设均值选择合适，该函数对所有数据点都返回正确的聚类结果。

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)

[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] WON'T CONVERT kmeans /tmp/ipykernel_2820163/1495106700.py line 3 
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] due to: 
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] Traceback (most recent call last):
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING]   File "/media/pc/data/tmp/cache/conda/envs/tvmz/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 541, in __bool__
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING]     from filelock import FileLock
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] ModuleNotFoundError: No module named 'filelock'
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] 
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING] 
[2023-10-25 16:14:04,224] torch._dynamo.convert_frame: [WARNING]

编译后的函数在单核运行时速度提升了 \(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 vs float32。许多 GPU，尤其是消费级 GPU，在执行 float64 运算时速度较慢。因此，将数据生成改为 float32，原始的 NumPy 代码只会稍微快一点，大约 \(9\%\)，但我们的 CUDA 代码会快 \(40\%\)，相对于普通的 NumPy 代码有 \(11\) 倍的速度提升。

torch.compile() 默认情况下遵循 NumPy 语义，因此它将所有创建操作的默认 dtype 设置为 np.float64。如前所述，这可能会影响性能，因此可以通过设置来更改此默认值。

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 以保留原始语义。这个问题在 NumP y中没有本地解决方案，因为 NumPy 没有设备的概念。话虽如此，我们可以通过为此函数创建一个包装器来解决这个问题，以便它接受 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 进行计算时，我们看到相对于初始的 float32 数组上的 NumPy 实现有 \(200\) 倍的速度提升。

混合使用 NumPy 和 PyTorch

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

目录