Expression templateとfmaについて

Fused Multiply Addという命令がある。これは、a * b + cという形の演算を1命令で処理するものだ。丸めが一度しか走らないので精度がよくなる。以前、理論値と数値計算を比較してテストするときにこいつが使われるかどうかでテストが通ったり落ちたりして困った思い出がある。

さて、必ずこれを使うようにするにはどうしたらよいか。簡単だ。C99ではmath.hfmafmaffmalが定義されており、この関数を使えばよい。C規格(N1256)の§F.9.10.1 "The fma functions"で、

fma(x,y,z) computes xy + z, correctly rounded once.

とのことなので、CPUがサポートしていなかったらソフトウェア側でエミュレーションするのだろう。実際FP_FAST_FMAというマクロがあり(§7.12 Mathematics <math.h>)、これは積と和を計算するよりもFMAの方が早い場合のみ定義される。FMAFとFMAL版もある。これはほぼ専用命令があれば定義されることと同義なので、これが定義されているかどうかで使い分けもできるだろう。もしCPUがFPUを持っておらず、浮動小数点数演算を全てソフトウェアエミュレーションしていれば(そしてエミュレーションの実装でFMAの方が速ければ)、FMA専用命令がなくとも定義されるかもしれない。いまどきそんなCPUあるか?

さて、C++C++11からstd::fmaがサポートされた。これはCのfmaと同等である。また、SIMDFMAでは以下のようなイントリンシックがある。

__m128 _mm_fmadd_ps(__m128 a, __m128 b, __m128 c);

FMA関係のイントリンシックはもちろんこの一つではなく、幅(__m128, __m256) かける 精度(_ps_pd) 足す 非SIMD_ss_sd) で10通りある。AVX512ではマスクがあるのでもっと色々ある。

さて、これを制御したいとする。ライブラリ側で、ユーザーには気づかれないうちに。

目的からして、ユーザーがどんな演算をするか何も仮定できない。とすると、どういう演算をするかという情報を取り出せる場所は一つしかない。ユーザーのコードだ。

先に書いておいたライブラリがユーザーのコードから情報を取り出す方法としては、まずExpression Template(ET)が思いつく。これは本来はBlitz++という線形代数ライブラリが必要のない要素の計算を避けるために開発した技法で、一種の遅延評価である。実装をかいつまんで書くと、まず以下のようなクラスを作って、演算子が計算結果の一次オブジェクトでなく、両辺への参照を持った特殊なクラスを返すようにする。

template<typename L, typename R>
struct added
{
    using value_type = decltype(
        std::declval<typename L::value_type>() +
        std::declval<typename R::value_type>()
    );
    value_type operator[](std::size_t i) const noexcept
    {
        return l[i] + r[i]};
    }

    L const& l;
    R const& r;
};

このクラスは2つの型(左辺のオブジェクトと右辺のオブジェクトになる)をとって、アクセスしたときに必要な値を計算して返すクラスだ。このクラスはアクセスされるまでは計算を行わない。よって、例えば巨大な行列の積を計算するが対角要素以外いらないとか、そういう時のために使われてきた。

さて、これを使うと、以下の式のdの型はどうなるだろうか。それぞれの変数はmatrix型だとする。

auto d = a * b + c;

まずa*bが行われるので、以下のようになる。

auto d = multiplied<matrix, matrix> + c;

その後、足し算が行われるので、以下のようになる。

auto d = added<multiplied<matrix, matrix>, matrix>;

というわけで、dの型はadded<multiplied<matrix, matrix>, matrix>だ。これは木構造になっている。というか、抽象構文木(AST)になっているのだ。

本来は実装の手間を省くために、std::plusなどを取る型パラメータを付け加えるためよりASTっぽくなるが、今回は例が複雑になるためやめた。

このように、ETを使うことでユーザーコードで出てくる変数の型をASTにすることができる。ここで、templateは部分特殊化、つまりパターンマッチが使えるということを思い出そう。すると、以下のような操作ができる。

template<typename L, typename M, typename R>
struct added<multiplied<L, M>, R>
{
    using value_type = decltype(
        std::declval<typename L::value_type>() *
        std::declval<typename M::value_type>() +
        std::declval<typename R::value_type>()
    );
    value_type operator[](std::size_t i) const noexcept
    {
        return std::fma(l[i], m[i], r[i]};
    }

    L const& l;
    M const& m;
    R const& r;
}:

これは、addedの左辺が何かしらのmultipliedだった場合への特殊化だ。そういうクラスが現れた時は、通常のaddedmultipliedではなくこちらで実体化される。

――という風にすれば、ユーザーが意識せずに書いてもfmaが勝手に生成されるという風にできるなあ、と思いつつも実装していない。面倒なのである。

だが、ライブラリ側でユーザーコードのASTを変換できる可能性というのは面白いのではないかなあと思う。コンパイラが最適化のためにいじっているのも何らかの中間表現なので、理論上は同等のものが実装できるはずだ。最近はETはコンパイラの自動SIMD化を阻害する(解析が難しくなる)ので下火という話も聞くが、こういう方向の発展もありなのではないか。

……それはコンパイラの最適化を阻害しつつ、自分でコンパイラがやってるのと同じ最適化をするってことだから、めっちゃくちゃ特殊な場合を除いて意味ないんじゃないの。