如何在Python中并行化GPU上的嵌套for循环？

Question

问题

我有三个嵌套for循环，并希望使用Python 3.6在我的（支持CUDA）GPU上并行运行它们。评估顺序并不重要，但必须存储输出。我怎样才能做到这一点？请注意，每个for循环的范围是不同的。

代码

def f(i, j, k):
    if i < j < k:
        return np.log(1 + i + j + k)
    else:
        return (i, j, k)

for i in range(1000):
    for j in range(2000):
        for k in range(3000):
            f(i, j, k) # Save the output???

额外细节

这甚至可以在GPU上并行化吗？我可以使用哪些包裹？如何存储输出而不必等待其他循环编译？我应该将[i, j, k, f(i, j, k)]附加到列表中，然后将它们放入有组织的数组中吗？如果f的输出为float16，则输出量很容易失控。我怎么能克服这个？上面的函数组成，因此它将具有混合形式参数，if语句，并且纯粹为了教育目的调用NumPy函数。 如果使用for循环不是正确的方法，请赐教。我是否应该创建一个所有可能(i, j, k)的列表并向GPU发送几千个来评估它？请考虑内存限制。即使在内存中存储100 x 200 x 300个整数也很昂贵。

在框外

如何将此问题转换为线性代数问题，以便利用Numpy或Numba？