numba-cpp benchmarks

bench.py

from benchmarks.bench_ising import IsingModel, setup
import time

if __name__ == '__main__':
    t = time.time()
    setup()
    im = IsingModel()
    print('compile', time.time()-t)
    t = time.time()
    im.time_ising()
    print('exec', time.time()-t)

benchmarks.bench_ising.py

"""
Ising model benchmark, adapted from
http://matthewrocklin.com/blog/work/2015/02/28/Ising/
"""

from math import exp, log, e, sqrt

import numpy as np
from numba import f8, u1, void, i8, prange


kT = 2 / log(1 + sqrt(2), e)

N = 1000

random = np.random.RandomState(0)

x_start = random.randint(2, size=(N, N)).astype('i8')
x_start[x_start == 0] = -1

N_iterations = 10


def setup():
    from numba import jit
    global _update, update

    @jit('void(i8[:,:],u1,u1)', nopython=True, cache=True, fastmath=True)
    def _update(x, i, j):
        n, m = x.shape
        dE = 2 * x[i, j] * (
            x[(i-1) % n, (j-1) % m]
            + x[(i-1) % n,  j]
            + x[(i-1) % n, (j+1) % m]

            + x[i, (j-1) % m]
            + x[i, (j+1) % m]

            + x[(i+1) % n, (j-1) % m]
            + x[(i+1) % n,  j]
            + x[(i+1) % n, (j+1) % m]
        )
        if dE <= 0 or exp(-dE / kT) > np.random.random():
            x[i, j] *= -1

    @jit('void(i8[:,:])', nopython=True, cache=True, fastmath=True)
    def update(x):
        n, m = x.shape
        for i in prange(n):
            for j in prange(0, m, 2):  # Even columns first to avoid overlap
                _update(x, j, i)

        for i in prange(n):
            for j in prange(1, m, 2):  # Odd columns second to avoid overlap
                _update(x, j, i)


class IsingModel:
    def time_ising(self):
        x = x_start.copy()
        for i in prange(N_iterations):
            update(x)

ising.cc

#pragma GCC optimize("O3")
#pragma GCC target("avx")

#define N 1000
#include <stdint.h>
#include <cmath>

float kT = 2 / log(1 + std::sqrt(2));
uint8_t x[N][N];

void _update(uint8_t x[N][N], uint8_t i, uint8_t j) {
  uint8_t m = N;
  uint8_t n = N;
  uint8_t dE = 2 * x[i][j] *
               (x[(i - 1) % n][(j - 1) % m] + x[(i - 1) % n][j] +
                x[(i - 1) % n][(j + 1) % m] + x[i][(j - 1) % m] +
                x[i][(j + 1) % m] + x[(i + 1) % n][(j - 1) % m] +
                x[(i + 1) % n][j] + x[(i + 1) % n][(j + 1) % m]);
  if (dE <= 0 || std::exp(-dE / kT)) x[i][j] *= -1;
};

void update(uint8_t x[N][N]) {
  uint8_t m = N;
  uint8_t n = N;
  for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j += 2) {
      _update(x, j, i);
    }
    for (int j = 1; j < N; j += 2) {
      _update(x, j, i);
    }
  }
};

int main(int argc, char const* argv[]) {
  for (int i = 0; i < 10; i++) update(x);

  return 0;
}

requirements.txt

numba==0.52.0

[email protected]

@numba.jit
compile 0.495[s]（2回目以降0.089[s]）
exec 0.599[s] 

@cpp -O3
time ./bin は 0.070[s]