cuda - 在CUDA中的另一个向量中找到最近的非零元素

Question

有一个M x N的矩阵A和B。(矩阵的实际大小是512 x 4096)
在A的每一行中，待处理的点都设置为1。而B的每一行都是通过特定操作得到的值。
基于每一行，我将做一个操作来获取最接近 A 中 1 的点的 B 的值。
例子如下图，我用MATLAB写的代码也记下来了。
这是我的想法:

用推力选取 A 的非零元素索引。并且对于每个元素，通过for循环从B的相应行中获取最接近的值。(如果A中有多个非零元素，预计会很慢。)

我想好好利用GPU的力量来做这个运算，你有什么更高效的想法吗？

[idxY,idxX] = find(A == 1);
for Point = 1:length(idxY)
    pointBuf = find(B(:,idxY(Point)) == 1); // find non-zero elements in Row of B
    if ~isempty(pointBuf)                   // there are non-zero elements in Row of B
        [MinValue, MinIndex] = min(abs(pointBuf - idxY(Point)));
        C(idxY(Point),idxX(Point)) = B(pointBuf(MinIndex(1)),RangeInd(Point)); // Get closest point in B
    else
        C(DopInd(Point),RangeInd(Point)) = 0; // if there is no non-zero elements in Row of B, just set to 0
    end
end

Answer 1

就像引用一个解决方案，它左右移动 4095。它与冒泡排序的变体有相似之处，同时向上和向下冒泡。

优点是它不依赖于B中非空元素的位置，可以很容易地在线程之间并行化。

但是转换为 2 条 SASS 指令的内部循环仍然太慢(调用频率太高):该程序在我的笔记本上需要 26 毫秒。

它会在输入矩阵的最佳和绝对最差情况下这样做。

它的部分和方法可能可以重复使用，因为它展示了一些 CUDA 编程方法。所以或多或少供引用，到底不是最终(够快)的解决方案:

__global__ void calcmatrix(bool* A, double* B, double* C)
{
// calculate row number
    int row = blockDim.x * gridDim.y + threadIdx.y;
    if (row >= 512)
        return;

// store index of valid double from B, this is moved up and down
// those indices are for the current thread. Each thread is responsible for 128 consecutive columns.
    int indices[128];

// prefill the indices with their own number (as if every double from B is valid)
#pragma unroll
    for (int i = 0; i < 128; i++)
        indices[i] = threadIdx.x * 128 + i;

// Store zero flags (4 * 32 bits) for our 128 elements
    unsigned int myzeroflags[4];

// For efficiently loading data from memory, we distribute the data in another way: thread 0 gets columns 0, 32, 64, 96, ...; thread 1 gets columns 1, 33, 65, 97, ...; thread 2 gets columns 2, 34, 66, 98, ...; and so on

#pragma unroll
    for (int i = 0; i < 128; i++) {
// load value from B
        double in = B[row * 4096 + i * 32 + threadIdx.x];
// compare to zero (!in) and combine all bool results from the 32 threads (__ballot_sync))
        unsigned int zeroflag = __ballot_sync(0xFFFFFFFF, !in);
// store the ones, which belong to us
        if (threadIdx.x == i / 4)
            myzeroflags[i & 3] = zeroflag;
    }

// go through our zero flags and set those indices to -1 (there is already a valid index "0", so we use a negative number to signify invalid)
#pragma unroll
    for (int i = 0; i < 4; i++)
#pragma unroll
        for (int j = 0; j < 32; j++)
            if (myzeroflags[i] & (1 << j))
                indices[i * 32 + j] = -1;

// main loop, do 4095 times
#pragma unroll 1
    for (int i = 0; i < 4095; i++) {
// move all elements to the left (if the index there is invalid)
// send index over thread boundaries
        int fromright = __shfl_down_sync(0xFFFFFFFF, indices[0], 1, 32);
#pragma unroll
// if left index is -1, set it to one index to the right
        for (int j = 0; j < 127; j++)
            if (indices[j] == -1)
                indices[j] = indices[j + 1];
// move over thread boundaries (except for the rightmost thread)
        if (threadIdx.x != 31 && indices[127] == -1)
            indices[127] = fromright;

// move to the right in the same way as to the left
        int fromleft = __shfl_up_sync(0xFFFFFFFF, indices[127], 1, 32);
#pragma unroll
        for (int j = 127; j > 0; j--)
            if (indices[j] == -1)
                indices[j] = indices[j - 1];
        if (threadIdx.x != 0 && indices[0] == -1)
            indices[0] = fromleft;
    }

// for the other distribution of elements for memory accesses, we have to redistribute the indices to the correct threads
// To not have bank conflicts, we define the shared memory array with 33 instead of 32 elements in the last dimension, but use only 32. With this method we can put threadIdx.x into the last and previous to last dimension without bank conflicts
    __shared__ short2 distribidx[8][32][33];

    int indices2[128];

// Redistribute first half; the index can go from 0..4095 (and also theoreticially -1, if there was no non-null element in this row). This fits into a short, convert for faster transfer
#pragma unroll
    for (int i = 0; i < 32; i++)
        distribidx[threadIdx.y][threadIdx.x][i] = { static_cast<short>(indices[i]), static_cast<short>(indices[i + 32]) };
    __syncwarp();
#pragma unroll
    for (int i = 0; i < 32; i++) {
        short2 idxback = distribidx[threadIdx.y][i][threadIdx.x];
        indices2[4 * i + 0] = idxback.x;
        indices2[4 * i + 1] = idxback.y;
    }

    __syncwarp();
// Redistribute second half

#pragma unroll
    for (int i = 0; i < 32; i++)
        distribidx[threadIdx.y][threadIdx.x][i] = { static_cast<short>(indices[i + 64]), static_cast<short>(indices[i + 96]) };
    __syncwarp();
#pragma unroll
    for (int i = 0; i < 32; i++) {
        short2 idxback = distribidx[threadIdx.y][i][threadIdx.x];
        indices2[4 * i + 2] = idxback.x;
        indices2[4 * i + 3] = idxback.y;
    }

// Do final calculation
#pragma unroll
    for (int i = 0; i < 128; i++) {
// Default value is zero
        double result = 0;
// Read only, if A is true and indices2 is valid
        if (A[row * 4096 + i * 32 + threadIdx.x] && indices2[i] != -1)
// Read B with calculated index (this read is not optimized/coalesced, because the indices can be wild, but hopefully was or can be cached)
            result = B[row * 4096 + indices2[i]];
// Store result in C
        C[row * 4096 + i * 32 + threadIdx.x] = result;
    }
}

int main()
{
    bool* A;
    double* B;
    double* C;
    cudaMalloc(&A, 2 * 512 * 4096);
    cudaMalloc(&B, 8 * 512 * 4096);
    cudaMalloc(&C, 8 * 512 * 4096);
// called in this fashion
    calcmatrix<<<(512 + 7) / 8, dim3(32, 8)>>>(A, B, C);

    return 0;
}