2
\$\begingroup\$

I have written a mergesort in Python/Numba:

import numba as nb
import numpy as np

@nb.jit( nopython=True )
def merge( x ):

    n = x.shape[0]
    width=1

    r   = x.copy()
    tgt = np.empty_like( r )
    while width<n:
        i=0
        while i<n:
            istart = i
            imid = i+width
            iend = imid+width
            # i has become i+2*width
            i = iend

            if imid>n:
                imid = n

            if iend>n:
                iend=n
            _merge( r, tgt, istart, imid, iend)

        # Swap them round, so that the partially sorted tgt becomes the result,
        # and the result becomes a new target buffer
        r, tgt = tgt, r
        width*=2

    return r

@nb.jit( nopython=True )
def _merge( src_arr, tgt_arr, istart, imid, iend ):
    """ The merge part of the merge sort """
    i0   = istart
    i1   = imid
    for ipos in range( istart, iend ):
        if ( i0<imid ) and ( ( i1==iend ) or ( src_arr[ i0 ] < src_arr[ i1 ] ) ):
            tgt_arr[ ipos ] = src_arr[ i0 ]
            i0+=1
        else:
            tgt_arr[ ipos ] = src_arr[ i1 ]
            i1+=1

I have written a test for it:

def test_merge_multi(self):

    n0 = 21
    n1 = 100

    for n in range( n0, n1 ):
        x = np.random.random_integers( 0, n, size=n )

        with Timer( 'nb' ) as t0:
            r = sas.merge( x )
        with Timer( 'np' ) as t1:
            e = np.sort( x, kind='merge' )
        #print( 'r:%s'%str(r))
        #print( 'e:%s'%str(e))
        print( 'nb/np performance %s'%(t0.interval/t1.interval ))
        np.testing.assert_equal( e, r )

I used this Timer class:

import time

class Timer:    

    def __init__(self,title=None):
        self.title=title

    def __enter__(self):
        if self.title:
            print( 'Beginning {0}'.format( self.title ) )
        self.start = time.clock()
        return self

    def __exit__(self, *args):
        self.end = time.clock()
        self.interval = self.end - self.start
        if self.title:
            print( '{1} took {0:0.4f} seconds'.format( self.interval, self.title ) )
        else:
            pass#
            #print( 'Timer took {0:0.4f} seconds'.format( self.interval ) )

The test results are:

nb/np performance 9307.846153856719
nb/np performance 1.1428571428616743
nb/np performance 0.7142857142925115
nb/np performance 0.8333333333302494
nb/np performance 0.9999999999814962
nb/np performance 0.9999999999777955
nb/np performance 0.8333333333456692
nb/np performance 0.8333333333302494
nb/np performance 1.0
nb/np performance 0.8333333333456692
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 0.8333333333456692
nb/np performance 0.9999999999814962
nb/np performance 1.0
nb/np performance 0.9999999999814962
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 1.0000000000185036
nb/np performance 1.2000000000044408
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 1.0000000000185036
nb/np performance 1.2000000000088817
nb/np performance 1.0
nb/np performance 1.1666666666512469
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 0.9999999999814962
nb/np performance 1.1666666666728345
nb/np performance 1.1666666666512469
nb/np performance 1.0
nb/np performance 1.0
nb/np performance 1.1666666666512469
nb/np performance 1.1666666666512469
nb/np performance 1.1666666666728345
nb/np performance 1.1666666666728345
nb/np performance 1.1666666666728345
nb/np performance 1.1666666666728345
nb/np performance 1.1666666666512469
nb/np performance 1.1666666666512469
nb/np performance 1.0
nb/np performance 1.1666666666728345
nb/np performance 1.3333333333456692
nb/np performance 1.3333333333024937
nb/np performance 1.3333333333456692
nb/np performance 1.1428571428435483
nb/np performance 1.3333333333209976
nb/np performance 1.1666666666728345
nb/np performance 1.3333333333456692
nb/np performance 1.3333333333209976
nb/np performance 1.000000000012336
nb/np performance 1.1428571428616743
nb/np performance 1.3333333333456692
nb/np performance 1.3333333333209976
nb/np performance 1.1428571428616743
nb/np performance 1.1428571428616743
nb/np performance 1.3333333333456692
nb/np performance 1.499999999990748
nb/np performance 1.2857142857074884
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857029569
nb/np performance 1.1428571428616743
nb/np performance 1.1428571428435483
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857233488
nb/np performance 1.2857142857029569
nb/np performance 1.1249999999895917
nb/np performance 1.2857142857029569
nb/np performance 1.2857142857233488
nb/np performance 1.4285714285623656
nb/np performance 1.249999999993061
nb/np performance 1.1250000000034694
nb/np performance 1.2857142857029569

Graphed results (from a different run):

enter image description here

Graphed results from a longer run:

enter image description here

Please note that for n<=20, numpy uses an insertion sort when the mergesort is called.

So you can see that for small values of n, the numba version of the mergesort beats the numpy version.

However, as n gets larger, numpy then consistently outperforms numba by a factor of 2.

Why is this? And how could I optimise the numba version to beat the numpy version for all n?

\$\endgroup\$
1
\$\begingroup\$
  • Stability

    The single most important feature of merge sort is stability: elements comparing equal retain their original order. As coded merge loses stability: shall src_arr[i0] and src_arr[i1] compared equal, src_arr[i1] is copied first. A standard remedy is to compare them backwards; in pseudocode:

    if (src_arr[i1] < src_arr[i0])
        copy src_arr[i1]
    else
        copy src_arr[i0]
    

    You may also check that numpy code at your link does exactly that.

  • Performance

    Fallback to insertion sort for small subarrays is important. The code is simpler and allows better optimization, particularly in reusing registers, which is very hard to achieve with recursive calls. It will be also interesting to inspect the generated code.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.