Aus Python-Code C-Code generieren mit einem Dekorator

C-Code aus Python-Code zu generieren ist eine schwierige Aufgabe, da die beiden Sprachen unterschiedliche Syntax und Funktionen haben. Es ist jedoch möglich, die Kluft zwischen den beiden Sprachen mit Hilfe eines einfachen Dekorators zu überbrücken.

Numba ist ein Paket für Python das in der Lage ist, Python-Code direkt in C-Code zu übersetzen. Dabei steht Numba unter der BSD-Lizenz und darf für kommerzielle Zwecke genutzt werden. Die Installation erfolgt mit Pip.

pip install numba

In der Praxis bedeutet das, dass ganz normaler Python-Code entwickelt werden kann, und durch das ergänzen einer einzigen Zeile wird eine Funktion unter den richtigen Voraussetzungen um eine Vielfaches schneller. Spoiler: Im Praxisbeispiel weiter unten erreichen wir eine Beschleunigung um 300%.

Numba kann seine Stärken am besten ausspielen, wenn Schleifen oder NumPy-Funktionen im Spiel sind. Wurde im bestehenden Python-Projekt eine langsame Funktion identifiziert, z. B. mit Hilfe des line profilers, und diese ist Schleifen- oder NumPy-lastig, kommt sie zur Optimierung mit Numba in Frage.

Schauen wir uns die folgende Funktion zur Berechnung von Matrix-Intervallen an.

def get_matrix_intervals(matrix, flag_or_value, flag=True):
    ''' Start and end indices for matrix preset flags '''
    presetBoolMat = ((matrix & flag_or_value).astype(bool)) if flag else matrix == flag_or_value
    fPad = np.zeros((presetBoolMat.shape[0], 1))
    presetBoolMatDiff = np.diff(np.hstack((fPad, presetBoolMat, fPad)).astype(int))
    starts = np.where(presetBoolMatDiff == 1)
    ends = np.where(presetBoolMatDiff == -1)
    return starts, ends

get_matrix_intervals berechnet für eine Matrix mit aufeinanderfolgenden, identischen Werten die Start- und End-Indizes. So werden z. B. für die Sequenz [0,1,1,1,1,0,1,1] die Intervalle (1 bis 4) und (6 bis 7) berechnet.

>>> get_matrix_intervals(np.array([[0,1,1,1,1,0,1,1]]), 1, flag=False)
((array([0, 0]), array([1, 6])), (array([0, 0]), array([5, 8])))

Wie schnell ist diese Funktion, wenn sie allein durch den Python-Interpreter ausgeführt wird?

>>> timeit.timeit('get_matrix_intervals(np.array([[0,1,1,1,1,0,1,1]]), 1, flag=False)', 'from __main__ import ' + ', '.join(globals()))
11.411816049000663

Sie braucht gut 11 Sekunden für eine Millionen Iterationen. Nun wenden wir den Dekorator von Numba an und verwenden Datentypen, die Numba versteht.

from numba import jit
from numba.types import bool_, int_

@jit(nopython=True)
def get_matrix_intervals(matrix, flag_or_value, flag=True):
    ''' Start and end indices for matrix preset flags '''
    presetBoolMat = ((matrix & flag_or_value).astype(bool_)) if flag else matrix == flag_or_value
    fPad = np.zeros((presetBoolMat.shape[0], 1))
    presetBoolMatDiff = np.diff(np.hstack((fPad, presetBoolMat, fPad)).astype(int_))
    starts = np.where(presetBoolMatDiff == 1)
    ends = np.where(presetBoolMatDiff == -1)
    return starts, ends

timeit.timeit('get_matrix_intervals(np.array([[0,1,1,1,1,0,1,1]]), 1, flag=False)', 'from __main__ import ' + ', '.join(globals()))
5.79446492399984

Die Funktion ist etwa doppelt so schnell wie vorher. Nicht schlecht dafür, dass wir fast nichts getan haben. Aber das ist nicht das Ende der Fahnenstange. Wendet man zusätzlich die Performance-Tips von Numba an, stellt man fest, dass die Funktion ungünstig implementiert wurde. Zum einen verschwendet sie durch zu großzügige Datentypen Arbeitsspeicher und zum anderen ist sie durch den flag-Parameter unnötig komplex, wodurch der generierte C-Code ineffizient wird. Wendet man das Wissen an, was man über die Eingabe-Parameter hat, kann man die Funktion wie folgt verbessern.

from numba import jit
from numba.types import bool_, int8, int64, uint16, UniTuple

@jit(UniTuple(UniTuple(int64[:], 2), 2)(uint16[:,:], uint16), nopython=True)
def get_matrix_intervals(matrix, flag):
    ''' Start and end indices for matrix preset flags '''
    presetBoolMat = (matrix & flag).astype(bool_)
    fPad = np.zeros((presetBoolMat.shape[0], 1), dtype=bool_)
    presetBoolMatDiff = np.diff(np.hstack((fPad, presetBoolMat, fPad)).astype(int8))
    starts = np.where(presetBoolMatDiff == 1)
    ends = np.where(presetBoolMatDiff == -1)
    return starts, ends

timeit.timeit('get_matrix_intervals(np.array([[0,1,1,1,1,0,1,1]], dtype="u2"), 1)', 'from __main__ import ' + ', '.join(globals()))
3.575752114000352

get_matrix_intervals ist nun mehr als 3 Mal so schnell wie vorher. Wir haben dafür den flag-Parameter entfernt und geben Numba Informationen über die Datentypen der Ein-und Ausgaben. Dadurch muss der durch Numba genutzte JIT diese Typen nicht selbst zur Laufzeit bestimmen. Zusätzlich kann man mit dem Parameter cache=True dafür sorgen, dass die einmal kompilierte Funktion im Dateisystem-Ordner __pycache__ hinterlegt wird und dadurch nicht bei jeder Ausführung neu erzeugt werden muss.