パック展開でリスト操作

先週はパック展開を使うとconstexpr関数が高速化するみたいなことを書きましたが、具体的な使い方をあまり書けなかったので、具体例を書いておきます。

map

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto map_f_impl( const T (&arr)[N], std::index_sequence<Indeces...> ) {
    return std::array { f( arr[Indeces] )... };
}
template<class T, std::size_t N>
constexpr auto map_f( const T (&arr)[N] ) {
    return map_f_impl( arr, std::make_index_sequence<N>() );
}

パック展開をそのまま使えばよいです。 C++17では推論ガイドのおかげでstd::arrayのテンプレート実引数を指定しなくてよくなりました。

fold

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto fold_impl( const T (&arr)[N], std::index_sequence<Indeces...> ) {
    return (arr[Indeces] + ...);
}
template<class T, std::size_t N>
constexpr auto fold( const T (&arr)[N] ) {
    return fold_impl( arr, std::make_index_sequence<N>() );
}

畳み込み式をそのまま使えばよいです。

scan (prefix sum)

template<class T>
constexpr auto sum( T& acc, T val ) {
    return acc += val;
}

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto prefix_sum_impl( const T (&arr)[N], T init, std::index_sequence<Indeces...> ) {
    return std::array { sum( init, arr[Indeces] )... };
}
template<class T, std::size_t N>
constexpr auto prefix_sum( const T (&arr)[N], T init = T() ) {
    return prefix_sum_impl( arr, init, std::make_index_sequence<N>() );
}

副作用を持つ関数を使ってパック展開すると、scanが作れます。あまり知られていないような気がします。

実践例

最小値の検索

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto min_impl( const T (&arr)[N], T init, std::index_sequence<Indeces...> ) {
    return ((init = std::min( init, arr[Indeces]), ...);
}
template<class T, std::size_t N>
constexpr auto min( const T (&arr)[N], T init ) {
    return min_impl( arr, init, std::make_index_sequence<N>() );
}

最小値の探索は一種の畳み込みであると考えることができます。 コンマ演算子で畳み込みすることでfor文の代わりにすることができます。 畳み込み式が使えないC++14の場合でも、Swallowの技法を使っても同じことができます。

行列乗算

template<class T, std::size_t N>
struct Matrix { 
    T data[N*N];
};

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto inner_product( const T (&lhs)[N*N], const T (&rhs)[N*N], std::size_t i, std::size_t k, std::index_sequence<Indeces...> ) {
    return ((lhs[i*N + Indeces] * rhs[Indeces*N + k]) + ...);
}

template<class T, std::size_t N, std::size_t... Indeces>
constexpr auto product_impl( const Matrix<T, N>& lhs, const Matrix<T, N>& rhs, std::index_sequence<Indeces...> ) {
    return Matrix<T, N> { inner_product<T, N>( lhs.data, rhs.data, Indeces / N, Indeces % N, std::make_index_sequence<N>() )... };
}

template<class T, std::size_t N>
constexpr auto product( const Matrix<T, N>& lhs, const Matrix<T, N>& rhs ) {
    return product_impl( lhs, rhs, std::make_index_sequence<4>() );
}

同時に複数のパック展開を行うことはできないので、関数ごとに分けて行います。 constexpr でニューラルネットワークを実装したというのが最近(6月)ありましたが、そこで使われていた行列乗算をこれに置き換えたところ8倍程度高速化した記憶があります。 前回説明した、

  • パック展開で初期化する
  • 生配列を使って関数呼び出しを減らす
  • 畳み込み式を使って変数の書き換えを減らす

をすべて使っています。