はじめに

AMD MI300Aにむけてコードの移植を行っています。移植作業においては、Unified memory/Managed memoryを使うことでデバイス－ホスト間のメモリ転送を、まずは気にしなくてよい、という点では作業効率が大変高くなります。特に、MI300AのようなネイティブにUnified memoryなハードウェア構成では、**「mallocやnewでhost側で確保した変数(配列)もGPU側から操作可能」**という特性により、さらに作業効率が上がります（数万行もあるコードだと、hipMallocManaged/cudaMallocManagedに全てを書き換えるだけでも大変です）。

ここで、AMDのUnified memoryの動作を司る機構としてXNACKというものがあります。メモリのページフォルトの管理をしてくれているものですが、この設定には

コンパイル時のオプションによる指定
実行時の環境変数による指定

の2か所で必要であり、この組み合わせが悪いと、最悪の場合クラッシュします。

クラッシュするくらいかな、動作する組み合わせさえ見つけられればそれでいいのかな、と思っていたのですが、試行錯誤している中で、動作してもパフォーマンスが著しく落ちるケースがありました。しかもそれが、メモリのアロケート方法に依存するというものでした。この点について、備忘録として纏めたいと思います。

2025/9/9追記しました

環境

GPU: AMD Instinct MI300A
ROCm 6.3.3
1ノード占有にて実行

テストコード

以下のコードでDGEMMのパフォーマンスを計測します。m * k行列と、k*n行列の行列積です。記載の設定では行列サイズm,n,kが正方から著しく離れているので、多少パフォーマンスは落ちますが、各自調整してみてください。


/*
main.hip
Example of using cuBLAS/rocBLAS for matrix multiplication (GEMM) with single/double precision.
*/

#include <cstdlib>
#include <cstdio>
#include <memory>
#include <random>
#include <iostream>

#include "gyield.h"
#include "gy_blas.h"

#include <chrono>

template <typename DURATION>
inline
double DoubleSec(DURATION d) {
    return 1.0e-9 * (double)std::chrono::duration_cast<std::chrono::nanoseconds>(d).count();
}


using high_f = double;


/*
mode == 0: device memory not managed.
mode == 1: managed memory.
mode == 2: managed memory which is alllocated by new on host (for native unified memory).
*/
int TestDGEMM(int mode){
    
    const int m = 901*2; // A, Cの行数
    const int n = 901*2; // B, Cの列数
    const int k = 80*80*80; // Aの列数、Bの行数
    

    // デバイス側のメモリ確保
    high_f *d_A, *d_B, *d_C;
    high_f *h_A, *h_B, *h_C;
    high_f alpha = 1.0;
    high_f beta = 0.0;

    auto begin_tm = std::chrono::high_resolution_clock::now();
    auto end_tm = begin_tm;

    if(mode == 0 ){
        // ホスト側の行列データ
        h_A = new high_f[(int64_t)m * (int64_t)k];//
        h_B = new high_f[(int64_t)k * (int64_t)n];//
        h_C = new high_f[(int64_t)m * (int64_t)n];//
        
        gyCheckError(gyMalloc((void**)&d_A, (int64_t)m * (int64_t)k * sizeof(high_f)), "Failed to allocate device memory for A");
        gyCheckError(gyMalloc((void**)&d_B, (int64_t)k * (int64_t)n * sizeof(high_f)), "Failed to allocate device memory for B");
        gyCheckError(gyMalloc((void**)&d_C, (int64_t)m * (int64_t)n * sizeof(high_f)), "Failed to allocate device memory for C");
        gyDeviceSynchronize();

        end_tm = std::chrono::high_resolution_clock::now();
        printf("malloc = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;


        for(int64_t i = 0; i < (int64_t)m * (int64_t)k; ++i) {
            h_A[i] = (high_f)(i / k) / (high_f)m;
        }
        for(int64_t i = 0; i < (int64_t)k * (int64_t)n; ++i) {
            h_B[i] = (high_f)(i / k) / (high_f)n;
        }
        for(int i = 0; i < m * n; ++i) {
            h_C[i] = (high_f)(i / k) / (high_f)(m*n);
        }
        
        printf("floating point = FP%d\n", (int)sizeof(high_f)*8);

        end_tm = std::chrono::high_resolution_clock::now();
        printf("initialize = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;
        
        

        // ホストからデバイスへデータ転送
        gyCheckError(gyMemcpy(d_A, h_A, (int64_t)m * (int64_t)k * sizeof(high_f), gyMemcpyHostToDevice), "Failed to copy A to device");
        gyCheckError(gyMemcpy(d_B, h_B, (int64_t)k * (int64_t)n * sizeof(high_f), gyMemcpyHostToDevice), "Failed to copy B to device");
        gyCheckError(gyMemcpy(d_C, h_C, (int64_t)m * (int64_t)n * sizeof(high_f), gyMemcpyHostToDevice), "Failed to copy C to device");
        gyDeviceSynchronize();
        end_tm = std::chrono::high_resolution_clock::now();
        printf("cpy_to_dev = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;
    }else {
        //managed memory//
        if(mode == 1){
        gyCheckError(gyMallocManaged((void**)&d_A, (int64_t)m * (int64_t)k * sizeof(high_f)), "Failed to allocate device memory for A");
        gyCheckError(gyMallocManaged((void**)&d_B, (int64_t)k * (int64_t)n * sizeof(high_f)), "Failed to allocate device memory for B");
        gyCheckError(gyMallocManaged((void**)&d_C, (int64_t)m * (int64_t)n * sizeof(high_f)), "Failed to allocate device memory for C");
        }else{                
            d_A = new high_f[(int64_t)m * (int64_t)k];//
            d_B = new high_f[(int64_t)k * (int64_t)n];//
            d_C = new high_f[(int64_t)m * (int64_t)n];//
        }
        h_A = d_A;
        h_B = d_B;
        h_C = d_C;
        

        end_tm = std::chrono::high_resolution_clock::now();
        printf("malloc = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;

        for(int64_t i = 0; i < (int64_t)m * (int64_t)k; ++i) {
            d_A[i] = (high_f)(i / k) / (high_f)m;
        }
        for(int64_t i = 0; i < (int64_t)k * (int64_t)n; ++i) {
            d_B[i] = (high_f)(i / k) / (high_f)n;
        }
        for(int i = 0; i < m * n; ++i) {
            d_C[i] = (high_f)(i / k) / (high_f)(m*n);
        }
        
        end_tm = std::chrono::high_resolution_clock::now();
        printf("initialize = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;
                
    }
            
    //first touch(rocmでは極端に初回が遅い場合があるので)
    blas_DGEMM_n(m, n, k,
            d_A, d_B, d_C, 
            m,   // Cのleading dimension
            alpha,
            beta);
    //追記:2025/9/9 (注意)
    //gyDeviceSynchronize();をここに書き忘れているのでfirst touchのタイムラグが時間計測に含まれてしまっている

    
    auto handle = gy::blas::SingletonHandle::Get();
    end_tm = std::chrono::high_resolution_clock::now();
    printf("gen handle = %f [s]\n", DoubleSec(end_tm - begin_tm));
    begin_tm = end_tm;

    const int NSTEP = 100;
    for(int istep=0; istep<NSTEP; ++istep) {
    
        // cuBLAS/rocBLASのdgemm呼び出し
        blas_DGEMM_n(m, n, k,
                d_A, d_B, d_C, 
                m,   // Cのleading dimension
                alpha,
                beta); 
    }
    gyDeviceSynchronize();
    end_tm = std::chrono::high_resolution_clock::now();
    
    
    printf("gemm = %f [s]\n", DoubleSec(end_tm - begin_tm));
    const double gflops = (double)NSTEP * (double)m * (double)n * (double)k / DoubleSec(end_tm - begin_tm) * 2.0e-9; // Gflops
    printf("GFLOPS = %f\n", gflops);
    begin_tm = end_tm;
    
    if(mode == 0 ){
        // デバイスからホストへ結果を転送
        gyCheckError(gyMemcpy(h_C, d_C, (int64_t)m * (int64_t)n * sizeof(high_f), gyMemcpyDeviceToHost), "Failed to copy C to host");
        gyDeviceSynchronize();
        end_tm = std::chrono::high_resolution_clock::now();
        printf("cpy to host = %f [s]\n", DoubleSec(end_tm - begin_tm));
        begin_tm = end_tm;
    }
    
    // 結果を表示
    std::cout << "Result matrix C:" << std::endl;
    for (int i = 0; i < 10; ++i) {
        for (int j = 0; j < 10; ++j) {
            std::cout << h_C[i + j * m] << " ";
        }
        std::cout << std::endl;
    }

    // リソースの解放
    if(mode <= 1 ){
        gyFree(d_A);
        gyFree(d_B);
        gyFree(d_C);
    }
    if(mode == 0 || mode == 2 ){
        delete[] h_A;
        delete[] h_B;
        delete[] h_C;
    }
    return 0;
}

int main(int argc, char* argv[]) {
    TestDGEMM(0);
    TestDGEMM(1);
    TestDGEMM(2);
    return 0;
}

一部に自作のラッパー関数を使っているので、そのままだとコンパイルは通りませんが、この文章に興味のある方なら、比較的簡単にご自身の環境に書き換えられると思います。

gy*で始まる関数は、同名のhip*やcuda*のラッパーです。
blas_DGEMM_nはrocBLAS/cuBLASにおけるGEMMのラッパーです。

比較設定１：コンパイルオプション

以下のコンパイルオプションでコンパイルしました。ソースファイルはmain.hipです。

hipcc -O3 -std=c++17 --offload-arch=gfx942:xnack+ -lrocblas -lamdhip64 main.hip -o a.exe

ここで--offload-arch=gfx942:xnack+とするとXNACKが有効なコードとなります。xnack+を明示的に書かずに--offload-arch=gfx942とだけ指定したコードと比較することにします。

比較設定２：実行時オプション

実行時には環境変数としてHSA_XNACK=1を設定するとXNACKが有効になるそうです。一方で、HSA_XNACK=0とすれば無効にできるそうです。この両者で比較します。

export HSA_XNACK=1   #もしくは 0 をセット
a.out

比較設定３：メモリのアロケートの比較

上述のコードでは、TestDGEMM()関数の引数としてmode==0, 1, 2の3通りの指定により、メモリの確保の仕方を切り替えています

mode = 0：hipMalloc()でデバイス専用メモリを確保
mode = 1：hipMallocManaged()でmanagedメモリを確保
mode = 2：newでhost用メモリを確保

ただし、MI300Aでは、ハードウェアがネイティブにUnified memoryであるためmode==2でも動作します。

比較結果

MI300Aで上記のプログラムを実行し、単位時間当たりの倍精度の浮動小数点演算回数(GFLOPS)を比較しました。あくまで計算時間からコード上で見積もった概算値です。

各ケースについては次のように命名しました。

Case A: コンパイル時指定--offload-arch=gfx942:xnack+、かつ、環境変数HSA_XNACK=1

Case B: コンパイル時指定--offload-arch=gfx942、かつ、環境変数HSA_XNACK=1

Case C: コンパイル時指定--offload-arch=gfx942:xnack+、かつ、環境変数HSA_XNACK=0

Case D: コンパイル時指定--offload-arch=gfx942、かつ、環境変数HSA_XNACK=0

それぞれの場合に、さらに配列のメモリ確保の方法をmode=0～2の範囲で変えて、実行性能GFLOPSを調べた結果を以下の表に示します。

メモリ確保	Case A	Case B	Case C	Case D
hipMalloc	36701	36728	48424	48529
hipMallocManaged	17476	17581	48495	48374
new	29318	29301	0	0

まず最もパフォーマンスの高いCase Dに比べて、Case A/Bでは約35%～75％程度のパフォーマンスになってしまいます。Case CとDの差や、AとBの差は誤差の範囲でしょう。

また、GFLOPSの数字が0になっている部分は、Memory access faulで落ちた場合を指します。Case C/Dのnewの部分で0になっていることから、C++のnewで確保したメモリをデバイスで触らせるには、環境変数HSA_XNACK=1を有効にする必要があるというのが分かります。おそらくC標準ののmallocでも同様でしょう。

議論

Case AとB、Case CとDはそれぞれほぼ同じ結果になっています。つまり、コンパイルオプションのxnack+の明記の有無では違いがなく、環境変数HSA_XNACKによって違いが出ていることが分かります。コンパイルオプションに明記しなかった場合にはXNACKの有効時・無効時両対応のコードが生成されるという話もあるので、そういうことかもしれません。コンパイルオプションには別途xnack-という明記も可能なようなので、試してみるのもよいかもしれません。

さて、移植作業の観点から見ると、Case A/BにおいてhipMallocManagedで確保した場合に最もパフォーマンスが落ちているというのは、コード移植の工程でhipMallocManagedを使うには悩ましい要素となります。特に、MI300A以外の環境へのポーティングも考えるなら、newからhipMallocManagedへの変更という作業をすることになるので、移植作業でパフォーマンスが落ちるのは困ります。一方で、パフォーマンスの高いCase C/Dとして動作させることができるのは、全カ所でnewからhipMallocManagedへ変更し終わった後です。したがって、最終的にはパフォーマンスが出ることを信じて、メモリ確保に関して移植しきるしかないという辛い工程にならざるを得ません。

Note

また、あらためてパフォーマンスに関しては、行列サイズを正方に近づければもう少し上がります。実際にm=n=k=5802としてみたところ、約72000GFLOPSとなりました。

注意(追記)

上記のコードでは、"first touch"の計算の後にDeviceSynchronizeをし忘れています。そのため、結果で示した計算時間にはfirst touch時の遅延が含まれてしまっています。

DeviceSynchronizeを正しく行った場合の結果も後日纏めますが、XNACKの有効／無効の差は多少改善します。

結論

環境変数HSA_XNACKの設定により演算パフォーマンスが大きく異なる。
環境変数HSA_XNACK=0の方が高いパフォーマンスが出るが、一方でnewなどのC++標準の方法で確保したホストメモリをデバイスから触るとMemory access faulで落ちる。
移植手順は、(1)環境変数HSA_XNACK=1として移植開始。(2)全てのnewをhipMallocManagedもしくはhipMallocに置き換えたら環境変数HSA_XNACK=0としてパフォーマンス確認。
それまでは信じて移植を進める

感想

Unified memoryという激甘の背後には、激辛（げきつら）の作業があるのですね。