在CUDA中对多维数组进行索引和切片

Question

我尝试使用CUDA为其他算法实现简单的张量类型。我想为这种张量类型添加索引和切片功能，类似于numpy中的索引。例如（在c ++中）

// 3-dimensional array with shape (2, 2, 2)
CUDAArray<double> a(2, 2, 2) = {0, 1, 2, 3, 4, 5, 6, 7};
CUDAArray<double> b = a.index(Range(0, 2), ALL, ALL);
/* b should be:

[[[0, 1],
 [2, 3]],

 [[4, 5],
 [6, 7]]]

(Calculated from numpy a[0:2, :, :])
*/

在这种情况下，涉及两个张量，即输出数组b和输入数组a。如何将按索引选择的元素映射到输出数组的索引，并使用CUDA线程并行执行映射？

我试图从Pytorch中读取源代码的相关部分，后者映射了线性索引，而Arrayfire使用了3-d线程块，但是可悲的是我无法理解其基本公式。我知道我可以使用那些完整的库，但是我真的很想了解它的工作原理。