# Numba mergesort

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):

Graphed results from a longer run:

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?

• 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.