在 Python 中使用 GPU 对大型数组执行 tan() 和 arctan() 的最快方法？

Question

这本质上是这个问题的扩展。那里提供的答案说使用数学库来实现三角函数。不幸的是，数学库仅适用于标量，不适用于数组。当我尝试将 njit 与 np.tan() 和 np.arctan() 一起使用时，它使用我的 CPU，而不是我的 GPU：

import numpy as np
import numba as nb

@nb.njit(fastmath=True, parallel=True)
def f1(a):
    np.tan(a)
    return np.arctan(a)

我有非常大的数组（8000 万行，300 万列），我需要同时执行 tan() 和 arctan() 操作。它不会一次全部装入内存，所以我很乐意一次将其分批分成几行。如果可能的话，我想使用我的 GPU 来执行所有这些三角函数。 Numba 会为此工作吗？像 PyTorch 这样的东西在这里可以工作吗？或者它是否有太多的开销？使用 GPU 对大型数组执行三角函数的最快方法是什么？

Answer 1

如果您执行的唯一计算是反正切，则从磁盘读取 720 TB 将成为瓶颈。最好将反正切嵌入到生成 720 TB 数据的算法或真正使用该数据的算法中。否则，仅仅为 arctan 传输 720 TB 的数据就会破坏 CPU 的所有缓存，成为 SSD 的瓶颈，并消耗大量能源。但是嵌入在数据生成算法最后一行的 arctan 将避免上述所有负面部分，除非必须更改现有代码。

假设 arctan 仅包含约 30 FLOPs，则每 4 个字节只需要 30 FLOPS，或者 4 GB/s 需要 30 GFLOP/s。如果您的SSD额定串行读取性能为4GB/s，那么您只需要30GFLOP/s的CPU核心。当前的处理器可为并行代码提供约 5 倍的性能，为标量版本提供约 0.5 倍的性能。可能需要第二个 CPU 核心，而 CPU 上的 Numba 已经很好地实现了这一点。