I want to write a function which will take an index lefts of shape (N_ROWS,) I want to write a function which will create a matrix out = (N_ROWS, N_COLS) matrix such that out[i, j] = 1 if and only if j >= lefts[i]. A simple example of doing this in a loop is here:
class Looped(Strategy):def copy(self, lefts):out = np.zeros([N_ROWS, N_COLS])for k, l in enumerate(lefts): out[k, l:] = 1return out
Now I wanted it to be as fast as possible, so I have different implementations of this function:
- The plain python loop
- numba with
@njit - A pure c++ implementation which I call with
ctypes
Here are the results of the average of 100 runs:
Looped took 0.0011599776260009093
Numba took 8.886413300206186e-05
CPP took 0.00013200821400096175
So numba is about 1.5 times than the next fastest implementation which is the c++ one. My question is why?
- I have heard in similar questions cython was slower because it wasn't being compiled with all the optimisation flags set, but the cpp implementation was compiled with
-O3is that enough for me to have all the possible optimisations the compiler will give me? - I do not fully understand how to hand the numpy array to c++, am I inadvertently making a copy of the data here?
# numba implementation@njit
def numba_copy(lefts):out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)for k, l in enumerate(lefts): out[k, l:] = 1.return outclass Numba(Strategy):def __init__(self) -> None:# avoid compilation time when timing numba_copy(np.array([1]))def copy(self, lefts):return numba_copy(lefts)
// array copy cppextern "C" void copy(const long *lefts, float *outdatav, int n_rows, int n_cols)
{ for (int i = 0; i < n_rows; i++) {for (int j = lefts[i]; j < n_cols; j++){outdatav[i*n_cols + j] = 1.;}}
}// compiled to a .so using g++ -O3 -shared -o array_copy.so array_copy.cpp
# using cpp implementationclass CPP(Strategy):def __init__(self) -> None:lib = ctypes.cdll.LoadLibrary("./array_copy.so")fun = lib.copyfun.restype = Nonefun.argtypes = [ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),ctypes.c_long,ctypes.c_long,]self.fun = fundef copy(self, lefts):outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )self.fun(lefts, outdata, N_ROWS, N_COLS)return outdata
The full code with the timing etc:
import time
import ctypes
from itertools import combinationsimport numpy as np
from numpy.ctypeslib import ndpointer
from numba import njitN_ROWS = 1000
N_COLS = 1000class Strategy:def copy(self, lefts):raise NotImplementedErrordef __call__(self, lefts):s = time.perf_counter()n = 1000for _ in range(n):out = self.copy(lefts)print(f"{type(self).__name__} took {(time.perf_counter() - s)/n}")return outclass Looped(Strategy):def copy(self, lefts):out = np.zeros([N_ROWS, N_COLS])for k, l in enumerate(lefts): out[k, l:] = 1return out@njit
def numba_copy(lefts):out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)for k, l in enumerate(lefts): out[k, l:] = 1.return outclass Numba(Strategy):def __init__(self) -> None:numba_copy(np.array([1]))def copy(self, lefts):return numba_copy(lefts)class CPP(Strategy):def __init__(self) -> None:lib = ctypes.cdll.LoadLibrary("./array_copy.so")fun = lib.copyfun.restype = Nonefun.argtypes = [ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),ctypes.c_long,ctypes.c_long,]self.fun = fundef copy(self, lefts):outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )self.fun(lefts, outdata, N_ROWS, N_COLS)return outdatadef copy_over(lefts):strategies = [Looped(), Numba(), CPP()]outs = []for strategy in strategies:o = strategy(lefts)outs.append(o)for s_0, s_1 in combinations(outs, 2):for a, b in zip(s_0, s_1):np.testing.assert_allclose(a, b)if __name__ == "__main__":copy_over(np.random.randint(0, N_COLS, size=N_ROWS))