[SciPy-user] The fastest kd-tree known to man?

Sturla Molden sturla@molden...
Fri Jan 23 22:16:56 CST 2009



Yesterday evening I was experimenting with Anne Archibald's cKDTree from
the latest SciPy superpack. The speed of this implementation is really
amazing. So I decided to try to make something even faster.


Building on Anne's code, I tried to variations:


1. Use multiprocessing (the official backport from Python 2.6) for
parallel queries. All the data was stored in shared memory (allocated
using multiprocessing.RawArray).

2. Modify the cKDTree C-code with OpenMP pragmas, and let the compiler do
the rest.



Here are some results from my dual core laptop:

http://folk.uio.no/sturlamo/kdtree/bench1.png
http://folk.uio.no/sturlamo/kdtree/bench2.png

Black: single-threaded cKDTree from SciPy superpack rc2
Blue: cKDTree + multiprocessing + shared memory
Red: cKDTree + OpenMP


As you can see, the OpenMP'd version is the fastest. The overhead from
using OpenMP seem to be negigible. It is even the faster option for the
smallest data sets. Not bad for a single line of code:


#pragma omp parallel for schedule(guided) private(__pyx_v_c)

right above the for loop on line 1870 in

http://www.scipy.org/scipy/scipy/browser/trunk/scipy/spatial/ckdtree.c?rev=4957


Using multiprocessing incures some more overhead, on the order of a few
seconds. But for more substantial work it scales almost as well as OpenMP.
Some overhead is expected when working with Python.


The code and results (incl. Windows binaries) are in the file:

http://folk.uio.no/sturlamo/kdtree/parallel numpy.zip

Which should have MD5 checksum
075f045592a9500c5ea3e48975094f71 *parallel numpy.zip

(Yes that is an empty space in the file name.)

The zipfile includes:

- parallel_kdtree.py:
parallel implemention using multiprocessing

- multiprocessing_utils.py:
a few useful helper functions for making numpy and multiprocessing
cooperate. It could be useful for the scipy cookbook.

-ckdtree.c:
source file with OpenMP pragma

- Win32 binary folder:
Compiled with gcc 4.4.0 ('mingw' binary from gfortran). The pthread DLL is
needed to run OpenMP and comes from the gfortran 'mingw' distro. And as
with all prebuilt binaries: I am a nice guy and don't write malware, but
you run it at your own risk. This code is for testing purposes only.



As for the cookbook, I think it is time to remove my two entries there.
Both of them are obsolete by now.




Regards,

Sturla Molden






More information about the SciPy-user mailing list