矢量化的基数排序与numpy - 它能击败np.sort吗？

Question

Numpy没有yet有基数排序，所以我想知道是否可以使用预先存在的numpy函数编写一个。到目前为止，我有以下内容，它确实有效，但比numpy的快速排序慢约10倍。

测试和基准：

a = np.random.randint(0, 1e8, 1e6)
assert(np.all(radix_sort(a) == np.sort(a))) 
%timeit np.sort(a)
%timeit radix_sort(a)

mask_b循环可以至少部分地进行矢量化，通过&的掩码进行广播，并使用cumsum和axis arg进行广播，但这最终会导致悲观情绪化，可能是由于内存占用增加。

如果有人能够看到一种方法来改进我所拥有的东西，我有兴趣听到，即使它仍然比np.sort慢......这更像是知识分子的情况好奇心和对numpy技巧的兴趣。

请注意，您can implement可以轻松快速计算排序，但这仅适用于小整数数据。

编辑1：将np.arange(n)移出循环会有所帮助，但这并不是非常令人兴奋。

编辑2： cumsum实际上是多余的（ooops！），但这个更简单的版本仅对性能有所帮助..

def radix_sort(a):
    bit_len = np.max(a).bit_length()
    n = len(a)
    cached_arange = arange(n)
    idx = np.empty(n, dtype=int) # fully overwritten each iteration
    for mask_b in xrange(bit_len):
        is_one = (a & 2**mask_b).astype(bool)
        n_ones = np.sum(is_one)      
        n_zeros = n-n_ones
        idx[~is_one] = cached_arange[:n_zeros]
        idx[is_one] = cached_arange[:n_ones] + n_zeros
        # next three lines just do: a[idx] = a, but correctly
        new_a = np.empty(n, dtype=a.dtype)
        new_a[idx] = a
        a = new_a
    return a

编辑3：而不是循环遍历单个位，如果您在多个步骤中构造idx，则可以一次循环两次或更多次。使用2位有点帮助，我没有尝试过更多：

idx[is_zero] = np.arange(n_zeros)
idx[is_one] = np.arange(n_ones)
idx[is_two] = np.arange(n_twos)
idx[is_three] = np.arange(n_threes)

编辑4和5： 4位似乎最适合输入I测试。此外，您可以完全摆脱idx步骤。现在只比np.sort（source available as gist）慢5倍，而不是10倍：

编辑6：这是上面的整理版本，但它也有点慢。 80％的时间花在repeat和extract上 - 如果只有一种方法可以广播extract :( ...

def radix_sort(a, batch_m_bits=3):
    bit_len = np.max(a).bit_length()
    batch_m = 2**batch_m_bits
    mask = 2**batch_m_bits - 1
    val_set = np.arange(batch_m, dtype=a.dtype)[:, nax] # nax = np.newaxis
    for _ in range((bit_len-1)//batch_m_bits + 1): # ceil-division
        a = np.extract((a & mask)[nax, :] == val_set,
                        np.repeat(a[nax, :], batch_m, axis=0))
        val_set <<= batch_m_bits
        mask <<= batch_m_bits
    return a

编辑7＆amp; 8：实际上，您可以使用as_strided中的numpy.lib.stride_tricks来广播摘录，但它似乎无助于提高性能：

最初这对我有意义，理由是extract将遍历整个数组batch_m次，因此CPU请求的缓存行总数将与之前相同（它只是在流程结束时它已经请求每个缓存行batch_m次）。然而the reality是extract不够聪明，无法迭代任意阶梯数组，并且必须在开始之前扩展数组，即无论如何重复都会完成。 事实上，看过extract的来源，我现在看到我们用这种方法做的最好的事情是：

a = a[np.flatnonzero((a & mask)[nax, :] == val_set) % len(a)]

略慢于extract。但是，如果len(a)是2的幂，我们可以用& (len(a) - 1)替换昂贵的mod操作，这最终会比extract版本快一点（现在约为4.9x {{ 1}} np.sort）。我想我们可以通过零填充使这两个长度的非幂次工作，然后在排序结束时裁剪额外的零...但是这将是一个悲观，除非长度已经接近于2。

Answer 1

我和Numba一起去了解基数排序有多快。 Numba（通常）表现良好的关键是写出所有循环，这非常有启发性。我最终得到了以下内容：

from numba import jit

@jit
def radix_loop(nbatches, batch_m_bits, bitsums, a, out):
    mask = (1 << batch_m_bits) - 1
    for shift in range(0, nbatches*batch_m_bits, batch_m_bits):
        # set bit sums to zero
        for i in range(bitsums.shape[0]):
            bitsums[i] = 0

        # determine bit sums
        for i in range(a.shape[0]):
            j = (a[i] & mask) >> shift
            bitsums[j] += 1

        # take the cumsum of the bit sums
        cumsum = 0
        for i in range(bitsums.shape[0]):
            temp = bitsums[i]
            bitsums[i] = cumsum
            cumsum += temp

        # sorting loop
        for i in range(a.shape[0]):
            j = (a[i] & mask) >> shift
            out[bitsums[j]] = a[i]
            bitsums[j] += 1

        # prepare next iteration
        mask <<= batch_m_bits
        # cant use `temp` here because of numba internal types
        temp2 = a
        a = out
        out = temp2

    return a

从4个内圈开始，很容易看出它是第4个，因此很难用Numpy进行矢量化。

解决这个问题的一种方法是从Scipy中引入一个特定的C ++函数：scipy.sparse.coo.coo_tocsr。它与上面的Python函数几乎完全相同的内部循环，因此可以滥用它来编写更快的＆＃34;矢量化＆＃34; Python中的基数排序。也许是这样的：

from scipy.sparse.coo import coo_tocsr

def radix_step(radix, keys, bitsums, a, w):
    coo_tocsr(radix, 1, a.size, keys, a, a, bitsums, w, w)
    return w, a

def scipysparse_radix_perbyte(a):
    # coo_tocsr internally works with system int and upcasts
    # anything else. We need to copy anyway to not mess with
    # original array. Also take into account endianness...
    a = a.astype('<i', copy=True)
    bitlen = int(a.max()).bit_length()
    radix = 256
    work = np.empty_like(a)
    _ = np.empty(radix+1, int)
    for i in range((bitlen-1)//8 + 1):
        keys = a.view('u1')[i::a.itemsize].astype(int)
        a, work = radix_step(radix, keys, _, a, work)
    return a

^{编辑：稍微优化一下这个功能..参见编辑历史。}

如上所述LSB基数排序的一个低效率是数组在RAM中完全洗牌多次，这意味着CPU缓存不能很好地使用。为了尝试减轻这种影响，可以选择先使用MSB基数排序进行传递，将项目放在大致正确的RAM块中，然后使用LSB基数排序对每个结果组进行排序。这是一个实现：

def scipysparse_radix_hybrid(a, bbits=8, gbits=8):
    """
    Parameters
    ----------
    a : Array of non-negative integers to be sorted.
    bbits : Number of bits in radix for LSB sorting.
    gbits : Number of bits in radix for MSB grouping.
    """
    a = a.copy()
    bitlen = int(a.max()).bit_length()
    work = np.empty_like(a)

    # Group values by single iteration of MSB radix sort:
    # Casting to np.int_ to get rid of python BigInt
    ngroups = np.int_(2**gbits)
    group_offset = np.empty(ngroups + 1, int)
    shift = max(bitlen-gbits, 0)
    a, work = radix_step(ngroups, a>>shift, group_offset, a, work)
    bitlen = shift
    if not bitlen:
        return a

    # LSB radix sort each group:
    agroups = np.split(a, group_offset[1:-1])
    # Mask off high bits to not undo the grouping..
    gmask = (1 << shift) - 1
    nbatch = (bitlen-1) // bbits + 1
    radix = np.int_(2**bbits)
    _ = np.empty(radix + 1, int)
    for agi in agroups:
        if not agi.size:
            continue
        mask = (radix - 1) & gmask
        wgi = work[:agi.size]
        for shift in range(0, nbatch*bbits, bbits):
            keys = (agi & mask) >> shift
            agi, wgi = radix_step(radix, keys, _, agi, wgi)
            mask = (mask << bbits) & gmask
        if nbatch % 2:
            # Copy result back in to `a`
            wgi[...] = agi
    return a

计时（我的系统上每个设置都具有最佳性能）：

def numba_radix(a, batch_m_bits=8):
    a = a.copy()
    bit_len = int(a.max()).bit_length()
    nbatches = (bit_len-1)//batch_m_bits +1
    work = np.zeros_like(a)
    bitsums = np.zeros(2**batch_m_bits + 1, int)
    srtd = radix_loop(nbatches, batch_m_bits, bitsums, a, work)
    return srtd

a = np.random.randint(0, 1e8, 1e6)
%timeit numba_radix(a, 9)
# 10 loops, best of 3: 76.1 ms per loop
%timeit np.sort(a)
#10 loops, best of 3: 115 ms per loop
%timeit scipysparse_radix_perbyte(a)
#10 loops, best of 3: 95.2 ms per loop
%timeit scipysparse_radix_hybrid(a, 11, 6)
#10 loops, best of 3: 75.4 ms per loop

按照预期，Numba的表现非常出色。而且，通过对现有C扩展的一些巧妙应用，可以击败numpy.sort。 IMO在优化级别上你已经得到了它的价值 - 它也考虑了Numpy的附加组件，但我不会真正考虑我的答案中的实现＆＃34;矢量化＃＆＃34 34;：大部分工作都是在外部专用功能中完成的。

令我印象深刻的另一件事是对基数选择的敏感性。对于我尝试的大多数设置，我的实现仍然比numpy.sort慢，所以在实践中需要某种启发式方法来提供全面的性能。

Answer 2

你能把它改成一次8位的计数/基数排序吗？对于32位无符号整数，创建字节字段出现次数的矩阵[4] [257]，使得对数组进行一次读取传递以进行排序。 matrix [] [0] = 0，matrix [] [1] =出现的数量为0，....然后将计数转换为索引，其中matrix [] [0] = 0，matrix [] [1] =字节数== 0，矩阵[] [2] =字节数== 0 +字节数== 1，......不使用最后一个计数，因为它会索引数组的结尾。然后进行4次基数排序，在原始数组和输出数组之间来回移动数据。每次工作16位需要一个矩阵[2] [65537]，但只需要2次传递。示例C代码：

size_t mIndex[4][257] = {0};            /* index matrix */
size_t i, j, m;
uint32_t u;
uint32_t *pData;                        /* ptr to original array */
uint32_t *pTemp;                        /* ptr to working array */
uint32_t *pSrc;                         /* working ptr */
uint32_t *pDst;                         /* working ptr */
/* n is size of array */
    for(i = 0; i < n; i++){             /* generate histograms */
        u = pData[i];
        for(j = 0; j < 4; j++){
            mIndex[j][1 + (size_t)(u & 0xff)]++; /* note [1 + ... */
            u >>= 8;
        }       
    }
    for(j = 0; j < 4; j++){             /* convert to indices */
        for(i = 1; i < 257; i++){       /* (last count never used) */
            mIndex[j][i] += mIndex[j][i-1]
        }       
    }
    pDst = pTemp;                       /* radix sort */
    pSrc = pData;
    for(j = 0; j < 4; j++){
        for(i = 0; i < count; i++){     /* sort pass */
            u = pSrc[i];
            m = (size_t)(u >> (j<<3)) & 0xff;
        /*  pDst[mIndex[j][m]++] = u;      split into 2 lines */
            pDst[mIndex[j][m]] = u;
            mIndex[j][m]++;
        }
        pTmp = pSrc;                    /* swap ptrs */
        pSrc = pDst;
        pDst = pTmp;
    }