| 通过 torch.compile 将 NumPy 代码编译为 C++ 或 CUDA | 您所在的位置：网站首页 › numpy并行算子 › 通过 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 |
| CopyRight 2018-2019 实验室设备网 版权所有 |