cuda中的嵌套循环模置换

Nested loops modulo permutation in cuda

我需要对从数组中取出的三元组执行一个函数并将结果添加到直方图中,但我想避免排列,因为该函数在 [F(i,j,k) = F( j,i,k) 等等]。

通常我会这样编码:

def F(int i, int j, int k){
    int temp_index;
    /* Do something */
    return temp_index;
}

for(int i=0;i<N;i++){
    for(int j=i+1;j<N;j++){
        for(int k=j+1;k<N;k++){
            hist[F(i,j,k)]++;
        }
    }
}

由于 N 很大(大约 10^5),我想使用 cuda 在 GPU 上执行此操作。

我已经编写了代码在 GPU 上调用此函数,但我不知道如何防止多次调用相同的索引三元组。到目前为止,我用 3 维网格调用 cuda,例如:

__global__ void compute_3pcf(float *theta, float *hist) {
  int i,j,k;
  i = blockIdx.x*blockDim.x + threadIdx.x;
  j = blockIdx.y*blockDim.y + threadIdx.y;
  k = blockIdx.z*blockDim.z + threadIdx.z;
  if(i>=j || j>=k) return;
  atomicAdd(&hist[F(i,j,k)],1);
}

int main(){
  /*
  Allocation of memory and cudaMemcpy
  */
  dim3 grid((N+15)/16,(N+7)/8,(N+7)/8);
  dim3 block(16,8,8);

  //Launch on GPU
  compute_3pcf<<<grid,block>>>(d_theta, d_hist);
}

但是,现在对于每个组合 (i,j,k),都会启动一个新线程,然后中止,这对我来说似乎效率很低,因为只有 1/6 的线程执行实际计算。我想要的是这样的:

__global__ void compute_3pcf(float *theta, float *hist) {
  int i,j,k,idx;
  idx = blockIdx.x*blockDim.x + threadIdx.x;
  i = H_i(idx);
  j = H_j(idx,i);
  k = H_k(idx,j);
  atomicAdd(&hist[F(i,j,k)],1);
}

int main(){
  /*
  Allocation of memory and cudaMemcpy
  */
  long long int N_combinations = N*(N-1)*(N-2)/6;
  long int grid = (N_combinations+1023)/1024;
  int block = 1024;
  //Launch on GPU
  compute_3pcf<<<grid,block>>>(d_theta, d_hist);
}

但是,我无法找到函数 H_i、H_j、H_k。如果有人能告诉我如何解决或避免这个问题,我将不胜感激。

编辑:直方图包含大约 10^6 个 bin,因此我不能在共享内存中为每个块提供一个直方图,就像 cuda 的示例代码中那样。相反,它位于 GPU 的全局内存中。

[免责声明 -- 这只是部分答案和正在进行的工作,并回答了相关问题,同时仅暗示了实际问题的解决方案]

在考虑算法和代码之前,了解问题的数学特征很有用。如果我们在 Python 中查看伪代码的输出(并注意这包括原始问题不包含的对角线条目),我们会在 5x5x5 的情况下看到:

N = 5
x0 = np.zeros((N,N,N), dtype=np.int)

idx = 1
for i in range(0,N):
    for j in range(i,N):
        for k in range(j,N):
            x0[i,j,k] = idx
            idx += 1

print(x0)

我们得到:

[[[ 1  2  3  4  5]
  [ 0  6  7  8  9]
  [ 0  0 10 11 12]
  [ 0  0  0 13 14]
  [ 0  0  0  0 15]]

 [[ 0  0  0  0  0]
  [ 0 16 17 18 19]
  [ 0  0 20 21 22]
  [ 0  0  0 23 24]
  [ 0  0  0  0 25]]

 [[ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0 26 27 28]
  [ 0  0  0 29 30]
  [ 0  0  0  0 31]]

 [[ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0 32 33]
  [ 0  0  0  0 34]]

 [[ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0 35]]]

即独特的条目形成一系列尺寸递减的堆叠上三角矩阵。正如评论中所确定的,非零条目的数量是 tetrahedral number, in this case for n = 5, the tetrahedral number Tr[5] = 5*(5+1)*(5+2)/6 = 35 entries, and the non zero entries fill a tetrahedral shaped region of the hypermatrix in three dimensions (best illustration here) 并且如原始问题中所述,问题中所有索引的排列在功能上都是相同的,这意味着有六个 (3P3) 在功能上相同立方超矩阵中的对称四面体区域。你可以自己确认一下:

x1 = np.zeros((N,N,N), dtype=np.int)
idx = 1
for i in range(0,N):
    for j in range(0,N):
        for k in range(0,N):
            if (i <= j) and (j <= k):
                x1[i,j,k] = idx
                x1[i,k,j] = idx
                x1[j,i,k] = idx
                x1[j,k,i] = idx
                x1[k,i,j] = idx
                x1[k,j,i] = idx
                idx += 1

print(x1)

给出:

[[[ 1  2  3  4  5]
  [ 2  6  7  8  9]
  [ 3  7 10 11 12]
  [ 4  8 11 13 14]
  [ 5  9 12 14 15]]

 [[ 2  6  7  8  9]
  [ 6 16 17 18 19]
  [ 7 17 20 21 22]
  [ 8 18 21 23 24]
  [ 9 19 22 24 25]]

 [[ 3  7 10 11 12]
  [ 7 17 20 21 22]
  [10 20 26 27 28]
  [11 21 27 29 30]
  [12 22 28 30 31]]

 [[ 4  8 11 13 14]
  [ 8 18 21 23 24]
  [11 21 27 29 30]
  [13 23 29 32 33]
  [14 24 30 33 34]]

 [[ 5  9 12 14 15]
  [ 9 19 22 24 25]
  [12 22 28 30 31]
  [14 24 30 33 34]
  [15 25 31 34 35]]]

这里应该很明显,您可以沿任何平面对超矩阵进行切片并得到一个对称矩阵,并且它可以通过同一基本四面体超矩阵的六个排列中的任何一个的一组反射来构造。

最后一部分很重要,因为我现在将重点关注您问题中的另一种排列。它在功能上是相同的(如上所示),但与问题中原始伪代码计算的上四面体相比,在数学和图形上更容易可视化。又是一些 Python:

N = 5
nmax = N * (N+1) * (N+2) // 6
x= np.empty(nmax, dtype=object)
x2 = np.zeros((N,N,N), dtype=np.int)

idx = 1
for i in range(0,N):
    for j in range(0,i+1):
        for k in range(0,j+1):
            x2[i,j,k] = idx
            x[idx-1] = (i,j,k)
            idx +=1
print(x)
print(x2)

产生

[(0, 0, 0) (1, 0, 0) (1, 1, 0) (1, 1, 1) (2, 0, 0) (2, 1, 0) (2, 1, 1)
 (2, 2, 0) (2, 2, 1) (2, 2, 2) (3, 0, 0) (3, 1, 0) (3, 1, 1) (3, 2, 0)
 (3, 2, 1) (3, 2, 2) (3, 3, 0) (3, 3, 1) (3, 3, 2) (3, 3, 3) (4, 0, 0)
 (4, 1, 0) (4, 1, 1) (4, 2, 0) (4, 2, 1) (4, 2, 2) (4, 3, 0) (4, 3, 1)
 (4, 3, 2) (4, 3, 3) (4, 4, 0) (4, 4, 1) (4, 4, 2) (4, 4, 3) (4, 4, 4)]

[[[ 1  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]]

 [[ 2  0  0  0  0]
  [ 3  4  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]]

 [[ 5  0  0  0  0]
  [ 6  7  0  0  0]
  [ 8  9 10  0  0]
  [ 0  0  0  0  0]
  [ 0  0  0  0  0]]

 [[11  0  0  0  0]
  [12 13  0  0  0]
  [14 15 16  0  0]
  [17 18 19 20  0]
  [ 0  0  0  0  0]]

 [[21  0  0  0  0]
  [22 23  0  0  0]
  [24 25 26  0  0]
  [27 28 29 30  0]
  [31 32 33 34 35]]]

你可以看到它是原始代码的一个变换,四面体的每个 "layer" 都是由一个大小递增的下三角矩阵构建的,而不是一个大小依次变小的上三角矩阵。

当您查看由此排列生成的四面体时,很明显每个下三角切片都从索引线性数组中的一个四面体数开始,下三角矩阵中的每一行都从 triangular number 相对于矩阵开始的偏移量。因此,索引方案是:

idx(i,j,k) = (i*(i+1)*(i+2)/6) + (j*(j+1)/2) + k

当数据排列时,第 k 个维度在内存中变化最快,第 i 个维度最慢。

现在进入实际问题。要从给定的 idx 值计算 (i,j,k) 需要计算 i 的整数立方根和 j 的整数平方根,这不是特别容易或高效,我不认为它会提供任何优势超过你现在所拥有的。但是,如果您的实现具有有限且已知的先验维度,您可以使用预先计算的四面体和三角形数并执行查找以取代计算根的需要。

玩具示例:

#include <cstdio>

__constant__ unsigned int tetdata[100] = 
{ 0, 1, 4, 10, 20, 35, 56, 84, 120, 165, 220, 286, 364, 455, 560, 680, 816, 969, 1140,
 1330, 1540, 1771, 2024, 2300, 2600, 2925, 3276, 3654, 4060, 4495, 4960, 5456, 5984,
 6545, 7140, 7770, 8436, 9139, 9880, 10660, 11480, 12341, 13244, 14190, 15180, 16215,
 17296, 18424, 19600, 20825, 22100, 23426, 24804, 26235, 27720, 29260, 30856, 32509,
 34220, 35990, 37820, 39711, 41664, 43680, 45760, 47905, 50116, 52394, 54740, 57155,
 59640, 62196, 64824, 67525, 70300, 73150, 76076, 79079, 82160, 85320, 88560, 91881,
 95284, 98770, 102340, 105995, 109736, 113564, 117480, 121485, 125580, 129766, 134044,
 138415, 142880, 147440, 152096, 156849, 161700, 166650 };

__constant__ unsigned int tridata[100] = 
 { 0, 1, 3, 6, 10, 15, 21, 28, 36, 45, 55, 66, 78, 91, 105, 120, 
 136, 153, 171, 190, 210, 231, 253, 276, 300, 325, 351, 378, 406, 
 435, 465, 496, 528, 561, 595, 630, 666, 703, 741, 780, 820, 861, 
 903, 946, 990, 1035, 1081, 1128, 1176, 1225, 1275, 1326, 1378, 1431,
 1485, 1540, 1596, 1653, 1711, 1770, 1830, 1891, 1953, 2016, 2080, 2145, 
 2211, 2278, 2346, 2415, 2485, 2556, 2628, 2701, 2775, 2850, 2926, 3003,
 3081, 3160, 3240, 3321, 3403, 3486, 3570, 3655, 3741, 3828, 3916, 4005,
 4095, 4186, 4278, 4371, 4465, 4560, 4656, 4753, 4851, 4950 };

__device__ unsigned int lookup(unsigned int&x, unsigned int n, const unsigned int* data)
{
    int i=0;
    while (n >= data[i]) i++;
    x = data[i-1];
    return i-1;
}

__device__ unsigned int tetnumber(unsigned int& x, unsigned int n) { return lookup(x, n, tetdata); }
__device__ unsigned int trinumber(unsigned int& x, unsigned int n) { return lookup(x, n, tridata); }

__global__ void kernel()
{
    unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;

    unsigned int x; 
    unsigned int k = idx;    
    unsigned int i = tetnumber(x, k); k -= x;
    unsigned int j = trinumber(x, k); k -= x;

    printf("idx = %d, i=%d j=%d k=%d\n", idx, i, j, k);
}

int main(void)
{
    cudaSetDevice(0);
    kernel<<<1,35>>>();
    cudaDeviceSynchronize();
    cudaDeviceReset();
    return 0;
}

与 python 做同样的事情(注意乱序打印输出):

$ nvcc -o tetrahedral tetrahedral.cu 
avidday@marteil2:~/SO$ cuda-memcheck ./tetrahedral 
========= CUDA-MEMCHECK
idx = 32, i=4 j=4 k=2
idx = 33, i=4 j=4 k=3
idx = 34, i=4 j=4 k=4
idx = 0, i=0 j=0 k=0
idx = 1, i=1 j=0 k=0
idx = 2, i=1 j=1 k=0
idx = 3, i=1 j=1 k=1
idx = 4, i=2 j=0 k=0
idx = 5, i=2 j=1 k=0
idx = 6, i=2 j=1 k=1
idx = 7, i=2 j=2 k=0
idx = 8, i=2 j=2 k=1
idx = 9, i=2 j=2 k=2
idx = 10, i=3 j=0 k=0
idx = 11, i=3 j=1 k=0
idx = 12, i=3 j=1 k=1
idx = 13, i=3 j=2 k=0
idx = 14, i=3 j=2 k=1
idx = 15, i=3 j=2 k=2
idx = 16, i=3 j=3 k=0
idx = 17, i=3 j=3 k=1
idx = 18, i=3 j=3 k=2
idx = 19, i=3 j=3 k=3
idx = 20, i=4 j=0 k=0
idx = 21, i=4 j=1 k=0
idx = 22, i=4 j=1 k=1
idx = 23, i=4 j=2 k=0
idx = 24, i=4 j=2 k=1
idx = 25, i=4 j=2 k=2
idx = 26, i=4 j=3 k=0
idx = 27, i=4 j=3 k=1
idx = 28, i=4 j=3 k=2
idx = 29, i=4 j=3 k=3
idx = 30, i=4 j=4 k=0
idx = 31, i=4 j=4 k=1
========= ERROR SUMMARY: 0 errors

显然,查找功能仅用于演示目的。在大尺寸情况下,二进制数组或基于散列的查找会快得多。但这至少表明,似乎可以按照您的设想去做,即使解决的问题和方法与您可能想到的略有不同。

请注意,对于此答案中的任何内容,我都没有正式的数学证明,也没有声称此处的任何代码或命题都是正确的。买家当心。


经过深思熟虑,通过相当高效的混合 search/calculation 例程扩展此方法是微不足道的:

#include <iostream>
#include <vector>
#include <cstdio>

typedef unsigned int uint;

__device__ __host__ ulong tetnum(uint n) { ulong n1(n); return n1 * (n1 + 1ull) * (n1 + 2ull) / 6ull; }
__device__ __host__ ulong trinum(uint n) { ulong n1(n); return n1 * (n1 + 1ull) / 2ull; }

typedef ulong (*Functor)(uint);
template<Functor F>
__device__ __host__ uint bounded(ulong& y, ulong x, uint n1=0, ulong y1=0)
{
    uint n = n1;
    y = y1;

    while (x >= y1) {
        y = y1; 
        n = n1++;
        y1 = F(n1);
    }
    return n;
}

__constant__ uint idxvals[19] = {
 0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 
 32768, 65536, 131072 };

__constant__ ulong tetvals[19] = { 
 0, 1, 4, 20, 120, 816, 5984, 45760, 357760, 2829056, 22500864, 179481600, 1433753600,
 11461636096, 91659526144, 733141975040, 5864598896640, 46914643623936, 375308558925824 };

__constant__ ulong trivals[19] = {
 0, 1, 3, 10, 36, 136, 528, 2080, 8256, 32896, 131328, 524800, 2098176, 8390656, 33558528,
 134225920, 536887296, 2147516416, 8590000128 };

__device__ __host__ uint lookup(ulong& x, uint n, const uint* abscissa, const ulong* data)
{
    uint i=0;
    while (n >= data[i]) i++;
    x = data[i-1];
    return abscissa[i-1];
}

 __device__ uint tetnumber(ulong& x, uint n)
{
    ulong x0;
    uint n0 = lookup(x0, n, idxvals, tetvals);
    return bounded<tetnum>(x, n, n0, x0);
}

 __device__ uint trinumber(ulong& x, uint n)
{
    ulong x0;
    uint n0 = lookup(x0, n, idxvals, trivals);
    return bounded<trinum>(x, n, n0, x0);
}


__global__ void kernel(uint3 *results, ulong Nmax)
{
    ulong idx = threadIdx.x + blockIdx.x * blockDim.x;
    ulong gridStride = blockDim.x * gridDim.x;

    for(; idx < Nmax; idx += gridStride) {
        ulong x, k1 = idx; 

        uint3 tuple;
        tuple.x = tetnumber(x, k1); k1 -= x;
        tuple.y = trinumber(x, k1); k1 -= x;
        tuple.z = (uint)k1;

        results[idx] = tuple;
    }
}

int main(void)
{
    cudaSetDevice(0);

    uint N = 500;
    ulong Nmax = tetnum(N);

    uint3* results_d; cudaMalloc(&results_d, Nmax * sizeof(uint3));

    int gridsize, blocksize;
    cudaOccupancyMaxPotentialBlockSize(&gridsize, &blocksize, kernel);
    kernel<<<gridsize, blocksize>>>(results_d, Nmax);
    cudaDeviceSynchronize();

    std::vector<uint3> results(Nmax);
    cudaMemcpy(&results[0], results_d, Nmax * sizeof(uint3), cudaMemcpyDeviceToHost);
    cudaDeviceReset();

    // Only uncomment this if you want to see 22 million lines of output
    //for(auto const& idx : results) {
    //    std::cout << idx.x << "  " << idx.y << "  " << idx.z << std::endl;
    //}

    return 0;
}

它执行此操作(请注意,如果您取消对最后一个循环的注释,它将发出 2100 万行输出):

$ module load use.own cuda9.2
$ nvcc -std=c++11 -arch=sm_52 -o tetrahedral tetrahedral.cu 
$ nvprof ./tetrahedral
==20673== NVPROF is profiling process 20673, command: ./tetrahedral
==20673== Profiling application: ./tetrahedral
==20673== Profiling result:
            Type  Time(%)      Time     Calls       Avg       Min       Max  Name
 GPU activities:   78.85%  154.23ms         1  154.23ms  154.23ms  154.23ms  kernel(uint3*, unsigned long)
                   21.15%  41.361ms         1  41.361ms  41.361ms  41.361ms  [CUDA memcpy DtoH]
      API calls:   41.73%  154.24ms         1  154.24ms  154.24ms  154.24ms  cudaDeviceSynchronize
                   30.90%  114.22ms         1  114.22ms  114.22ms  114.22ms  cudaMalloc
                   15.94%  58.903ms         1  58.903ms  58.903ms  58.903ms  cudaDeviceReset
                   11.26%  41.604ms         1  41.604ms  41.604ms  41.604ms  cudaMemcpy
                    0.11%  412.75us        96  4.2990us     275ns  177.45us  cuDeviceGetAttribute
                    0.04%  129.46us         1  129.46us  129.46us  129.46us  cuDeviceTotalMem
                    0.02%  55.616us         1  55.616us  55.616us  55.616us  cuDeviceGetName
                    0.01%  32.919us         1  32.919us  32.919us  32.919us  cudaLaunchKernel
                    0.00%  10.211us         1  10.211us  10.211us  10.211us  cudaSetDevice
                    0.00%  5.7640us         1  5.7640us  5.7640us  5.7640us  cudaFuncGetAttributes
                    0.00%  4.6690us         1  4.6690us  4.6690us  4.6690us  cuDeviceGetPCIBusId
                    0.00%  2.8580us         4     714ns     393ns  1.3680us  cudaDeviceGetAttribute
                    0.00%  2.8050us         3     935ns     371ns  2.0030us  cuDeviceGetCount
                    0.00%  2.2780us         1  2.2780us  2.2780us  2.2780us  cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags
                    0.00%  1.6720us         1  1.6720us  1.6720us  1.6720us  cudaGetDevice
                    0.00%  1.5450us         2     772ns     322ns  1.2230us  cuDeviceGet

该代码在我的 GTX970 上计算并存储了 500 x 500 x 500 搜索 space(大约 2100 万个值)的唯一 (i,j,k) 对,用时 150 毫秒。也许这对你有用。

this wikipedia page ("Finding the k-combination for a given number") 上给出了一种可能的方法,用于将线性索引转换为唯一 C(n,3) 组合的封闭形式解决方案。

但是会涉及计算平方根和立方根,所以是"non-trivial"。我什至提到它的理由有两个:

  1. 如果每个线程要节省的工作量很大,那么此方法提出的额外负担可能会被它抵消。但是,对于给出的示例,每个线程节省的工作量只是一些简单的 if 测试。

  2. 处理器的趋势是计算成本下降得比其他处理器更快。内存访问成本。由于这种方法不涉及内存访问,如果未来的处理器趋势继续这种趋势,这种方法可能会变得更加 palatable.

这种方法的另一个特点是没有迭代穷举 table 搜索。然而,正如另一个答案中所指出的,对于那里给出的规定,目前几乎肯定比这种方法更可取。

如前面提到的维基页面所示,一般方法是:

  1. 找出小于当前索引(N)的最大C(n,3)数。与此 C(n,3) 数关联的 n 值成为我们第一个 "choice" 索引 n1.

  2. 的序数值
  3. 从当前索引中减去C(n,3)数。对余数和 C(n,2) 重复该过程。与适合余数的最大 C(n,2) 数相关联的 n 值成为我们的第二个 "choice" 索引 n2.

  4. 从第 2 步中找到余数,然后确定我们最终的 C(n,1) 选择 (C(n,1) = n = n3)。

为了得出第 1 步的封闭式解决方案,我们必须:

  • 找出与三次方程相关的关系 N 和 C(N,3)

  • 用三次多项式的解来识别N(在浮动 观点)。

  • T运行计算值N,得到我们的"largest" N.

  • 围绕这一点执行整数搜索,寻找正确的解决方案,解决浮点问题

可以对步骤 2(二次)和步骤 3(线性)重复类似的过程。

我不打算特别详细地涵盖所有数学,但是可以在网上轻松找到封闭形式的三次多项式方程的解(例如 here)和推导步骤 1 的控制三次方程很简单。我们简单地使用问题中已经给出的选项总数的公式,再加上特定的线程索引:

  n(n-1)(n-2)/6 = N  -> n(n-1)(n-2)/6 - N = 0

重新排列:

  (n^3)/6 - (n^2)/2 + n/3 - N = 0

据此我们可以获得 a、b、c、d 系数以输入到我们的三次求解方法中。

  a = 1/6, b = -1/2, c = 1/3, d = -N

(注意这里的 N 实际上是我们全局唯一的一维线程索引。我们正在求解 n,这给了我们第一个 "choice" 索引。)

研究三次解的公式,我们注意到线程之间唯一不同的项目是 d 系数。这允许在 运行 时间减少一些算术。

接下来是一个有效的例子。它没有经过全面测试,因为我的目的是确定解决方法,而不是经过全面测试的解决方案:

$ cat t1485.cu
#include <stdio.h>
#include <math.h>

typedef float ct;
const int STEP_DOWN = 2;
// only float or double template types allowed
template <typename ft>
struct CN3{
  __host__ __device__
  int3 operator()(size_t N){
    int3 n;
    if (N == 0) {n.x = 2; n.y = 1; n.z = 0; return n;}
    if (N == 1) {n.x = 3; n.y = 1; n.z = 0; return n;}
    if (N == 2) {n.x = 3; n.y = 2; n.z = 0; return n;}
    if (N == 3) {n.x = 3; n.y = 2; n.z = 1; return n;}
    if (N == 4) {n.x = 4; n.y = 1; n.z = 0; return n;}
    ft x, x1;
    // identify n.x from cubic
    // compiler computed
    const ft a = 1.0/6;
    const ft b = -1.0/2;
    const ft c = 1.0/3;
    const ft p1 = (-1.0)*(b*b*b)/(27.0*a*a*a) + b*c/(6.0*a*a);
    const ft p2 = c/(3.0*a) - (b*b)/(9.0*a*a);
    const ft p3 = p2*p2*p2;
    const ft p4 = b/(3.0*a);
    // run-time computed
    //const ft d = -N;
    const ft q0 = N/(2.0*a);  // really should adjust constant for float vs. double
    const ft q1 = p1 + q0;
    const ft q2 = q1*q1;
    if (sizeof(ft)==4){
      x1 = sqrtf(q2+p3);
      x =  cbrtf(q1+x1) + cbrtf(q1-x1) - p4;
      n.x = truncf(x);}
    else {
      x1 = sqrt(q2+p3);
      x =  cbrt(q1+x1) + cbrt(q1-x1) - p4;
      n.x = trunc(x);}
    /// fix floating-point errors
    size_t tn = n.x - STEP_DOWN;
    while ((tn)*(tn-1)*(tn-2)/6 <= N) tn++;
    n.x = tn-1;
    // identify n.y from quadratic
    // compiler computed
    const ft qa = 1.0/2;
    //const ft qb = -qa;
    const ft p5 = 1.0/4;
    const ft p6 = 2.0;
    // run-time computed
    N = N - (((size_t)n.x)*(n.x-1)*(n.x-2))/6;
    if (sizeof(ft)==4){
      x = qa + sqrtf(p5+p6*N);
      n.y = truncf(x);}
    else {
      x = qa + sqrt(p5+p6*N);
      n.y = trunc(x);}
    /// fix floating-point errors
    if ((n.y - STEP_DOWN) <= 0) tn = 0;
    else tn = n.y - STEP_DOWN;
    while ((((tn)*(tn-1))>>1) <= N) tn++;
    n.y = tn-1;
    // identify n3
    n.z = N - ((((size_t)n.y)*(n.y-1))>>1);
    return n;
    }
};

template <typename T>
__global__ void test(T f, size_t maxn, int3 *res){

  size_t idx = threadIdx.x+((size_t)blockDim.x)*blockIdx.x;
  if (idx < maxn)
    res[idx] = f(idx);
}

int3 get_next_C3(int3 prev){
  int3 res = prev;
  res.z++;
  if (res.z >= res.y){
    res.y++;  res.z = 0;
    if (res.y >= res.x){res.x++; res.y = 1; res.z = 0;}}
  return res;
}

int main(int argc, char* argv[]){
  size_t n = 1000000000;
  if (argc > 1) n *= atoi(argv[1]);
  const int nTPB = 256;
  int3 *d_res;
  cudaMalloc(&d_res, n*sizeof(int3));
  test<<<(n+nTPB-1)/nTPB,nTPB>>>(CN3<ct>(), n, d_res);
  int3 *h_gpu = new int3[n];
  int3 temp;
  temp.x = 2; temp.y = 1; temp.z = 0;
  cudaMemcpy(h_gpu, d_res, n*sizeof(int3), cudaMemcpyDeviceToHost);
  for (int i = 0; i < n; i++){
    if ((temp.x != h_gpu[i].x) || (temp.y != h_gpu[i].y) || (temp.z != h_gpu[i].z))
      {printf("mismatch at index %d: cpu: %d,%d,%d gpu: %d,%d,%d\n", i, temp.x,temp.y,temp.z, h_gpu[i].x, h_gpu[i].y, h_gpu[i].z); return 0;}
    temp = get_next_C3(temp);}
}
$ nvcc -arch=sm_70 -o t1485 t1485.cu
$ cuda-memcheck ./t1485 2
========= CUDA-MEMCHECK
========= ERROR SUMMARY: 0 errors
[user2@dc10 misc]$ nvprof ./t1485
==6128== NVPROF is profiling process 6128, command: ./t1485
==6128== Profiling application: ./t1485
==6128== Profiling result:
            Type  Time(%)      Time     Calls       Avg       Min       Max  Name
 GPU activities:   99.35%  4.81251s         1  4.81251s  4.81251s  4.81251s  [CUDA memcpy DtoH]
                    0.65%  31.507ms         1  31.507ms  31.507ms  31.507ms  void test<CN3<float>>(float, int, int3*)
      API calls:   93.70%  4.84430s         1  4.84430s  4.84430s  4.84430s  cudaMemcpy
                    6.09%  314.89ms         1  314.89ms  314.89ms  314.89ms  cudaMalloc
                    0.11%  5.4296ms         4  1.3574ms  691.18us  3.3429ms  cuDeviceTotalMem
                    0.10%  4.9644ms       388  12.794us     317ns  535.35us  cuDeviceGetAttribute
                    0.01%  454.66us         4  113.66us  103.24us  134.26us  cuDeviceGetName
                    0.00%  65.032us         1  65.032us  65.032us  65.032us  cudaLaunchKernel
                    0.00%  24.906us         4  6.2260us  3.2890us  10.160us  cuDeviceGetPCIBusId
                    0.00%  8.2490us         8  1.0310us     533ns  1.5980us  cuDeviceGet
                    0.00%  5.9930us         3  1.9970us     381ns  3.8870us  cuDeviceGetCount
                    0.00%  2.8160us         4     704ns     600ns     880ns  cuDeviceGetUuid
$

备注:

  • 如上所述,我已经测试了前 20 亿个结果的准确性
  • 上面的实现解释了浮点三次方程和​​二次方程的解引入错误的事实。这些错误是 "fixed" 通过围绕浮点计算给定的起点创建局部整数搜索来产生正确答案。
  • 如前所述,在我的 Tesla V100 上,内核在 运行 秒内耗时约 30 毫秒,获得 10 亿个结果 (10^9)。如果该方法可以正确地扩展到 10^15 个结果,我没有理由认为它不会花费至少 0.03*10^6 秒,或者超过 8 小时(!)
  • 我没有 运行 测试,但我怀疑在简单生成完整域 (10^15) 然后丢弃 ~5 的问题中提出的简单案例的快速基准测试/6 的 space 不适用,会更快。

出于好奇,我创建了一个替代测试用例,用于测试每 32 个值中的 31 个,跨越更大的 space。

这是代码和测试:

$ cat t1485.cu
#include <stdio.h>
#include <math.h>

typedef float ct;
const int nTPB = 1024;
const int STEP_DOWN = 2;
// only float or double template types allowed
template <typename ft>
struct CN3{
  __host__ __device__
  int3 operator()(size_t N){
    int3 n;
    if (N == 0) {n.x = 2; n.y = 1; n.z = 0; return n;}
    if (N == 1) {n.x = 3; n.y = 1; n.z = 0; return n;}
    if (N == 2) {n.x = 3; n.y = 2; n.z = 0; return n;}
    if (N == 3) {n.x = 3; n.y = 2; n.z = 1; return n;}
    if (N == 4) {n.x = 4; n.y = 1; n.z = 0; return n;}
    ft x, x1;
    // identify n.x from cubic
    // compiler computed
    const ft a = 1.0/6;
    const ft b = -1.0/2;
    const ft c = 1.0/3;
    const ft p1 = (-1.0)*(b*b*b)/(27.0*a*a*a) + b*c/(6.0*a*a);
    const ft p2 = c/(3.0*a) - (b*b)/(9.0*a*a);
    const ft p3 = p2*p2*p2;
    const ft p4 = b/(3.0*a);
    // run-time computed
    //const ft d = -N;
    const ft q0 = N/(2.0*a);  // really should adjust constant for float vs. double
    const ft q1 = p1 + q0;
    const ft q2 = q1*q1;
    if (sizeof(ft)==4){
      x1 = sqrtf(q2+p3);
      x =  cbrtf(q1+x1) + cbrtf(q1-x1) - p4;
      n.x = truncf(x);}
    else {
      x1 = sqrt(q2+p3);
      x =  cbrt(q1+x1) + cbrt(q1-x1) - p4;
      n.x = trunc(x);}
    /// fix floating-point errors
    size_t tn = n.x - STEP_DOWN;
    while ((tn)*(tn-1)*(tn-2)/6 <= N) tn++;
    n.x = tn-1;
    // identify n.y from quadratic
    // compiler computed
    const ft qa = 1.0/2;
    //const ft qb = -qa;
    const ft p5 = 1.0/4;
    const ft p6 = 2.0;
    // run-time computed
    N = N - (((size_t)n.x)*(n.x-1)*(n.x-2))/6;
    if (sizeof(ft)==4){
      x = qa + sqrtf(p5+p6*N);
      n.y = truncf(x);}
    else {
      x = qa + sqrt(p5+p6*N);
      n.y = trunc(x);}
    /// fix floating-point errors
    if ((n.y - STEP_DOWN) <= 0) tn = 0;
    else tn = n.y - STEP_DOWN;
    while ((((tn)*(tn-1))>>1) <= N) tn++;
    n.y = tn-1;
    // identify n3
    n.z = N - ((((size_t)n.y)*(n.y-1))>>1);
    return n;
    }
};

__host__ __device__
int3 get_next_C3(int3 prev){
  int3 res = prev;
  res.z++;
  if (res.z >= res.y){
    res.y++;  res.z = 0;
    if (res.y >= res.x){res.x++; res.y = 1; res.z = 0;}}
  return res;
}

template <typename T>
__global__ void test(T f){

  size_t idx = threadIdx.x+((size_t)blockDim.x)*blockIdx.x;
  size_t idy = threadIdx.y+((size_t)blockDim.y)*blockIdx.y;
  size_t id = idx + idy*gridDim.x*blockDim.x;
  int3 temp = f(id);
  int3 temp2;
  temp2.x = __shfl_up_sync(0xFFFFFFFF, temp.x, 1);
  temp2.y = __shfl_up_sync(0xFFFFFFFF, temp.y, 1);
  temp2.z = __shfl_up_sync(0xFFFFFFFF, temp.z, 1);
  temp2 = get_next_C3(temp2);
  if ((threadIdx.x & 31) != 0)
    if ((temp.x != temp2.x) || (temp.y != temp2.y) || (temp.z != temp2.z)) printf("%lu,%d,%d,%d,%d,%d,%d\n", id, temp.x, temp.y, temp.z, temp2.x, temp2.y, temp2.z);
}

int main(int argc, char* argv[]){
  const size_t nbx = 200000000ULL;
  const int nby = 100;
  dim3 block(nbx, nby, 1);
  test<<<block,nTPB>>>(CN3<ct>());
  cudaDeviceSynchronize();
  cudaError_t e = cudaGetLastError();
  if (e != cudaSuccess) {printf("CUDA error %s\n", e); return 0;}
  printf("tested space of size: %lu\n",  nbx*nby*nTPB);
}
$ nvcc -arch=sm_70 -o t1485 t1485.cu
$ time ./t1485
tested space of size: 20480000000000

real    25m18.133s
user    18m4.804s
sys     7m12.782s

在这里我们看到 Tesla V100 花了大约 30 分钟来准确测试 space 的 20480000000000 个结果(大约 2 * 10^13)。