C
assembly
BLAS
HPC

密行列積演算高速化と科学計算ライブラリについて

これは eeic Advent Calender 2017 (https://qiita.com/advent-calendar/2017/eeic) の3日目の記事です。

はじめに

この記事を読んでくださりありがとうございます。

電子情報工学科に所属しておりますhigucheeseと申します。

突然ですが皆さんは、
お手持ちのラップトップPCに積まれているCPUがどれほど利用されているか、
考えたことはありますでしょうか。

メジャーなOSのデスクトップ環境ではGUIでPCのパフォーマンスを表示(WindowsならTask Manager, OS XならActivity Monitor)出来るため、
見たことがある人も多いと思います。
他にアプリケーションを起動していなければCPU稼働率が数%ほど、
起動中や稼働中では10%overほどになるのが普通でしょう。
処理が重たいアプリケーションであれば一時的に50%を超えることもありますが、
常に100%付近ということはまずないでしょう。

このように十分に発達したOSがインストールされた環境では、
無数のアプリケーション(タスク)が同時に存在するのにも関わらず、
それらによってCPU利用が競合することもなく共存が出来ています。

なぜこのようなことが可能であるかというと、
OSとCPUが協力して

  • 計算に使うデータがない(データ読み込み・ダウンロード中など)
  • CPUに与える命令がない
  • 他のシグナルを待っている(人間の入力待ちなど)

時はCPUが仕事をしないようにしているからです。
OSの基本的な機能の一つです。

そしていよいよここからが本題ですが、

  • 計算に使うデータがある
  • CPUに与える命令がある
  • 待つべき他のシグナルがない

更に

  • 極めて大量の計算が必要となる

場合は、
プログラマはアプリケーションが存在するCPU資源を使い切りたいはずです。
そのときに利用されるのが科学計算ライブラリです。

科学計算ライブラリはその種類も多岐に渡りますが、
ここでは線形代数の実装を最適化したものを指すことにします。
BLAS(Basic Linear Algebra Subprograms)に代表され、
LAPACK(BLASの拡張),Eigen,Intel MKL,cuBLAS,OpenBLASなど様々な実装がされています。

これらに共通することは、

  • プロセッサのアーキテクチャレベルでの最適化が行われている
  • 様々なデータサイズや計算係数に対応している
  • 浮動小数点型演算誤差のテストが行われている

ということです。
それらがどのようなものなのかを実例も交えつつ説明していきます。

BLASライブラリと勝負する

元祖BLASが持つ機能は少なく以下の三段階に分かれています。
(URL: http://www.netlib.org/blas/#_blas_routines)

  • Level1: ベクトルの処理またはベクトル同士の演算
  • Level2: ベクトルと行列同士の演算
  • Level3: 行列同士の演算

今回は勝負する題材として密行列同士の演算を選びます。
より詳しくルールを説明しますと、

  • 行列の要素に一切の偏りを仮定しない
  • 行列の精度は単精度(c/c++ではfloat型)
  • 精度保証は適当に行う

となります。
計算式で書き表すと

C = A \times B

です。
最初のルールがわかりにくいので補足すると、
計算に用いる行列が疎行列(要素の大部分が0)/三角行列/対称行列ではないということです。

密行列積演算はHPC(High Performance Computing)の分野でベンチマークの一種として用いられます。
例としてTOP500(URL: https://www.top500.org)というスーパーコンピューターのランキングがあり、
この順位付けは密行列積演算によって算出されるFLOPS(Floating-point Operations Per Second)で行われます。
ちなみに東大と筑波大が共同で管理する「Oakforest-PACS」というスパコンシステムは記事投稿時(2017/12/3)9位です。

これがベンチマークとして利用される理由は、
行列積はデータ量に対する計算量(計算密度)が高く計算資源をより有効に使える上、
画像処理/データ解析/ニューラルネットワークなどの様々な分野で出てくるためだと自分は理解しています。

測定環境

プログラムの実行・計測には上のOakforest-PACSシステムを利用しました。
以下は簡単なシステムの構成です。

  • プロセッサ名: Intel XeonPhi 7250
  • 理論性能(単精度): 約6TFLOPS
  • メモリ:96GB(DDR4)+16GB(MCDRAM Cache Mode)
  • OS: CentOS 7
  • コンパイラ: icpc 17.0.4

コンパイルオプション
-std=c++11 -xMIC-AVX512 -O3 -qopenmp

プログラムの並列化はスレッド単位で行い、
プロセスレベルの並列化は今回は取り上げません。

ベースライン

それではここから実際に行列積のコードを書いていきたいと思います。
言語はC++を使いますが内部はすべてC言語です。

まず肝心な行列の形状ですが、
コードの改造を簡単にするため$A,B,C$すべてを$N\times N$の正方行列とします。
次に計算回数$FLO$(Floating-point Operations)を以下の計算式で求めます。

FLO = 2 \times N^3

この$FLO$を計算にかかった時間で割れば$FLOPS(Floating-point Operations Per Second)$が算出できます。

それではまずは最も数式に忠実なコードを書いてみます。

void mm_square(float* A, float* B, float* C, int N) {
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            for (int k = 0; k < N; k++) {
                C[i*N+j] += A[i*N+k] * B[k*N+j];
            }
        }
    }
}

N=128として計算させるとその性能は..

# Flops:          0.43759 [GFLOPS]

低い..理論性能の1万分の1も出ていませんね。
それに対してIntelの提供するBLASライブラリ(Intel MKL2017)を用いると、

void MulMat::multiply() {
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,N,N,N,alpha,A,LDA,B,LDB,beta,C,LDC);
}

このような実装になりN=26800として性能を測定すると、

# Flops:   4740.67 [GFLOPS]

理論性能の80%ほど出せていることがわかります。
最も単純な実装と比較すると1万倍の差があることもわかります。
この性能が今回の目標とするところになります。

プログラムの並列化

マルチコアであるプロセッサの性能を最大限引き出すには必ず並列化が必要です。
ここで少し並列プログラムをの書き方を知っている人はOpenMPというライブラリを知っていると思います。
OpenMPの説明はこちらのサイト(http://www.rccm.co.jp/development/parallel/openmp.html) に日本語で簡単に書いてあるので見てみてください。
要はfor文の中で行う計算を自動的に生成したスレッドに分散して行うということです。

早速pragmaディレクティブを追加して並列化してみます。(スレッド数68, スレッドの配置は各物理コア割当)

void mm_square(float* A, float* B, float* C, int N) {
#pragma omp parallel for
    for (int i = 0; i < N; i++) {

同じN=128で

# Flops:          0.101562 [GFLOPS]

性能が下がりました。
実はこのような性能劣化は並列化を行うプログラマにとっては珍しいことでも何でもなく、
並列化に伴うオーバーヘッドが悪さを行う場面は多々あります。

プログラムを並列に実行するには、

  • スレッドの生成
  • データ(とプログラム)の各コアへの転送
  • スレッドの削除

が必要ですが今回はこの2番目の要素が問題となっていると思われます。
データが小さく転送コストをかけるくらいならコア1つでやったほうが速い、
といったことが起こっているわけです。

キャッシュの階層構造

先程は簡単に並列化をしてみて性能劣化が起こりましたが、
一旦並列化をやめて元のプログラムでNのサイズを変化させてみます。
(単位はGFLOPS)

N=128 N=256 N=512 N=1024
0.438 0.449 0.115 0.081

なぜこのような性能差が発生するのかはキャッシュの階層構造で説明が可能です。
一般的にCPUとメモリの間にはデータ転送の遅延を隠蔽するために、
キャッシュと呼ばれるメモリの一時コピー用記憶装置があります。
今回の実験に使ったシステムではL1,L2,L3キャッシュ(MCDRAM)が実装されており、
アクセス速度はレジスタ>L1>L2>L3=メモリとなっています。(一般的にはL3>メモリ)
レジスタはその容量が極めて小さいので考えないとして、
それぞれの1物理コアあたり利用可能なキャッシュの容量は、

L1 Cache L2 Cache L3 Cache
32KiB 1MiB/2 16GiB/68

L2キャッシュは隣接する2コアで共有しています。

行列積はデータの再利用が大量に行われる演算なので、
一度利用したデータをキャッシュに入れて再利用ができると性能が向上します。

Nの各サイズで利用するデータ量はそれぞれ、

N=128 N=256 N=512 N=1024
192KiB 768KiB 3MiB 12MiB

となりN=512からL2キャッシュまでに乗り切らないことになります。
キャッシュに乗り切らないということは演算の途中で頻繁にメモリ(またはL3キャッシュ)にアクセスするため、
上で確認できたように性能の劣化が起こります。

行列積のタイリング

ここまでで

  • 並列化で良い結果を得るにはデータサイズを大きくしたい
  • 行列積で良い結果を得るにはデータサイズを大きくしたくない

という矛盾がすでに発生していますが、
これに対する答えの一つが行列のタイリングとなります。
一つの行列を複数の行列に分割して計算することで、
データのサイズを大きくしつつ計算に必要なデータを再利用することができます。

A = [a_{ij}]_{m \times k},B = [b_{ij}]_{k \times n}

であるとき、

  A = \left[
    \begin{array}{rrrr}
      A_{11} & A_{12} & \dots & A_{1k_t}\\
      A_{21} & A_{22} &  & A_{2k_t}\\
      \vdots &  & \ddots & \vdots\\
      A_{m_t1} & A_{m_t2} & \dots & A_{m_tk_t}\\
    \end{array}
  \right]
,
  B = \left[
    \begin{array}{rrrr}
      B_{11} & B_{12} & \dots & B_{1n_t}\\
      B_{21} & B_{22} &  & B_{2n_t}\\
      \vdots &  & \ddots & \vdots\\
      B_{k_t1} & B_{k_t2} & \dots & B_{k_tn_t}\\
    \end{array}
  \right]
,
C_{ij} = \sum^{k_t}_k A_{i k} \times B_{k j}

この演算の組み換えをキャッシュの階層構造に合わせて何段かに分けて行います。
以下はその実装です。

const int AVX512Fs = 16; // 512bit/32bit(float)=16
const int sssubN = 32;
const int ssubN = 128;
const int subN = 1024;

inline
void mm_dep3(float* A, float* B, float* C, int N) {
    for (int i = 0; i < sssubN; i++) {
        for (int j = 0; j < sssubN; j++) {
            for (int k = 0; k < sssubN; k++) {
                C[i*N+j] += A[i*N+k] * B[k*N+j];
            }
        }
    }
}

void mm_dep2(float* A, float* B, float* C, int N) {
    for (int i = 0; i < ssubN; i += sssubN) {
        for (int j = 0; j < ssubN; j += sssubN) {
            for (int k = 0; k < ssubN; k += sssubN) {
                float* Ai = &A[i*N+k];
                float* Bi = &B[k*N+j];
                float* Ci = &C[i*N+j];
                mm_dep3(Ai, Bi, Ci, N);
            }
        }
    }
}

void mm_dep1(float* A, float* B, float* C, int N) {
    for (int i = 0; i < subN; i += ssubN) {
        for (int j = 0; j < subN; j += ssubN) {
            for (int k = 0; k < subN; k += ssubN) {
                float* Ai = &A[i*N+k];
                float* Bi = &B[k*N+j];
                float* Ci = &C[i*N+j];
                mm_dep2(Ai, Bi, Ci, N);
            }
        }
    }
}

void mm_square(float* A, float* B, float*C, int N) {
    for (int i = 0; i < N; i += subN) {
        for (int j = 0; j < N; j += subN) {
            for (int k = 0; k < N; k += subN) {
                float* Ai = &A[i*N+k];
                float* Bi = &B[k*N+j];
                float* Ci = &C[i*N+j];
                mm_dep1(Ai, Bi, Ci, N);
            }
        }
    }
}

計算の順序は変わっていますが、
最内ループで行われる演算は全く同じです。
このプログラムをN=1024で走らせてみると、

# Flops:          0.10535 [GFLOPS]

このような結果になります。
なかなかに努力が報われないような結果ですがこの理由を示すのは難しいです。
律速となるのがデータフローではなく計算部分であるときこのような結果になると思われます。
ですが理論的には確実に性能向上に近づいているはずです。

命令間の依存関係

ここまでで行列のタイリングが行われ、
行列の計算本体は最内ループのみに注目することができるようになりました。

inline
void mm_dep3(float* A, float* B, float* C, int N) {
    for (int i = 0; i < sssubN; i++) {
        for (int j = 0; j < sssubN; j++) {
            for (int k = 0; k < sssubN; k++) {
                C[i*N+j] += A[i*N+k] * B[k*N+j];
            }
        }
    }
}

この最内ループの更に内側のループに着目します。

for (int k = 0; k < sssubN; k++) {
    C[i*N+j] += A[i*N+k] * B[k*N+j];
}

この最内forループの問題点は、

C[i*N+j] += A[i*N+k] * B[k*N+j];

の実行が終わらないと

C[i*N+j] += A[i*N+(k+1)] * B[(k+1)*N+j];

の計算が始められないという点です。
これを「命令間に依存関係が存在する」と表現します。
この場合は2番目の計算が1番目の計算結果に依存しています。

連続した命令間に依存関係が存在しない場合は、
プロセッサは命令の終了を待たず次の命令を開始できます。

次のように少し変更を加えるだけで依存関係のある命令を遠ざけることができます。

for (int k = 0; k < sssubN; k++) {
    for (int j = 0; j < sssubN; j++) {
        C[i*N+j] += A[i*N+k] * B[k*N+j];
    }
}

通常計算には複数クロックかかりますが計算ユニットは同時に複数の計算を実行できるため、
このように書けばパイプライン処理が働き計算ユニットの全体を活用できるでしょう。

同様にN=1024で計測してみます。

# Flops:          4.51573 [GFLOPS]

先程の結果と比較して45倍となりました。
(実際には依存関係が妨げていた自動ベクトル化が行われたということも速度向上の理由と思われます。)

ベクトル演算

今時のCPUにはベクトル演算用のユニットが搭載されています。
計測で用いるXeonPhiはAVX512命令用のベクトル演算ユニットが搭載されています。
Intelの提供する内で最新のベクトル命令セットです。

AVX512は同時に512bit幅のベクトル演算処理が可能です。
ソースコードで例えると、

float A[16], B[16], C[16];
for (int i = 0; i < 16; i++) {
    C[i] += A[i] * B[i];
}

この演算がFMADD(Fused Multiply-ADD)というベクトル演算1回で行えます。
もちろん演算は6クロック(https://software.intel.com/sites/landingpage/IntrinsicsGuide) かかってしまいますが、
パイプライン処理で隠蔽します。

組み込み関数

コンパイラには自動ベクトル化機能があります。
上記のプログラムをobjdump -dで逆アセンブルして確認してみました。
ここが最内関数内の最内ループの一部です。

     16a:   62 b2 7d 48 18 14 9a    vbroadcastss (%rdx,%r11,4),%zmm2
     171:   62 91 7c 48 10 0c ba    vmovups (%r10,%r15,4),%zmm1
     178:   62 91 7c 48 10 44 ba    vmovups 0x40(%r10,%r15,4),%zmm0
     17f:   01 
     180:   62 f2 6d 48 a8 0b       vfmadd213ps (%rbx),%zmm2,%zmm1
     186:   62 f1 7c 48 11 0b       vmovups %zmm1,(%rbx)
     18c:   62 f2 7d 48 a8 53 01    vfmadd213ps 0x40(%rbx),%zmm0,%zmm2
     193:   62 f1 7c 48 11 53 01    vmovups %zmm2,0x40(%rbx)

vfmadd213psという命令が密集しているほどパイプライン処理が効き、
ループアンローリングもうまくいっていると言えるのですが、
このプログラムではあまりうまくいっていないようです。
%zmm?というのが512bit幅をもつベクトルレジスタで合計32本あります。
ただし通常の自動ベクトル化ではそれをフル活用させるのは難しいです。

自動ベクトル化がどのように難しいかについてはここでは説明できません。
しかし自動ベクトル化では自分が求める、
ベクトルレジスタとベクトル命令を使い切るようなバイナリは出力してくれないようです。
このような場合にコンパイラに無理やりベクトル演算命令を利用してもらうために、
組み込み関数を利用します。

inline
void mm_dep3(float* A, float* B, __m512* vecC, int N) {
    __m512 vecB;
    for (int k = 0; k < 16; k++) {
        vecB = _mm512_load_ps(&B[k*N]);
        _mm_prefetch((const char*)&B[(k+1)*N], _MM_HINT_T0);
#define CALC(i) _mm512_fmadd_ps(_mm512_set1_ps(A[i*N+k]), vecB, vecC[i])
        vecC[0] = CALC(0);
        vecC[1] = CALC(1);
        vecC[2] = CALC(2);
        vecC[3] = CALC(3);
        vecC[4] = CALC(4);
        vecC[5] = CALC(5);
        vecC[6] = CALC(6);
        vecC[7] = CALC(7);
        vecC[8] = CALC(8);
        vecC[9] = CALC(9);
        vecC[10] = CALC(10);
        vecC[11] = CALC(11);
        vecC[12] = CALC(12);
        vecC[13] = CALC(13);
        vecC[14] = CALC(14);
        vecC[15] = CALC(15);
    }
}

void mm_dep2(float* A, float* B, float* C, int N) {
    for (int i = 0; i < ssubN; i += sssubN) {
        for (int j = 0; j < ssubN; j += sssubN) {
            _mm_prefetch((const char*)&B[0], _MM_HINT_T0);
            __m512 vecC[16];
            for (int t = 0; t < 16; t++)
                vecC[t] = _mm512_load_ps(&C[(i+t)*N+j]);
            for (int k = 0; k < ssubN; k += sssubN) {
                float* Ai = &A[i*N+k];
                float* Bi = &B[k*N+j];
                mm_dep3(Ai, Bi, vecC, N);
            }
            for (int t = 0; t < 16; t++)
                _mm512_store_ps(&C[(i+t)*N+j], vecC[t]);
        }
    }
}

__mm512型で直接ベクトルレジスタを表現でき、
_mm512から始まる組み込み関数は直接特定の命令を使うよう支持しています。
適度にソフトウェアプリフェッチも利用して少しでもキャッシュ利用率を上げています。
処理の一部がmm_dep2にも侵食していますが、
最内関数のC(16x16)をすべてベクトルレジスタに載せることで、
組み込み関数を利用した行列積記述を簡単にしています。

性能はN=1024で

# Flops:          10.4587 [GFLOPS]

組み込み関数を利用する前と比べて2倍ほどになりました。
早速バイナリの一部を見てみます。

    1753:   62 52 0d 58 b8 0c 91    vfmadd231ps (%r9,%rdx,4){1to16},%zmm14,%zmm9
    175a:   62 52 0d 58 b8 24 97    vfmadd231ps (%r15,%rdx,4){1to16},%zmm14,%zmm12
    1761:   62 52 0d 58 b8 04 94    vfmadd231ps (%r12,%rdx,4){1to16},%zmm14,%zmm8
    1768:   62 d2 0d 58 b8 3c 93    vfmadd231ps (%r11,%rdx,4){1to16},%zmm14,%zmm7
    176f:   62 f2 0d 58 b8 34 91    vfmadd231ps (%rcx,%rdx,4){1to16},%zmm14,%zmm6
    1776:   62 f2 0d 58 b8 24 93    vfmadd231ps (%rbx,%rdx,4){1to16},%zmm14,%zmm4
    177d:   62 f2 0d 58 b8 1c 96    vfmadd231ps (%rsi,%rdx,4){1to16},%zmm14,%zmm3
    1784:   62 f2 0d 58 b8 14 97    vfmadd231ps (%rdi,%rdx,4){1to16},%zmm14,%zmm2
    178b:   62 d2 0d 58 b8 0c 90    vfmadd231ps (%r8,%rdx,4){1to16},%zmm14,%zmm1
    1792:   62 52 0d 58 b8 2c 92    vfmadd231ps (%r10,%rdx,4){1to16},%zmm14,%zmm13

先程と比べればマシに見えます。
しかし注目するべきはvfmadd213psの第一引数です。
1753アドレスの例を使うと、
これは4倍した%r9を%rdxに足した位置の値が16個並んだものと扱え、
という意味ですが各命令ごとに利用されるレジスタが違います。

これは行列を細かく分割してもデータの位置は変わらないため、
コンパイル時に決定不能なNを利用してインデックスを計算しなくてはならなくなり、
結果として大量の余分な計算とレジスタの利用が発生していることを意味しています。

プリフェッチ

CPUにはメモリのアクセスパターンから推測して次にアクセスされるであろうデータを事前にキャッシュに保存する機能があります。
それをハードウェアプリフェッチと呼びます。
ハードウェアプリフェッチを利用することでCPUから特別な命令が発行される、
またはページフォールトを起こすことなしにキャッシュにデータをロードできます。
ハードウェアプリフェッチの中でも連続アクセスに対応するものを利用します。

例として次の2つのメモリアクセスパターンの異なるプログラムを走らせて速度を測ってみます。

#include <stdlib.h>
int main(void) {
    long N = 50000;
    float* A = (float*)malloc((N*N)*sizeof(float));
    for (long i = 0; i < N; i++) {
        for (long j = 0; j < N; j++) {
            volatile float p = A[i*N+j];
        }
    }
    return 0;
}
#include <stdlib.h>
int main(void) {
    long N = 50000;
    float* A = (float*)malloc((N*N)*sizeof(float));
    for (long j = 0; j < N; j++) {
        for (long i = 0; i < N; i++) {
            volatile float p = A[i*N+j];
        }
    }
    return 0;
}

結果は

データに対して連続にアクセスするプログラム

% time ./a.out
./a.out 1.53s user 0.01s system 99% cpu 1.541 total

データに対してNごとにアクセスするプログラム

% time ./a.out
./a.out 5.46s user 0.01s system 99% cpu 5.465 total

前者のプログラムのほうが素早くデータにアクセスできていることがわかります。
パフォーマンスチューニングではデータを演算ユニットに如何に効率よく流し込むかが問題であるので、
なるべく連続にアクセスできるようなデータ構造が良いと言えます。

並べ替え

データにアクセスする順番が大事であると述べました。
これを行列積演算に応用するためには行列内部の並べ替えが必要になります。

A,B,Cに対応する変数をtA,tB,tCとして、
A,B,Cにアクセスする順番を考慮して並べ替えます。

void reorder_A(float* A, float* tA, int N) {
    long indextA = 0;

//#pragma omp parallel for collapse(2)
    for (int i = 0; i < N; i += subN) {
        for (int j = 0; j < N; j += subN) {
            for (int ii = 0; ii < subN; ii += ssubN) {
                for (int jj = 0; jj < subN; jj += ssubN) {
                    for (int iii = 0; iii < ssubN; iii += sssubN) {
                        for (int jjj = 0; jjj < ssubN; jjj += sssubN) {
                            float64 ttA[sssubN][sssubN];
                            long indexA = (long)(i + ii + iii) * N + (long)(j + jj + jjj);
                            for (int ti = 0; ti < sssubN; ti++) {
                                for (int tj = 0; tj < sssubN; tj++) {
                                    ttA[tj][ti] = A[indexA + ti * N + tj];
                                }
                            }
                            for (int t = 0; t < sssubN; t++) {
                                _mm512_store_ps(&tA[indextA], _mm512_load_ps(ttA[t]));
                                indextA += sssubN;
                            }
                        }
                    }
                }
            }
        }
    }
}

void reorder_B(float* B, float* tB, int N) {
    long indextB = 0;

//#pragma omp parallel for collapse(2)
    for (int i = 0; i < N; i += subN) {
        for (int j = 0; j < N; j += subN) {
            for (int ii = 0; ii < subN; ii += ssubN) {
                for (int jj = 0; jj < subN; jj += ssubN) {
                    for (int iii = 0; iii < ssubN; iii += sssubN) {
                        for (int jjj = 0; jjj < ssubN; jjj += sssubN) {
                            long indexB = (long)(i + ii + iii) + (long)(j + jj + jjj) * N;
                            for (int t = 0; t < sssubN; t++) {
                                _mm512_store_ps(&tB[indextB], _mm512_load_ps(&B[indexB]));
                                indexB += N;
                                indextB += sssubN;
                            }
                        }
                    }
                }
            }
        }
    }   
}

void reorder_C(float* C, float* tC, int N) {
    long indextC = 0;

//#pragma omp parallel for collapse(2)
    for (int i = 0; i < N; i += subN) {
        for (int j = 0; j < N; j += subN) {
            for (int ii = 0; ii < subN; ii += ssubN) {
                for (int jj = 0; jj < subN; jj += ssubN) {
                    for (int iii = 0; iii < ssubN; iii += sssubN) {
                        for (int jjj = 0; jjj < ssubN; jjj += sssubN) {
                            long indexC = (long)(i + ii + iii) * N + (long)(j + jj + jjj);
                            for (int t = 0; t < sssubN; t++) {
                                _mm512_store_ps(&C[indexC], _mm512_load_ps(&tC[indextC]));
                                indexC += N;
                                indextC += sssubN;
                            }
                        }
                    }
                }
            }
        }
    }
}

並べ替えにもベクトル演算は用いています。
時間計測には当然並べ替えの時間も含まれていますが、
並べ替えに必要な計算量はO(N^2)であるのに対し、
行列の計算はO(N^3)であるためNが大きいほど並べ替えにかかる時間は比較的小さくなります。

この並べ替えによって行列積本体のプログラムも変更されます。
大きな変化は行列データのアクセスにNを利用しなくても良くなったところです。

#include <cstdio>
#include <cstdlib>
#include <immintrin.h>
#include <sys/time.h>
#include <omp.h>

#include "mymulmat.h"

//#define DEBUG_MM

long get_us(void) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_usec + 1000000 * tv.tv_sec;
}

inline
void mm_dep3(float* A, float* B, float* C) {
#ifdef DEBUG_MM
    int indexA = 0;
    for (int k = 0; k < sssubN; k++) {
        for (int i = 0; i < sssubN; i++, indexA++) {
            for (int j = 0; j < sssubN; j++) {
                C[i*sssubN+j] += A[indexA] * B[k*sssubN+j];
            }
        }
    }
#else
    __m512 vecC[16];
    __m512 vecB;
    int indexA;
    for (int i = 0; i < 16; i++) {
        vecC[i] = _mm512_load_ps(&C[i*16]);
    }
#define FMADD(i) _mm512_fmadd_ps(_mm512_set1_ps(A[indexA+i]), vecB, vecC[i]);
#define CALC \
    vecC[0] = FMADD(0); \
    vecC[1] = FMADD(1); \
    vecC[2] = FMADD(2); \
    vecC[3] = FMADD(3); \
    vecC[4] = FMADD(4); \
    vecC[5] = FMADD(5); \
    vecC[6] = FMADD(6); \
    vecC[7] = FMADD(7); \
    vecC[8] = FMADD(8); \
    vecC[9] = FMADD(9); \
    vecC[10] = FMADD(10); \
    vecC[11] = FMADD(11); \
    vecC[12] = FMADD(12); \
    vecC[13] = FMADD(13); \
    vecC[14] = FMADD(14); \
    vecC[15] = FMADD(15)

    for (int k = 0; k < 16; k += 2) {
        indexA = k * 16;
        vecB = _mm512_load_ps(&B[k*16]);
        CALC;
        indexA = (k+1) * 16;
        vecB = _mm512_load_ps(&B[(k+1)*16]);
        CALC;
    }
    for (int i = 0; i < 16; i++) {
        _mm512_store_ps(&C[i*16], vecC[i]);
    }
#endif
}

void mm_dep2(float* A, float* B, float* C, int N) {
    long indexC = 0;
    for (int i = 0; i < S2; i++) {
        for (int j = 0; j < S2; j++) {
            long indexA = i * S2 * sssubN2;
            long indexB = j * S2 * sssubN2;
        float* Ci = &C[indexC];
            for (int k = 0; k < S2; k++) {
                float* Ai = &A[indexA];
                float* Bi = &B[indexB];
                mm_dep3(Ai, Bi, Ci);
                indexA += sssubN2;
                indexB += sssubN2;
            }
            indexC += sssubN2;
        }
    }
}

void mm_dep1(float* A, float* B, float* C, int N) {
    long indexC = 0;
    for (int i = 0; i < S1; i++) {
        for (int j = 0; j < S1; j++) {
            long indexA = i * S1 * ssubN2;
            long indexB = j * S1 * ssubN2;
            float* Ci = &C[indexC];
            for (int k = 0; k < S1; k++) {
                float* Ai = &A[indexA];
                float* Bi = &B[indexB];
                mm_dep2(Ai, Bi, Ci, N);
                indexA += ssubN2;
                indexB += ssubN2;
            }
        indexC += ssubN2;
        }
    }
}

void mm_square(float* A, float* B, float*C, int N, float* tA, float* tB, float* tC) {
    long t0 = get_us();

    reorder_A(A, tA, N);

    long t1 = get_us();

    reorder_B(B, tB, N);

    long t2 = get_us();

    const int S0 = N / subN;
    long indexC = 0;

//#pragma omp parallel for collapse(2)
    for (int i = 0; i < S0; i++) {
        for (int j = 0; j < S0; j++) {
            long indexA = i * S0 * subN2;
            long indexB = j * S0 * subN2;
            float* Ci = &tC[indexC];
            for (int k = 0; k < S0; k++) {
                float* Ai = &tA[indexA];
                float* Bi = &tB[indexB];
                mm_dep1(Ai, Bi, Ci, N);
                indexA += subN2;
                indexB += subN2;
            }
        indexC += subN2;
        }
    }

    long t3 = get_us();

    reorder_C(C, tC, N);

    long t4 = get_us();

    printf("%ld %ld %ld %ld\n", t1-t0, t2-t1, t3-t2, t4-t3);
}

ここまでやってようやくN=1024で

# Flops:          40.9958 [GFLOPS]

という性能が出せました。
もちろん1コアで出した記録です。
元のナイーブな実装と比較して100倍ほどの速度向上が達成できました。

早速アセンブリも見てみましょう。

  5f:   62 51 7c 48 10 3c b1    vmovups (%r9,%rsi,4),%zmm15
  66:   62 51 7c 48 10 74 b1    vmovups 0x40(%r9,%rsi,4),%zmm14
  6d:   01 
  6e:   49 89 cb                mov    %rcx,%r11
  71:   49 89 c2                mov    %rax,%r10
  74:   32 db                   xor    %bl,%bl
  76:   62 51 7c 48 10 6c b1    vmovups 0x80(%r9,%rsi,4),%zmm13
  7d:   02 
  7e:   62 51 7c 48 10 64 b1    vmovups 0xc0(%r9,%rsi,4),%zmm12
  85:   03 
  86:   62 51 7c 48 10 5c b1    vmovups 0x100(%r9,%rsi,4),%zmm11
  8d:   04 
  8e:   62 51 7c 48 10 54 b1    vmovups 0x140(%r9,%rsi,4),%zmm10
  95:   05 
  96:   62 51 7c 48 10 4c b1    vmovups 0x180(%r9,%rsi,4),%zmm9
  9d:   06 
  9e:   62 51 7c 48 10 44 b1    vmovups 0x1c0(%r9,%rsi,4),%zmm8
  a5:   07 
  a6:   62 d1 7c 48 10 7c b1    vmovups 0x200(%r9,%rsi,4),%zmm7
  ad:   08 
  ae:   62 d1 7c 48 10 74 b1    vmovups 0x240(%r9,%rsi,4),%zmm6
  b5:   09 
  b6:   62 d1 7c 48 10 6c b1    vmovups 0x280(%r9,%rsi,4),%zmm5
  bd:   0a 
  be:   62 d1 7c 48 10 64 b1    vmovups 0x2c0(%r9,%rsi,4),%zmm4
  c5:   0b 
  c6:   62 d1 7c 48 10 5c b1    vmovups 0x300(%r9,%rsi,4),%zmm3
  cd:   0c 
  ce:   62 d1 7c 48 10 54 b1    vmovups 0x340(%r9,%rsi,4),%zmm2
  d5:   0d 
  d6:   62 d1 7c 48 10 4c b1    vmovups 0x380(%r9,%rsi,4),%zmm1
  dd:   0e 
  de:   62 d1 7c 48 10 44 b1    vmovups 0x3c0(%r9,%rsi,4),%zmm0
  e5:   0f 
  e6:   45 32 ff                xor    %r15b,%r15b
  e9:   4f 8d 34 90             lea    (%r8,%r10,4),%r14
  ed:   4e 8d 2c 9f             lea    (%rdi,%r11,4),%r13
  f1:   45 33 e4                xor    %r12d,%r12d
  f4:   62 81 7c 48 10 04 a6    vmovups (%r14,%r12,4),%zmm16
  fb:   62 81 7c 48 10 4c a6    vmovups 0x40(%r14,%r12,4),%zmm17
 102:   01 
 103:   41 80 c7 01             add    $0x1,%r15b
 107:   62 12 7d 50 b8 7c a5    vfmadd231ps 0x0(%r13,%r12,4){1to16},%zmm16,%zmm15
 10e:   00 
 10f:   62 12 7d 50 b8 74 a5    vfmadd231ps 0x4(%r13,%r12,4){1to16},%zmm16,%zmm14
 116:   01 
 117:   62 12 7d 50 b8 6c a5    vfmadd231ps 0x8(%r13,%r12,4){1to16},%zmm16,%zmm13
 11e:   02 
 11f:   62 12 7d 50 b8 64 a5    vfmadd231ps 0xc(%r13,%r12,4){1to16},%zmm16,%zmm12
 126:   03 
 127:   62 12 7d 50 b8 5c a5    vfmadd231ps 0x10(%r13,%r12,4){1to16},%zmm16,%zmm11
 12e:   04 
 12f:   62 12 7d 50 b8 54 a5    vfmadd231ps 0x14(%r13,%r12,4){1to16},%zmm16,%zmm10
 136:   05 
 137:   62 12 7d 50 b8 4c a5    vfmadd231ps 0x18(%r13,%r12,4){1to16},%zmm16,%zmm9
 13e:   06 
 13f:   62 12 7d 50 b8 44 a5    vfmadd231ps 0x1c(%r13,%r12,4){1to16},%zmm16,%zmm8
 146:   07 
 147:   62 92 7d 50 b8 7c a5    vfmadd231ps 0x20(%r13,%r12,4){1to16},%zmm16,%zmm7
 14e:   08 
 14f:   62 92 7d 50 b8 74 a5    vfmadd231ps 0x24(%r13,%r12,4){1to16},%zmm16,%zmm6
 156:   09 
 157:   62 92 7d 50 b8 6c a5    vfmadd231ps 0x28(%r13,%r12,4){1to16},%zmm16,%zmm5
 15e:   0a 
 15f:   62 92 7d 50 b8 64 a5    vfmadd231ps 0x2c(%r13,%r12,4){1to16},%zmm16,%zmm4
 166:   0b 
 167:   62 92 7d 50 b8 5c a5    vfmadd231ps 0x30(%r13,%r12,4){1to16},%zmm16,%zmm3
 16e:   0c 
 16f:   62 92 7d 50 b8 54 a5    vfmadd231ps 0x34(%r13,%r12,4){1to16},%zmm16,%zmm2
 176:   0d 
 177:   62 92 7d 50 b8 4c a5    vfmadd231ps 0x38(%r13,%r12,4){1to16},%zmm16,%zmm1
 17e:   0e 
 17f:   62 92 7d 50 b8 44 a5    vfmadd231ps 0x3c(%r13,%r12,4){1to16},%zmm16,%zmm0
 186:   0f 
 187:   62 12 75 50 b8 7c a5    vfmadd231ps 0x40(%r13,%r12,4){1to16},%zmm17,%zmm15
 18e:   10 
 18f:   62 12 75 50 b8 74 a5    vfmadd231ps 0x44(%r13,%r12,4){1to16},%zmm17,%zmm14
 196:   11 
 197:   62 12 75 50 b8 6c a5    vfmadd231ps 0x48(%r13,%r12,4){1to16},%zmm17,%zmm13
 19e:   12 
 19f:   62 12 75 50 b8 64 a5    vfmadd231ps 0x4c(%r13,%r12,4){1to16},%zmm17,%zmm12
 1a6:   13 
 1a7:   62 12 75 50 b8 5c a5    vfmadd231ps 0x50(%r13,%r12,4){1to16},%zmm17,%zmm11
 1ae:   14 
 1af:   62 12 75 50 b8 54 a5    vfmadd231ps 0x54(%r13,%r12,4){1to16},%zmm17,%zmm10
 1b6:   15 
 1b7:   62 12 75 50 b8 4c a5    vfmadd231ps 0x58(%r13,%r12,4){1to16},%zmm17,%zmm9
 1be:   16 
 1bf:   62 12 75 50 b8 44 a5    vfmadd231ps 0x5c(%r13,%r12,4){1to16},%zmm17,%zmm8
 1c6:   17 
 1c7:   62 92 75 50 b8 7c a5    vfmadd231ps 0x60(%r13,%r12,4){1to16},%zmm17,%zmm7
 1ce:   18 
 1cf:   62 92 75 50 b8 74 a5    vfmadd231ps 0x64(%r13,%r12,4){1to16},%zmm17,%zmm6
 1d6:   19 
 1d7:   62 92 75 50 b8 6c a5    vfmadd231ps 0x68(%r13,%r12,4){1to16},%zmm17,%zmm5
 1de:   1a 
 1df:   62 92 75 50 b8 64 a5    vfmadd231ps 0x6c(%r13,%r12,4){1to16},%zmm17,%zmm4
 1e6:   1b 
 1e7:   62 92 75 50 b8 5c a5    vfmadd231ps 0x70(%r13,%r12,4){1to16},%zmm17,%zmm3
 1ee:   1c 
 1ef:   62 92 75 50 b8 54 a5    vfmadd231ps 0x74(%r13,%r12,4){1to16},%zmm17,%zmm2
 1f6:   1d 
 1f7:   62 92 75 50 b8 4c a5    vfmadd231ps 0x78(%r13,%r12,4){1to16},%zmm17,%zmm1
 1fe:   1e 
 1ff:   62 92 75 50 b8 44 a5    vfmadd231ps 0x7c(%r13,%r12,4){1to16},%zmm17,%zmm0
 206:   1f 
 207:   49 83 c4 20             add    $0x20,%r12
 20b:   41 80 ff 08             cmp    $0x8,%r15b
 20f:   0f 82 df fe ff ff       jb     f4 <_Z7mm_dep2PfS_S_i+0xc4>
 215:   62 51 7c 48 11 3c b1    vmovups %zmm15,(%r9,%rsi,4)
 21c:   80 c3 01                add    $0x1,%bl
 21f:   62 51 7c 48 11 74 b1    vmovups %zmm14,0x40(%r9,%rsi,4)
 226:   01 
 227:   49 81 c3 00 01 00 00    add    $0x100,%r11
 22e:   62 51 7c 48 11 6c b1    vmovups %zmm13,0x80(%r9,%rsi,4)
 235:   02 
 236:   49 81 c2 00 01 00 00    add    $0x100,%r10
 23d:   62 51 7c 48 11 64 b1    vmovups %zmm12,0xc0(%r9,%rsi,4)
 244:   03 
 245:   80 fb 08                cmp    $0x8,%bl
 248:   62 51 7c 48 11 5c b1    vmovups %zmm11,0x100(%r9,%rsi,4)
 24f:   04 
 250:   62 51 7c 48 11 54 b1    vmovups %zmm10,0x140(%r9,%rsi,4)
 257:   05 
 258:   62 51 7c 48 11 4c b1    vmovups %zmm9,0x180(%r9,%rsi,4)
 25f:   06 
 260:   62 51 7c 48 11 44 b1    vmovups %zmm8,0x1c0(%r9,%rsi,4)
 267:   07 
 268:   62 d1 7c 48 11 7c b1    vmovups %zmm7,0x200(%r9,%rsi,4)
 26f:   08 
 270:   62 d1 7c 48 11 74 b1    vmovups %zmm6,0x240(%r9,%rsi,4)
 277:   09 
 278:   62 d1 7c 48 11 6c b1    vmovups %zmm5,0x280(%r9,%rsi,4)
 27f:   0a 
 280:   62 d1 7c 48 11 64 b1    vmovups %zmm4,0x2c0(%r9,%rsi,4)
 287:   0b 
 288:   62 d1 7c 48 11 5c b1    vmovups %zmm3,0x300(%r9,%rsi,4)
 28f:   0c 
 290:   62 d1 7c 48 11 54 b1    vmovups %zmm2,0x340(%r9,%rsi,4)
 297:   0d 
 298:   62 d1 7c 48 11 4c b1    vmovups %zmm1,0x380(%r9,%rsi,4)
 29f:   0e 
 2a0:   62 d1 7c 48 11 44 b1    vmovups %zmm0,0x3c0(%r9,%rsi,4)

かなり理想に近いかたちをしています。

  1. データのロード(vmovups addr zmm)
  2. 計算(vfmadd231ps (addr) zmm zmm)
  3. データのストア(vmovups zmm addr)

という流れがまとまってかつ余計な処理がほとんど行われていません。
行列並び替えなしのアセンブラと比較すると、
fmadd231psの引数アドレスの並びが連続であるため、
連続データアクセスに対応したハードウェアプリフェッチが十分に効くことが推測されます。
こうなればソフトウェアプリフェッチも不要になります。

再度並列に

ようやくOpenMPを利用して意味がある段階になったと思います。
OpenMPを有効にして性能を改めて測定してみました。(上段N,下段GFlops)

1024 4096 8192 16384 26624
40.99 680.532 2661.64 2738.29 3031.19

最終的に理論性能の約50%ほど出せたことになります。
データが大きいほど性能が出るというのは、
データアクセスの遅延が再利用によって隠蔽されてゆくことが主な理由です。

それに対するIntel MKL2017の性能は

1200 4000 8000 16000 26800
45.77 140.49 1105.45 4039.71 4740.67

結果としてNが小さい場合は性能でIntelが提供する科学計算ライブラリの性能を抜くことができていますが、
大きくなると逆転されてしまいます。
これを勝ちと見るか負けと見るかは..個人によるのでは。

ただしBLASはメモリを2倍利用することもありませんし、
そもそも計算する行列を正方行列に限定しません。

Intel MKLの本気

実はMKLのBLASは内部で行列をバッファリングしています。
そのため二回目以降の実行は速度が急激に上がり、
N>5000で4800GFLOPS以上の性能を出すことができます。

まとめ

如何に行列積演算を速くするかという話ではなく、
BLASを始めとした科学計算ライブラリは素人のにわか仕込みよりはるかに素晴らしく、
ぜひその存在と有用性を知ってもらいたいという話でした。

素人が素人なりに頑張ればそこそこにプロセッサの性能を使うことはできますが、
プロセッサの構成とメモリー周りの構造、命令セットすべてを把握してプログラムするなんて勘弁してほしいですよね。
逆にそこを理解しようとせずに低級な高級言語でプログラムを書いても、
正しく科学計算ライブラリを利用した高級な高級言語にパフォーマンスで劣るということは容易に考えられることです。

手軽に高速なプログラムを書くといえばPython内部のNumPy,SciPyの利用などがありますが、
これらも内部ではBLAS派生のAtlasライブラリによって高速化が行われています。

もちろんIntel MKLやその他の科学計算ライブラリは行列積だけでなく、
通常のベクトル演算・高速フーリエ変換などもサポートしているのでぜひ使ってみてください。

最後に

記事に間違いがありましたらコメントをぜひお願い致します。