NumPy 2.5 macht searchsorted bis zu 25-mal schneller, indem Binärsuchen im Gleichschritt laufen
Victoria Shinder2 Min. Lesezeitnp.searchsorted – die Binärsuche von NumPy, die unter Histogrammen, Intervallsuchen und Teilen von SciPy und scikit-learn arbeitet – wurde in NumPy 2.5 grundlegend neu geschrieben. Ein Beitrag im Scientific-Python-Blog beschreibt den Weg dorthin und meldet in den Benchmarks der Autoren bis zu 25-fache Beschleunigung gegenüber NumPy 2.4.
Was langsam war. NumPy 2.4 führte für jeden Suchschlüssel eine klassische Binärsuche aus. Jeder Schritt hängt vom vorherigen Vergleich ab; die CPU wartet also auf jeweils einen Speicherzugriff, und bei großen Arrays verfehlen diese Zugriffe meist den Cache. Die einzige Optimierung – die Grenzen der vorigen Suche wiederverwenden, wenn die Schlüssel sortiert kommen – hilft bei sortierten Eingaben und bremst zufällige leicht.
Die Idee. Alle Suchen gemeinsam laufen lassen. Die Autoren schrieben das zunächst in NumPy selbst: lo und hi als Arrays mit einem Eintrag pro Schlüssel, die in jedem Schritt alle vorrücken. Schon diese vektorisierte Python-Fassung war bei vielen Schlüsseln um eine Größenordnung schneller als der native Code von 2.4, weil unabhängige Suchen der CPU mehrere gleichzeitige Speicherzugriffe erlauben. Ohne die Maske für noch aktive Suchen – jede Suche läuft genau ceil(log2(n)) Runden, bereits fertige bewegen sich einfach nicht mehr – wurde es nochmals bis zu doppelt so schnell.
Damit es in NumPy passt. Zwei Zustandsarrays pro Anfrage kosten Speicher proportional zur Zahl der Schlüssel; für eine Kernroutine hielten die Autoren das für inakzeptabel. Die Lösung: Jedes Intervall schrumpft pro Schritt um denselben Betrag, dann muss nur die linke Grenze gespeichert werden, und die Länge ist eine gemeinsame Zahl – die verzweigungsfreie Variante, die auch die Fallstudie von Algorithmica beschreibt. In C++ ist die erste Runde ausgerollt: Alle Schlüssel werden mit demselben Median verglichen, der nur einmal gelesen wird, und das Ergebnisarray braucht keinen Initialisierungsdurchlauf. Über die Ausgabe hinaus fällt O(1) Speicher an. Die Änderung kam als PR #30517.
Messwerte und ihre Grenzen. Gemessen wurde auf einem MacBook Pro mit Apple M1 Pro, mit 32-Bit-Ganzzahlen und Arrays bis 2^30 Elemente; berichtet wird der beste von 50 Läufen. Im Vergleich mit JAX, TensorFlow und PyTorch nennen die Autoren NumPy konkurrenzfähig; die beiden Letzteren verteilen Suchen auf Threads statt sie zu bündeln, und sobald das Array den Cache übersteigt, liegen alle nah beieinander. Als nächsten Schritt schlagen sie ein Eytzinger-Speicherlayout vor.
Die 2.5-Linie selbst erhielt am 10. Oktober NumPy 2.5.4, ein Patch-Release, das unter anderem linalg.svd und linalg.lstsq bei Eingaben mit Unendlich-Werten einen Fehler werfen lässt, statt Unsinn zu liefern oder zu hängen.
Quellen: Scientific-Python-Blog — https://blog.scientific-python.org/numpy/searchsorted ; Release Notes NumPy 2.5.4 — https://github.com/numpy/numpy/releases/tag/v2.5.4