我试过用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函数几乎相同，因此可以滥用它来在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的附加组件，但我不会真正考虑我的答案中的实现“矢量化”：大部分工作是在外部专用功能中完成的。
另一件让我印象深刻的事情是对基数选择的敏感性。对于我尝试的大多数设置，我的实现仍然比numpy.sort慢，因此在实践中需要某种启发式方法来提供全面的良好性能。

你能把它改成一次8位的计数/基数排序吗？对于32位无符号整数，创建一个字节字段出现计数的矩阵[4][257]，对要排序的数组进行一次读取。矩阵[][0]= 0，矩阵[][1]= 0的占用数，.然后将计数转换成索引，其中矩阵[][0]= 0，矩阵[][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;
    }

我实际上使用Cython创建了一个基数排序。从我的测试来看，它比Rust或C/C++中的基数实现快了大约5%。而且它比np. sort快了很多（我想大概是10%）。地址：https://github.com/Ohmagar/Radix_cython/blob/main/parallel_radix_5.pyx
我做了一些漂亮的事情来减少处理，通过number_of_digits对元素进行预排序，确保数字只在有数字要排序时才被排序到桶中。因此，“10”将只被处理两次而不是8次（如果9_999_999 < max_element < 10_000_000）。我从头开始在python中构建它作为POC，并对它做了越来越多的工作。一旦我无法获得更多的速度，我就在Cython中重写了它，并开始进行更多的修补。最后一步是对每个“digit_chunk”进行并行处理，这最终导致我的实现比任何可比的东西都快，特别是numpy. sort。
我刚刚看到我可以让它更快一点，可能在处理器功能中，通过让预排序也并行完成。不知道我怎么没发现。
你可以去看看。