通过 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)

/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()

编译后的函数在单核运行时速度提升了 \(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 以保留原始语义。这个问题在 NumPy 中没有本地解决方案，因为 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)。

import torch
import numpy as np

@torch.compile
def numpy_fn(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = np.random.randn(1024, 64)
Y = np.random.randn(1024, 64)
with torch.device("cuda"):
    Z = numpy_fn(X, Y)
assert isinstance(Z, np.ndarray)

如果在同一程序运行中多次执行此函数，可能希望避免所有这些相当昂贵的内存复制。为此，只需要调整 numpy_fn，使其接受 cuda 张量并返回张量。可以使用 torch.compiler.wrap_numpy() 来做到这一点。

@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")
Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
assert Z.device.type == "cuda"

在这里，显式地在 CUDA 内存中创建张量，并将它们传递给函数，该函数在 CUDA 设备上执行所有计算。wrap_numpy 负责在 torch.compile 级别将任何 torch.Tensor 输入标记为具有 np.ndarray 语义的输入。在编译器内部标记张量是非常廉价的操作，因此在运行时不会发生任何数据复制或数据移动。

使用此装饰器，还可以通过 NumPy 代码进行微分！

@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.mean(np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1)))

X = torch.randn(1024, 64, device="cuda", requires_grad=True)
Y = torch.randn(1024, 64, device="cuda")
Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
Z.backward()
# X.grad now holds the gradient of the computation
print(X.grad)

tensor([[-0.0024, -0.0024, -0.0024,  ..., -0.0024, -0.0024, -0.0024],
        [ 0.0014,  0.0014,  0.0014,  ...,  0.0014,  0.0014,  0.0014],
        [ 0.0010,  0.0010,  0.0010,  ...,  0.0010,  0.0010,  0.0010],
        ...,
        [-0.0034, -0.0034, -0.0034,  ..., -0.0034, -0.0034, -0.0034],
        [-0.0118, -0.0118, -0.0118,  ..., -0.0118, -0.0118, -0.0118],
        [-0.0030, -0.0030, -0.0030,  ..., -0.0030, -0.0030, -0.0030]],
       device='cuda:0')

一直在使用 fullgraph=True，因为图中断在此上下文中存在问题。当发生图中断时，需要具体化 NumPy 数组。由于 NumPy 数组没有 device 或 requires_grad 的概念，因此此信息在图中断期间会丢失。

无法通过图中断传播梯度，因为图中断代码可能会执行不知道如何微分的任意代码。另一方面，在 CUDA 执行的情况下，可以像在第一个示例中那样解决此问题，方法是使用 torch.device("cuda") 上下文管理器。

@torch.compile
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    prod = X[:, :, None] * Y[:, None, :]
    print("oops, a graph break!")
    return np.sum(prod, axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")

with torch.device("cuda"):
    Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
assert Z.device.type == "cuda"

oops, a graph break!

在图中断期间，中间张量仍然需要移动到 CPU，但是当在图中断后恢复跟踪时，图的其余部分仍然在 CUDA 上跟踪。鉴于此 CUDA <> CPU 和 CPU <> CUDA 移动，图中断在 NumPy 上下文中相当昂贵，应避免，但至少它们允许跟踪复杂代码段。

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

目录

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

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

将 NumPy 代码编译成 CUDA 代码#

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