経緯

唐突に、「pybind11＋SIMD命令＋スレッドを使えば、numpyを使うより速く計算できるようになるんじゃね？」
と、思い立ち試してみました。

試したこと

まずは簡単そうなnumpy.innerから挑戦
とりあえず、１次元配列同士の内積を求めることしかできません。

使い方

パッケージ名はmylibという名前で、下記のようにコンパイルすると、mylib.cpython-37m-darwin.so という名前のファイルができます。このファイルをカレントに置いておけば、pyファイルからmylibとしてimport可能になります。１次元以外の配列を渡すとどうなるかは試してません。

clang++ -mavx2 -mfma -O3 -Wall -shared -std=c++11 -fPIC `python -m pybind11 --includes` -undefined dynamic_lookup mylib.cpp -o mylib`python3-config --extension-suffix`

Pybind11 C++ のコード

スレッドを４個使用し、SIMD命令を使って書いてあります。

# include <immintrin.h>
# include <type_traits>
# include <thread>
# include <future>
# include <pybind11/pybind11.h>
# include <pybind11/numpy.h>

using namespace std;
namespace py = pybind11;

//T,Uの組み合わせの確認マクロ
# define type_check() static_assert(((is_same<double,U>::value) && (is_same<__m256d,T>::value)) || \
    ((is_same<float,U>::value) && (is_same<__m256,T>::value)),"type mismatch")

//
//simd組み込み関数をdouble/float両方で使えるようにするための、template関数
//
template <typename T,typename U>
static inline T m256_load(const U* d1){
    type_check();
    if(is_same<double,U>::value){
        return (T)_mm256_load_pd((const double*)d1);
    }else{
        return (T)_mm256_load_ps((const float*)d1);
    }
}
template <typename T,typename U>
static inline void m256_store(U* d1, T reg){
    type_check();
    if(is_same<double,U>::value){
        _mm256_store_pd((double*)d1, (__m256d)reg);
    }else{
        _mm256_store_ps((float*)d1, (__m256)reg);
    }
}
template <typename T>
static inline T m256_fmadd(T d1, T d2, T d3){
    if(is_same<__m256d,T>::value){
        return (T)_mm256_fmadd_pd((__m256d)d1,(__m256d)d2,(__m256d)d3);
    }else{
        return (T)_mm256_fmadd_ps((__m256)d1,(__m256)d2,(__m256)d3);
    }
}
template <typename T>
static inline T m256_setzerp(){
    if(is_same<__m256d,T>::value){
        return (T)_mm256_setzero_pd();
    }else{
        return (T)_mm256_setzero_pd();
    }
}
//simd無しの内積を求める関数
template <typename U>
static U calc_inner_nosimd(const U* v1, const U* v2, ssize_t size){
    U result = 0;
    for(ssize_t i = 0; i < size; i++){
        result += v1[i] * v2[i];
    }
    return result;
}
//simd命令を使って、内積を求める。
template <typename T,typename U>
static U calc_inner(const U* v1, const U* v2, ssize_t size){
    type_check();
    constexpr int step = (is_same<float,U>::value) ? 8 : 4;//simdレジスタ１個に入るデータの数
    if( size < step ){
        return calc_inner_nosimd<U>(v1,v2,size);//配列が短すぎる場合には普通に計算
    }
    auto remain = size % step;//simdに入りきらないあまり分。後で別に計算する。
    size -= remain;
    auto sum = m256_setzerp<T>();
    for(ssize_t i = 0; i < size; i+= step){
        auto d1 = m256_load<T,U>(v1);
        auto d2 = m256_load<T,U>(v2);
        sum = m256_fmadd<T>(d1,d2,sum);//sum=d1*d2+sum
        v1 += step;
        v2 += step;
    }
    U result = calc_inner_nosimd<U>(v1,v2,remain);//４個または８個のあまりの計算を行う。
    
    U results[step];
    m256_store<T,U>(results, sum);//計算結果を配列に書き出す。
    for(int i=0; i < step; i++){//計算結果の集計
        result += results[i];
    }
    return result;
}
//実験用。使ってない。
template <typename T, typename U>
U inner_single_thread(py::array_t<U> arr1, py::array_t<U> arr2) {
    auto size = arr1.size();
    auto *d1 = arr1.data();
    auto *d2 = arr2.data();
    auto result = calc_inner<T,U>(d1,d2,size);
    return result;
}
//NUM_THREADSの数のスレッドに計算を割り当てて計算する関数。
//future/asyncを使用
template <typename T, typename U, unsigned NUM_THREADS>
U inner(py::array_t<U> arr1, py::array_t<U> arr2) {
    auto size1 = arr1.size();
    if( size1 != arr2.size() ){
        throw invalid_argument("shapes not aligned");//２個の配列の長さが違う。
    }
    if( size1 < 16 ){//配列が短い場合にはスレッドを起こさない。この数は環境によって要調整
        return calc_inner<T,U>(arr1.data(), arr2.data(), size1);
    }
    auto remain = size1 % NUM_THREADS;//スレッドに割り振れないあまり。
    auto size = size1 / NUM_THREADS;//スレッド１個あたりの計算数
    auto *d1 = arr1.data();
    auto *d2 = arr2.data();
    vector<future<U>> futures;
    if( size > 0 ){
        for(unsigned i = 0; i < NUM_THREADS-1; i++){//自分を除いた数のスレッドを実行
            ssize_t start = i * size;
            auto d1start = d1 + start;
            auto d2start = d2 + start;
            futures.push_back(async(launch::async,[d1start,d2start,size](){//別スレッドでcalc_innerを実行
                return calc_inner<T,U>(d1start, d2start, size);
            }));
        }
    }
    ssize_t start = (NUM_THREADS-1) * size;
    U result = calc_inner<T,U>(d1 + start, d2 + start, size + remain);//自スレッド分の計算
    for(auto&& f : futures){//他のスレッドの計算結果を集計
        result += f.get();
    }
    return result;
}

PYBIND11_MODULE(mylib, m) {
    m.doc() = "pybind11 mylib";
    m.def("inner", &inner<__m256d,double, 4>);//double配列の場合
    m.def("inner", &inner<__m256,float, 4>);//float配列の場合
}

計測プログラム

mylib.cpython-37m-darwin.so と同じ場所に置いて、実行すれば結果が表示されます。

import numpy
import mylib
import time

# 配列の長さ。１次元の配列です。
maxsize = 256 * 256 * 256 * 16

# １回の計算にかかる時間を測る。
def laptime(f,vec1,vec2):
    t = time.time()
    result = f(vec1, vec2)
    t = time.time() - t
    return t

def print_time_inner(v1, v2):
    time_numpy_inner = 0
    time_mylib_inner = 0

    #warming up
    result_numpy_inner = numpy.inner(v1,v2)
    result_mylib_inner = mylib.inner(v1,v2)
    #できるだけ公平になるように、ループ内で交互に呼ぶ
    for i in range(16):
        time_mylib_inner += laptime(mylib.inner,v1,v2)
        time_numpy_inner += laptime(numpy.inner,v1,v2)
    #結果の表示
    print("numpy.inner %5f(sec)" % time_numpy_inner,      "result=",result_numpy_inner)
    print("mylib.inner %5f(sec)" % time_mylib_inner,      "result=",result_mylib_inner)

print("double array size=",maxsize)
vector1 = numpy.random.rand(maxsize)
vector2 = numpy.random.rand(maxsize)
print_time_inner(vector1,vector2)

print("float array size=",maxsize)
vector1 = vector1.astype('float32')
vector2 = vector2.astype('float32')
print_time_inner(vector1,vector2)

計測結果

1個目の数字が268435456(256 * 256 * 256 * 16)回を１６回ループにかかった時間(秒)です。

double array size= 268435456
numpy.inner 2.596892(sec) result= 67115867.58791062
mylib.inner 2.569095(sec) result= 67115867.58790968

float array size= 268435456
numpy.inner 1.791725(sec) result= 66669584.0
mylib.inner 1.277165(sec) result= 66669976.0

多少はnumpyよりは、少しだけ速くなったようです。
floatで1.35倍
doubleで1.02倍

timeitを使ってない理由

かなり大きな配列を作っているためか、timeitを使うと実行順序によって時間が変化したため。

計測した環境について

numpy           1.16.1
Python 3.7.2

clang++
Apple LLVM version 10.0.0 (clang-1000.11.45.5)
Target: x86_64-apple-darwin18.2.0
Thread model: posix
InstalledDir: /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin

ハードの情報

  機種名:	MacBook Pro
  機種ID:	MacBookPro13,2
  プロセッサ名:	Intel Core i5
  プロセッサ速度:	2.9 GHz
  プロセッサの個数:	1
  コアの総数:	2
  二次キャッシュ（コア単位）:	256 KB
  三次キャッシュ:	4 MB
  メモリ:	16 GB