cuda gpu CTAMergesort

注:下面所提到的代码实现，均在moderngpu2.0中被作者所重写。

gpu上-并行归并排序,modern gpu上CTAMergesort源码实现:

template<int NT, int VT, bool Stable, bool HasValues, typename KeyType, 
    typename ValType, typename Comp>
MGPU_DEVICE void CTAMergesort(KeyType threadKeys[VT], ValType threadValues[VT],
    KeyType* keys_shared, ValType* values_shared, int count, int tid, 
    Comp comp) {


    // Stable sort the keys in the thread.
    if(VT * tid < count) {
        if(Stable)
            OddEvenTransposeSort<VT>(threadKeys, threadValues, comp);
        else
            OddEvenMergesort<VT>(threadKeys, threadValues, comp);
    }   


    // Store the locally sorted keys into shared memory.
    DeviceThreadToShared<VT>(threadKeys, tid, keys_shared);


    // Recursively merge lists until the entire CTA is sorted.
    CTABlocksortLoop<NT, VT, HasValues>(threadValues, keys_shared, 
        values_shared, tid, count, comp);
}

该方法的实现逻辑：

1.将要排序的数据列表,先将每个线程自己要排序的数据存储到线程本地变量上(寄存器上)

2.每个线程先使用奇偶排序排序本地变量数据。

3.将本地排序好的变量数据，统一存放会原来共享内存的对应位置上，供下面CTABlocksortLoop进行最终数据的排序(共享内存的数据是SM内的所有线程共享的，而SM是以block进行统一调度的，即该排序是block级别的排序)

4.CTABlocksortLoop协调每个线程自己的排序工作。

CTABlocksortLoop源码:

template<int NT, int VT, bool HasValues, typename KeyType, typename ValType,
    typename Comp>
MGPU_DEVICE void CTABlocksortLoop(ValType threadValues[VT],
    KeyType* keys_shared, ValType* values_shared, int tid, int count,
    Comp comp) {


    #pragma unroll
    for(int coop = 2; coop <= NT; coop *= 2) {
        int indices[VT];
        KeyType keys[VT];
        CTABlocksortPass<NT, VT>(keys_shared, tid, count, coop, keys,
            indices, comp);


        if(HasValues) {
            // Exchange the values through shared memory.
            DeviceThreadToShared<VT>(threadValues, tid, values_shared);
            DeviceGather<NT, VT>(NT * VT, values_shared, indices, tid,
                threadValues);
        }


        // Store results in shared memory in sorted order.
        DeviceThreadToShared<VT>(keys, tid, keys_shared);
    }
}

coop为线程合作数,且每个线程各自排序的数据为VT数据量。

每一次迭代结束，均表明每一块为coop * VT的数据均是排序好的数据。

CTABlocksortPass源码实现:

template<int NT, int VT, typename T, typename Comp>
MGPU_DEVICE void CTABlocksortPass(T* keys_shared, int tid, int count,
    int coop, T* keys, int* indices, Comp comp) {


    int list = ~(coop - 1) & tid;
    int diag = min(count, VT * ((coop - 1) & tid));
    int start = VT * list;
    int a0 = min(count, start);
    int b0 = min(count, start + VT * (coop / 2));
    int b1 = min(count, start + VT * coop);


    int p = MergePath<MgpuBoundsLower>(keys_shared + a0, b0 - a0,
        keys_shared + b0, b1 - b0, diag, comp);


    SerialMerge<VT, true>(keys_shared, a0 + p, b0, b0 + diag - p, b1, keys,
        indices, comp);
}

这段代码便是能够使多个线程能并行排序各自的数据，并且各自数据排序完后，相关联的线程下的数据便是连续有序的数据了。这听起来有点神奇。

这里有点注明下，假设现在有N段的已经排序好的数据，那么需要将数据合并成N/2端排序好的数据,即需要将相邻的两端排序好的数据进行合并工作，而每每两段的合并工作的

协作线程数便是coop。也就是不管协调线程数coop不管有多少，其所负责的都只是2相邻两端的排序好的数据,也就是说我们需要将这两端的合并排序的工作，平分给这coop个线程。

下面便说这段代码是如何做到的，其实这里用到了

Merge Path-Parallel Merging Made Simple,有兴趣的可以去搜搜看具体内容。

下面说说我对这个的理解。

首先，请先回忆下cpu上的两个已排序的数据的合并算法，时间耗时是O(n)即可。

1.A和B数组，同时下标i,下标j代表A和B数据的索引

2.A[i] <= B[j]时， 将A[i]的值放到C[index]中, i+1

3.B[j] < A[i]时， 将B[j]的值放到C[index]中, j+1

4.将剩余的A或者B数组元素直接拷贝到C末尾中

这里用一个矩阵进行表示：

从i = 0, j = 0(左上角)开始其代表的数字为3,17,当A[i]小则走向向下,当B[j]小则走向向右，最终到达右下角结束。

那条橙黄色的连续箭头便是合并的路径(merge path)。

而这个合并路径有一个很重要的属性-在这个路径上的随意一点均可将两段已排序的数据的合并工作分成独立的两个合并工作。

以上图为例,中间的这条45度的分割线,与Merge path相交的点,即A(17 ~ 35), B(3 ~ 45)和A(73 ~ 99)，B(64 ~ 82)这两部分的合并工作是独立，不会相互影响各自的结果的。

这有个通俗的原因，是因为MergePath上的点刚好可以将两个已排序的数组的数分成全是较小部分和全是较大部分，即对于前一部分的每一个数据都是小于后一部分的数据的，因此这两部分的数据排序不会影响最终结果的位置，因此可以并行排序。

以此发散，可以将这个图同时画上多条线

这3条线，便把这两个排序数组的合并工作分成4等份工作了,记得前面提到的coop写作线程数吗，这里的4等份便是对应于coop=4时，这4个协同线程是如何协同工作的。

那么有一个就是该如何确定这个交点呢,细心的应该可以发现，交点是在这条斜线所经过的块1和块0的交界处(有两个情况可能会觉得不太一样,当斜线均经过0或者经过1的时候，

其实这两种情况在编写代码的时候均可归类到查询块1和块0的交界处)

因此确定了交界点，并可以正确的将上述的两个数组的排序划分成相等的几个独立的排序工作了。

而实际代码编写中,是每个协调线程查询出自己的起始位置，和自己处理的第几块合并数据，便可以独立完成工作了，如下图

这里会有个diag的限制,其实diag0 diag1 diag2 diag3之间的差是每个线程的计算数据量(一般会相同,VT表示)

而第一个协调线程diag0,起始位置是(17,3),计算量VT

而第二个协调线程diag1,起始位置是(29, 22),计算量VT

而第三个协调线程diag2,起始位置是(73, 64),计算量VT

而第四个协调线程diag3,起始位置是(86, 82),计算量VT

最终A和B便合并完成排序。

注:该方法便是查询每个协调线程的起始位置

int p = MergePath<MgpuBoundsLower>(keys_shared + a0, b0 - a0,
        keys_shared + b0, b1 - b0, diag, comp);

注:该方法是协调线程合并计算量VT的排序数据

SerialMerge<VT, true>(keys_shared, a0 + p, b0, b0 + diag - p, b1, keys,
        indices, comp);