Python: Richtige JIT-Optimierung – damit nichts schiefgeht

JIT-optimierter Code beschleunigt Python-Programme auf GPUs rasant. Doch wer nicht aufpasst, zerstört schnell den Performancegewinn oder rechnet sogar falsch.

Lesezeit: 17 Min.
In Pocket speichern
vorlesen Druckansicht Kommentare lesen 24 Beiträge

(Bild: Shutterstock)

Von
  • Eldar Sultanow
Inhaltsverzeichnis

Forscherinnen und Entwickler sehen sich heute oftmals mit Fragestellungen konfrontiert, die immer komplexer werden – auch die Datenmengen nehmen an Komplexität zu. Die Leistung der Hardware wird aber ebenfalls immer besser und steht sogar per Knopfdruck in der Cloud zur Verfügung. Allerdings ist es gar nicht so einfach, diese Leistung optimal auszuschöpfen. Nicht selten laufen Berechnungsroutinen auf leistungsstarker Hardware, nutzen dabei aber nur einen Bruchteil dieser Leistung. Wer on Demand beispielsweise in der AWS-Cloud die EC2 P3-Instanz vom Typ p3.8xlarge (diese beinhaltet vier GPUs) verwendet, bezahlt pro Stunde 12,24 US-Dollar.

Wie sträflich wäre es, dort eine Berechnungsroutine laufen zu lassen, die weder die JIT-Optimierung noch die GPU nutzt: Das dauert unnötig länger und kostet dementsprechend mehr Geld. Von welchen Faktoren hier die Rede ist, wie sehr viel schneller eine JIT- oder GPU-optimierte Routinen läuft und wie Entwicklerinnen und Entwickler sie mit Python umsetzen können, zeigt dieser Artikel. Ausgewählte Beispiele veranschaulichen die Fallstricke, die auftreten können.

Die Programmiersprachen C++ und Python sind sehr gut geeignet, um rechenintensive, leistungsoptimierte Routinen zu schreiben – erst im letzten Jahr hat Nvidia Neuerungen für beide Programmiersprachen eingeführt. Dieser Artikel konzentriert sich auf Python, betrachtet dabei aber auch Alternativen über den Tellerrand hinaus (Mathematica, C++). Daher benötigen Entwicklerinnen das entsprechende Handwerkszeug:

  • Eine Entwicklungsumgebung (IDE): Die hier vorgestellten Beispiele, Programme und Codefragmente wurden in der IDE Visual Studio Code entwickelt. Es gibt auch weitere gut geeignete Umgebungen wie PyCharm.
  • Eine Python-Installation: Die hier verwendete Installation umfasst die im November letzten Jahres veröffentlichte Anaconda-Version Anaconda Individual Edition 2021.11, die Python 3.9.7 mitbringt.
  • Entsprechende Pakete: zum Ausprobieren der hier entwickelten Routinen müssen die Pakete Numba (JIT-Compiler zur Übersetzung von Python-Code in schnellen Maschinencode), CUDA Toolkit (GPU-beschleunigte Bibliotheken) Pyculib (Python bindings für CUDA), math (Paket mit mathematischen Funktionen), pandas (Paket zur Verarbeitung und Speicherung/Export von Daten), NumPy (Paket für numerische Berechnungen), multiprocessing (Paket für Prozess-basierte Parallelisierung) und gmpy2 (Paket für hoch performante Multipräzisionsarithmetik) vorhanden sein; time und sys für Zeitmessung und Systemfunktionen sind standardmäßig in Python vorhanden.
  • Codeassistent und Qualitätsprüfer: Pylint erkennt Probleme, Fehler im Code sowie Verstöße gegen die Konventionen und Empfehlungen der Python-Community und weist auf diese hin.

Angeregt durch den französischen Mathematiker Jacques Ozanam beschäftigte sich sein italienischer Kollege Pietro Mengoli in den 1670er Jahren mit speziellen diophantischen Gleichungen – sie lassen lediglich ganzzahlige Lösungen zu. Ein bekanntes Problem, das von solchen Gleichungen handelt, ist das Sechs-Quadrate-Problem: Finde drei natürliche Zahlen x, y, z, deren Differenzen z-y, z-x, y-x jeweils Quadratzahlen sind, wobei die Differenzen z2-y2, z2-x2, y2-x2 der Quadrate ebenfalls Quadratzahlen sind.

Die Ergebnisse umfassen sehr große Zahlen. Ozanam hatte damals schon eine Lösung präsentiert, nämlich x=1873432, y=2288168 und z=2399057. Derartige Probleme in höheren Zahlenregionen sind als Anschauungsmaterial ideal, da es Fallstricke in der JIT-optimierten Programmierung mit sehr großen Zahlen gibt, wie sich später noch zeigt.

Bemerkenswert ist, dass Ozanam bereits im 17. Jahrhundert derartig großen Zahlen gefunden hat, für die diese Gleichung erfüllt ist. Hatte er doch keine technischen Hilfsmittel wie Computer, Programmiersprachen oder Ähnliches zur Hand. Unsere heutigen Werkzeuge sind in der Lage, solch aufwendige Probleme zu lösen. So hat beispielsweise ein Team am Massachussetts Institute of Technology (MIT) im Jahr 2019 mit x=-80538738812075974, y=80435758145817515 und z=12602123297335631 eine Lösung für die diophantische Gleichung x3+y3+z3=42 gefunden. Unter den Zahlen bis 100 war die 42 die letzte noch offene Zahl, die als Summe dreier Kubikzahlen darstellbar ist. Es wundert nicht, dass das Interesse an diesem Rätsel auch bei Fans des Science-Fiction-Klassikers "Per Anhalter durch die Galaxis" groß war. Das Team hat für die Berechnung die ungenutzte Leistung von mehr als einer halben Million Heim-PCs verwendet. Derartige Gleichungen gehören zu dem seit 160 Jahren ungelösten "Summe dreier Kubikzahlen"-Problem. Es gibt weitere Fragestellungen, die ohne hohe Rechenpower auch in Hunderten von Jahren nicht zu bewältigen wären. Hierzu zählen vor allem Aufgaben aus dem Umfeld der künstlichen Intelligenz (KI) wie Spracherkennung, Bilderkennung oder Mustererkennung in massenweise numerischen Daten.

Zur Vereinfachung liegt der Fokus auf den drei Quadratzahlen x2, y2 und z2, deren Differenz wieder Quadratzahlen sind. Hier gibt es auch Ergebnisse, die nicht ganz so groß sind: x=153, y=185 und z=697 oder x=520, y=533 und z=925. Ein kleiner Quercheck 9252-5332 ergibt 571536, was tatsächlich eine Quadratzahl ist, nämlich 7562. Nachzuprüfen, ob die anderen Differenzen ebenfalls stimmen, sei an dieser Stelle der Leserschaft überlassen. Eine große Anzahl solcher Treffer bis in die Zahlenregion von 12 Millionen wurden mithilfe einer Just-in-Time-optimierten Routine (JIT) ermittelt und ist in einem GitHub-Repository (siehe Datei pythagorean_12000000.csv) verfügbar. Der dafür zuständige Algorithmus ist auf StackExchange Mathematics beschrieben, für JIT optimiert und ebenfalls im GitHub-Repository (siehe Datei pythagorean_gendata_jit.py) verfügbar.

@jit('void(uint64, uint64[:])')
def generateData(limit: np.uint64, triples: np.ndarray) -> np.uint32:
    A=np.zeros(limit, dtype=np.uint64)
    B=np.zeros(limit, dtype=np.uint64)
    rows = np.uint32(0)
    for w in np.arange(1, limit+1, dtype=np.uint64):
        count = np.uint32(0)
        for a in np.arange(1, w+1, dtype=np.uint64):
            sqr = np.uint64(sqrt(w*w-a*a))
            if sqr*sqr == w*w-a*a:
                count+=1
                A[count]=a
                B[count]=np.uint64(sqrt(w*w-a*a))
        if count>1:
            for i in np.arange(1, count+1, dtype=np.uint64):
                for j in np.arange(1, count+1, dtype=np.uint64):
                    if i!=j:
                        x=np.uint64(A[i])
                        t=np.uint64(B[i])
                        s=np.uint64(A[j])
                        z=np.uint64(B[j])
                        if z > x:
                            y = np.uint64(sqrt(z*z-x*x))
                            if y*y == z*z-x*x:
                                arr = np.array([x, y, z, s, t, w], dtype=np.uint64)
                                triples[rows] = arr
                                rows+=1
    return rows

Listing 1: Optimierter Algorithmus zur Lösung des Sechs-Quadrate-Problems

Wer die Routine in Listing 1 zur Suche von Lösungen startet, die Werte bis zu 50000 annehmen können (limit=50000), benötigt dazu auf einem HP Z-Book (32 GByte RAM, Intel Core i7-9750H, 2.60GHz, 6 Kerne, 12 logische, NVIDIA Quadro T1000) 4,8 Sekunden und findet 1074 Tripel für x, y und z. Führt man die Routine ohne Optierung auf demselben Rechner aus, dann benötigt sie 1082,8 Sekunden (18 Minuten), um diese 1074 Tripel zu ermitteln. Demnach hat die JIT-Optimierung eine Beschleunigung um etwa den Faktor 225 erbracht.