前回の記事（行列乗算の最適化入門 - よーる）では、L3キャッシュに乗りきる程度のサイズ（だいたい1000×1000くらい）の行列積について、シングルスレッドの理論性能の84%以上を出すことができるコードを紹介しました。 4000×4000くらいのもう少し大きな行列の乗算もしたくなったので、マルチスレッド化しました。非常に簡単なコードでマルチスレッドの理論性能の89%以上を出せることがわかったので、紹介します。

実験環境

Ubuntu 24.04 LTS
g++ (Ubuntu 13.2.0-23ubuntu4) 13.2.0
Ryzen9 7950X
- Zen4世代のCPUで、AVX512拡張命令セットが使用可能。doubleが8個入るzmmレジスタに対するfma命令を1cycleに1回実行できる
- zmmレジスタは32本ある
- L1Dキャッシュは32KiB
- L2キャッシュは1MiB、レイテンシは14サイクル
- L3キャッシュは4MiB×16コア、レイテンシは50サイクル
メインメモリは32GiB×2、DDR5 4800MT/s

OpenMP並列

OpenMPは、共有メモリ型のマルチスレッド実行を手軽に記述することができる標準化されたAPIです。 #pragma omp parallel forなどのディレクティブを挿入することで、いいかんじに並列化してくれます。

ここで考えなければいけないことは、どのように仕事を分割して各スレッドに割り当てるかです。出力c[i][j]の計算を複数のスレッドで分担すると、途中結果をどこにおいておくとか、どのスレッドが途中結果を足し合わせて最終結果を取りまとめるかとか、そういったややこしい問題が生じます。ここでは単純に、一つの出力c[i][j]に対しては担当スレッドは一つ、という方式で行きたいと思います。また分担はごく単純に、「スレッド0はc[0～249][0～3999]を担当」「スレッド1はc[250～499][0～3999]を担当」……などとしました。

また、面倒な問題として、OpenMPを使ってコンパイルするとその部分が関数に切り出されるという問題があります。これにより、ポインタのエイリアス解析がうまく働かなくなります。 C言語であればrestrictをつけることでエイリアスがないことを手軽にコンパイラに伝えることができるので、行列積コードはC言語で書くことにしました。

#include <iostream>
#include <chrono>
#include <random>
#include <omp.h>
#include "config.h"

extern "C" { void mm(double*,double*,double*); }

int main() {
    std::mt19937_64 mt;
    std::uniform_real_distribution dist(-1.0, 1.0);

    alignas(64) double ah[N*N];
    alignas(64) double bh[N*N];
    alignas(64) double ch[N*N];
    for( int i = 0; i < N*N; ++i ) {
        ah[i] = dist( mt );
        bh[i] = dist( mt );
        ch[i] = dist( mt );
    }
    std::cout << "initialized." << std::endl;
    const auto start = std::chrono::system_clock::now();

#pragma omp parallel for
    for( int tid = 0; tid < 16; ++tid )
    {
        int i0 = tid % 16 * ThreadBlockSizeI;
        int j0 = tid / 16 * ThreadBlockSizeJ;
        mm( &ah[i0*N], &bh[j0], &ch[i0*N+j0] );
    }

    const auto finish = std::chrono::system_clock::now();
    const double s = std::chrono::duration_cast<std::chrono::nanoseconds>( finish - start ).count() * 1e-9;
    static constexpr double flop_per_fma = 2.0;
    static constexpr double insn_per_fma = 1.0 / 8.0;
    std::cout << s << " seconds, " << N*N*N*flop_per_fma/s * 1e-9 << " GFLOPS, " << N*N*N*insn_per_fma/s << " FPIns/s" << std::endl;
}

行列積コード

基本的には前回のコードと同様です。ただし、前回は全てのデータがL3キャッシュに乗ると仮定していましたが、今回はそうではないため、L3キャッシュのブロッキングに相当するループを追加しました。

#include <stddef.h>
#include <math.h>
#include "config.h"

void mm( double*restrict a, double*restrict b, double*restrict c ) {
    for( int i1 = 0; i1 < ThreadBlockSizeI; i1 += L3CacheBlockSizeI )
    for( int k1 = 0; k1 < ThreadBlockSizeK; k1 += L3CacheBlockSizeK )
    for( int j1 = 0; j1 < ThreadBlockSizeJ; j1 += L3CacheBlockSizeJ )
    for( int i2 = 0; i2 < L3CacheBlockSizeI; i2 += L1DCacheBlockSizeI )
    for( int k2 = 0; k2 < L3CacheBlockSizeK; k2 += L1DCacheBlockSizeK )
    for( int j2 = 0; j2 < L3CacheBlockSizeJ; j2 += L1DCacheBlockSizeJ )
    for( int i3 = 0; i3 < L1DCacheBlockSizeI; i3 += RegisterBlockSizeI )
    for( int k3 = 0; k3 < L1DCacheBlockSizeK; k3 += RegisterBlockSizeK )
    for( int j3 = 0; j3 < L1DCacheBlockSizeJ; j3 += 1 )
    for( int i4 = 0; i4 < RegisterBlockSizeI; i4 += 1 )
    for( int k4 = 0; k4 < RegisterBlockSizeK; k4 += 1 )
    {
        int i = i1 + i2 + i3 + i4;
        int k = k1 + k2 + k3 + k4;
        int j = j1 + j2 + j3;

        c[i*N+j] = fma( a[i*N+k], b[k*N+j], c[i*N+j] );
    }
}

実験の結果、(i3,k3,j3)の三重ループにおいて、L1Dキャッシュに収まるようにブロッキングするのは最適ではなく、L1Dキャッシュを少しはみ出すくらい（L2キャッシュに余裕で入るくらい）のブロックサイズにしたほうがいいことがわかりました。これはZen4コアにおけるL2キャッシュのレイテンシは十分に短く、強力なアウトオブオーダーエンジンにより十分隠蔽できる範囲であるからだと思われます。逆にL1Dキャッシュをはみ出す量が多くなると性能が下がるのは、L1Dキャッシュ←→L2キャッシュのバンド幅は無尽蔵ではなく、ある程度はL1Dキャッシュでヒットしないとさばききれないから、と説明できるでしょう。

また、(i2,k2,j2)の三重ループは、本来の意図としてはL3キャッシュに収まる程度のブロッキングを行うために追加したものですが、最適なブロックサイズはL2キャッシュに余裕で入るくらいのものとなりました。これについては、以下のような説明が可能ですが、この解釈で正しいのでしょうか。

L3キャッシュはコア間で共有しているため、それに頼るようなブロッキングを行ってしまうと、他のスレッドのデータで追い出されてしまったりして都合がよくない
したがって、各スレッドのブロッキングに使うキャッシュはL2キャッシュが主力
他のスレッドがメインメモリからL3キャッシュに持ってきたおかげでヒットになる公算が高いとはいえ、それに頼りすぎるのではなく、ヒットしたらラッキー程度に考えておいたほうが良い
したがって、L2キャッシュミスした場合のレイテンシは読めないものと思ったほうがよく、かなり積極的なプリフェッチが行われる（？）
プリフェッチが積極的に行われてL2キャッシュにどんどんデータが入ってくるとすると、L2キャッシュ容量ぎりぎりまで使ったブロッキングを行ってはダメで、容量に余裕を持っておいたほうがよい

ブロッキングに関する情報をまとめたconfig.hは以下のようになります。

static const size_t N = 4000;
static const size_t ThreadBlockSizeI = 250;  // 5周（×16スレッド）
static const size_t ThreadBlockSizeK = 4000; // 50周
static const size_t ThreadBlockSizeJ = 4000; // 50周
static const size_t L3CacheBlockSizeI = 50;  // 1周
static const size_t L3CacheBlockSizeK = 80;  // 2周
static const size_t L3CacheBlockSizeJ = 80;  // 2周
static const size_t L1DCacheBlockSizeI = 50; // 10周
static const size_t L1DCacheBlockSizeK = 40; // 10周
static const size_t L1DCacheBlockSizeJ = 40; // SIMD方向8要素×5周
static const size_t RegisterBlockSizeI = 5;  // 5レジスタ並列に
static const size_t RegisterBlockSizeK = 4;  // fma連鎖4回

レジスタブロッキングは大量にあるzmmレジスタを生かすため、5(i4)×4(k4)×8(SIMD方向)としました。 L1Dキャッシュブロッキングは50(i3)×40(k3)×40(j3)（5600要素、43.75KiB）としました。 L3キャッシュブロッキングは50(i2)×80(k2)×80(j2)（14400要素、112.5KiB）としました。つまり、i2は一周しか実行しないので実はループではありません。結局、スレッドあたりの行列乗算コードは十重ループ、その外側にOpenMP並列のループが一重、という感じになっています。

性能

Ryzen9 7950Xはフル稼働させると、4.9GHzくらいしか出ません。周波数が4.9GHzだとすると、理論性能は4.9GHz×16FLOP/cycle×16コア＝1254.4GFLOPSということになります。

これに対し、私のコードは1117GFLOPSの性能を示しました。これは理論性能の89.04%に相当します。一番高い性能を示したときは1120GFLOPSに達したこともありました（4000×4000行列の行列乗算が0.1142秒で完了しました）。 0.1142秒の間にSIMD長8のfma命令が80億回も行われたことになります。

まとめ

4000×4000行列の行列乗算しかできませんが、数十行のコードで理論性能の89%以上を出すことができました。ここまで簡単にそれなりの性能を出せるとは思っていなかったので、コンパイラの自動SIMD化とOpenMPの威力を知るいい機会になりました。