jit wrapped numpy sort & unique is more than 5x slower

[game-difficulty]

example:
@numba.njit
def jit_sort(arr):
return np.sort(arr)

vs

np.sort(arr)

[game-difficulty]

Issue Description
I have encountered a significant performance drop when using np.sort and np.unique within functions decorated with numba.jit. The performance is more than 5 times slower compared to their execution without the JIT compilation, which is contrary to the expected optimization usually provided by Numba.

Steps to Reproduce
Below is a minimal example that demonstrates the slower performance of the jitted functions compared to the non-jitted versions:

import numpy as  np
from numba import jit
import time

data = np.random.randint(0,2**64,int(1e7),dtype='uint64')

# Non-jitted functions
def sort_array(data):
    return np.sort(data)

def unique_array(data):
    return np.unique(data)

# Jitted functions
@jit(nopython=True)
def sort_array_jitted(data):
    return np.sort(data)

@jit(nopython=True)
def unique_array_jitted(data):
    return np.unique(data)

# Timing non-jitted functions
start = time.time()
sort_array(data)
print("Non-jitted sort time: ", time.time() - start)

start = time.time()
unique_array(data)
print("Non-jitted unique time: ", time.time() - start)

# Timing jitted functions
start = time.time()
sort_array_jitted(data)
print("Jitted sort time: ", time.time() - start)

start = time.time()
unique_array_jitted(data)
print("Jitted unique time: ", time.time() - start)

Expected Behavior
I expect that using numba.jit to decorate functions involving np.sort and np.unique should not degrade performance, considering Numba’s purpose to accelerate Python code, especially for numerical computations.

Actual Behavior
The execution times of the jitted versions of np.sort and np.unique are substantially longer than those of the non-jitted versions:
Non-jitted sort time: 0.076s
Jitted sort time: 0.685s
Non-jitted unique time: 0.116s
Jitted unique time: 0.799s

Environment Information
Operating System: Windows 11
Python Version: Python 3.11.8
Numba Version: Numba 0.60.0rc1
NumPy Version: NumPy 2.0.0rc2
(About twice as slow with numba 0.59 and numpy 1.26)

[rlav440]

In the minimal example you have given, the time you are reporting includes the time taken to compile the code.
Do you still see the same slowdown if you call the jit'd functions once, to allow them to compile, before you time their execution?

[game-difficulty]

In the minimal example you have given, the time you are reporting includes the time taken to compile the code. Do you still see the same slowdown if you call the jit'd functions once, to allow them to compile, before you time their execution?

I know that jit compiled code will be slower the first few times it runs, and I took that into account in my testing. In fact, the execution time given above is after the 3rd call to the function

[game-difficulty]

numpy 2.0.0 release seems to have optimized sorting significantly, perhaps numba needs some changes to go along with this update?

[esc]

@game-difficulty thank you for reporting this. From memory: there are duplicates to this on the issue tracker, it might be worth a look.

IIRC np.sort has historically suffered from severe performance issues but there was never enough time/resources/interest for anyone to fix this -- i.e. this is waiting for a champion to tackle the issue and pull-requests would be very welcome! I think you would make a lot of people happy if this was finally resolved. 🚀

[esc]

This is mentioned in the docs:

[game-difficulty]

I see. But I believe numba 0.60 needs to do something to follow up on numpy 2.0's optimization of sorting

[esc]

I see. But I believe numba 0.60 needs to do something to follow up on numpy 2.0's optimization of sorting

The way Numba works is that all NumPy functions are reimplemented in pure python so they can be compiled. Essentially someone will need to look at the Numba implementation of the numpy sorting function. This will not happen unless someone from the community volunteers and fixes it. Everyone realizes this is an important use-case, so polished contributions are very likely to attract reviewer attention.

[esc]

FWIW: the sorting code to be able to compile np.sort begins here:

https://github.com/numba/numba/blob/main/numba/np/arrayobj.py#L6364

Hope that helps!