AMD MI300A(APU)のメモリ確保とrocBLAS出力の正否

目的

AMD MI300Aにむけてコードの移植を行っています。先日の記事でも四苦八苦したように、APUであるMI300AはハードウェアネイティブなUnified memory/Managed memoryであることも影響して、メモリ確保の方法によっては正しい動作をしないことがあります。特に、XNACKというAMD独自の機構の影響もあります。前回は計算がクラッシュする場合についてレポートしました。

今回は、計算は実行できるものの、正しい結果が得られないケースが見つかったので調査してみました。ここで扱い計算は、rocBLASを利用した行列行列積(DGEMM)になります。

環境

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

テストコード

以下のコードでDGEMMのパフォーマンスを計測します。m * k行列と、k*n行列の行列積です。記載の設定では、元々本問題に遭遇した条件としてm=8,n=32,k=3375を設定しました。また、先日の記事とは異なり、このソースコード単体で実行できるように、独自ライブラリなどは排しました。

/*
main07.hip
Example for memory allocation vs. XNACK on rocBLAS.
*/


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

#include <hip/hip_runtime.h>
#include <rocblas/rocblas.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();
}



/*
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).
*/
double* ModeAlloc(int mode, size_t size){
    if(mode == 0 ){
        double* buf;
        hipMalloc((void**)&buf, sizeof(double)*size);
        return buf;
    }else if(mode == 1){
        double* buf;
        hipMallocManaged((void**)&buf, sizeof(double)*size);
        return buf;
    }else if(mode == 2){
        return new double[size];
    }else{
        return nullptr;
    }
}

void ModeFree(int mode, double* ptr){
    if(mode == 0 || mode == 1){
        hipFree(ptr);
    }else if(mode == 2){
        delete [] ptr;
    }
}
    
int TestDGEMM2(int modeA, int modeB, int modeC, double** p_reference){
    
    const int m = 8; // A, Cの行数
    const int n = 32; // B, Cの列数
    const int k = 3375;// Aの列数、Bの行数


    // デバイス側のメモリ確保
    double* d_A = ModeAlloc(modeA, m * k);
    double* d_B = ModeAlloc(modeB, k * n);
    double* d_C = ModeAlloc(modeC, m * n);
    double alpha = 1.0;
    double beta = 0.0;


    printf("floating point = FP%d\n", (int)sizeof(double)*8);
    printf("Allocation mode: %d  %d  %d\n", modeA, modeB, modeC);

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

    //APUを前提にダイレクトアクセスによる初期化//
    for(int64_t i = 0; i < (int64_t)m * (int64_t)k; ++i) {
        d_A[i] = (double)(i / k) / (double)m;
    }
    for(int64_t i = 0; i < (int64_t)k * (int64_t)n; ++i) {
        d_B[i] = (double)(i / k) / (double)n;
    }
    for(int i = 0; i < m * n; ++i) {
        d_C[i] = (double)(i / k) / (double)(m*n);
    }
    
    

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

    rocblas_handle handle;
    rocblas_create_handle(&handle);

    //first touch(rocmでは極端に初回が遅い場合があるので)     
    rocblas_dgemm(handle, rocblas_operation_none, rocblas_operation_none,
                m, n, k, &alpha, d_A, m, d_B, k, &beta, d_C, m);
            
    
    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呼び出し
        rocblas_dgemm(handle, rocblas_operation_none, rocblas_operation_none,
                m, n, k, &alpha, d_A, m, d_B, k, &beta, d_C, m);
    }
    hipDeviceSynchronize();
    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);fflush(stdout);
    begin_tm = end_tm;


    // 結果を比較//
    if(*p_reference == nullptr){
        double* ref = new double[m*n];
        *p_reference = ref;
        for (int i = 0; i < m * n; ++i) {
            ref[i] = d_C[i];
        }
        //初回のみ表示//
        std::cout << "Result matrix C:" << std::endl;
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < m; ++j) {
                std::cout << d_C[i + j * m] << " ";
            }
            std::cout << std::endl;
        }
    }else{
        const double THRESHOLD = 1.0e-6;
        const double* ref = *p_reference;
        double max_diff = 0.0;
        for (int i = 0; i < m * n; ++i) {
            const double d = fabs(ref[i] - d_C[i]);
            if(max_diff < d ) max_diff = d;
        }

        if(max_diff < THRESHOLD){
            printf("Result(%d,%d,%d) is Consistent, max_diff = %g\n", modeA, modeB, modeC, max_diff);
        }else{
            printf("Result(%d,%d,%d) is Inconsistent, max_diff = %g\n", modeA, modeB, modeC, max_diff);
        }
        fflush(stdout);
    }
    

    // リソースの解放
    rocblas_destroy_handle(handle);
    ModeFree(modeA, d_A);
    ModeFree(modeB, d_B);
    ModeFree(modeC, d_C);    
    
    return 0;
}

int main(int argc, char* argv[]) {
    double* reference = nullptr; 
    TestDGEMM2(0,0,0, &reference);
    printf("CHECK\n");fflush(stdout);

#if 1
    //for HSA_XNACK=1    
    TestDGEMM2(0,0,1, &reference);
    TestDGEMM2(0,0,2, &reference);
    TestDGEMM2(0,1,0, &reference);
    TestDGEMM2(0,1,1, &reference);
    TestDGEMM2(0,1,2, &reference);
    TestDGEMM2(0,2,0, &reference);
    TestDGEMM2(0,2,1, &reference);
    TestDGEMM2(0,2,2, &reference);
    TestDGEMM2(1,0,0, &reference);
    TestDGEMM2(1,0,1, &reference);
    TestDGEMM2(1,0,2, &reference);
    TestDGEMM2(1,1,0, &reference);
    TestDGEMM2(1,1,1, &reference);
    TestDGEMM2(1,1,2, &reference);
    TestDGEMM2(1,2,0, &reference);
    TestDGEMM2(1,2,1, &reference);
    TestDGEMM2(1,2,2, &reference);
    TestDGEMM2(2,0,0, &reference);
    TestDGEMM2(2,0,1, &reference);
    TestDGEMM2(2,0,2, &reference);
    TestDGEMM2(2,1,0, &reference);
    TestDGEMM2(2,1,1, &reference);
    TestDGEMM2(2,1,2, &reference);
    TestDGEMM2(2,2,0, &reference);
    TestDGEMM2(2,2,1, &reference);
    TestDGEMM2(2,2,2, &reference);
#else
   //for HSA_XNACK=0 
    TestDGEMM2(0,0,1, &reference);
    TestDGEMM2(0,1,0, &reference);
    TestDGEMM2(0,1,1, &reference);
    TestDGEMM2(1,0,0, &reference);
    TestDGEMM2(1,0,1, &reference);
    TestDGEMM2(1,1,0, &reference);
    TestDGEMM2(1,1,1, &reference);
#endif
    delete[] reference;
    return 0;
}

コンパイルオプションと実行時オプション

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

hipcc -O2 -std=c++17 -lrocblas -lamdhip64 --offload-arch=gfx942 main07.hip -o a.exe

ここで--offload-arch=gfx942の指定のところには、xnack+は加えませんでした。先日の記事にて記述の有無で動作の違いはないと思われるためです。今後、動作に違いが見られたら、また追記レポートします。

実行時には環境変数としてHSA_XNACK=1を設定しました。

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

また、後述しますが、ホストメモリの確保のケースを省いて、HSA_XNACK=0の場合にも試しています。

メモリ確保の比較方法

DGEMMの計算では、行列A, B, Cに対して、
$$
C = \alpha A B + \beta C
$$
を行います。ここで、3つの行列A, B, Cに対して、それぞれメモリ確保の方法を次の3パターンの中で変化させました。

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

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

3つの行列に対して、それぞれ3パターンのメモリ確保の方法を試すため、合計で27ケースのテストを行いました。

それぞれのケースにおいてDGEMMの計算結果である出力Cが正しいか否かをチェックします。正しい出力C（リファレンス）はA, B, C全てのメモリをmode = 0hipMalloc()として確保した場合の出力Cとしました。

計算を実行し、grepコマンドで標準出力からInconを検索すれば、次のような結果が得られます。

$ grep Incon log.txt
Result(1,1,1) is Inconsistent, max_diff = 11.25
Result(1,2,1) is Inconsistent, max_diff = 1.33203
Result(2,1,1) is Inconsistent, max_diff = 4.6875
Result(2,1,2) is Inconsistent, max_diff = 78.75

この場合は、メモリ確保の組み合わせとして4ケースで出力Cの結果が間違っていたことが分かります。

比較結果

さらに、厄介なことに、実行するタイミングによって、計算結果の正否が変わってしまうことが分かりました。そのため、10回程度の試験を行って、一度でも誤った結果が出たものは、以下の表において「否」と表示しました。

以下の表に環境変数HSA_XNACK=1をセットした場合の結果を示します。第1列～第3列にはそれぞれ行列A, B, Cのメモリ確保の方法を示しており、第4列にはその場合の出力Cの正否をY/ERRORで示します。

Aのメモリ確保	Bのメモリ確保	Cのメモリ確保	出力Cの正否
0	0	0	Y
0	0	1	Y
0	0	2	Y
0	1	0	Y
0	1	1	Y
0	1	2	ERROR
0	2	0	Y
0	2	1	Y
0	2	2	Y
1	0	0	Y
1	0	1	Y
1	0	2	Y
1	1	0	Y
1	1	1	ERROR
1	1	2	Y
1	2	0	ERROR
1	2	1	ERROR
1	2	2	ERROR
2	0	0	Y
2	0	1	Y
2	0	2	Y
2	1	0	ERROR
2	1	1	ERROR
2	1	2	ERROR
2	2	0	Y
2	2	1	Y
2	2	2	Y

また、環境変数をHSA_XNACK=0とした場合にも試しました。この場合、メモリ確保の方法でホストメモリ(mode=2)を指定するとクラッシュします。よって、メモリ確保の方法としては0 or 1の2通り（A, B, C）に対して計8通りで試しました。その結果、やはり、(1,1,1)の場合だけで出力Cが間違うことがある、という上記の表と同様の結果になりました。

議論

なんとも系統だった結果ではないのですが、ともかく複数のケースにおいて出力エラーとなることが分かりました。

そもそも、A, B, C三つの変数でメモリ確保の方法が異なるケースはコードとして行儀が良いのか？といえば、良くないのですが、移植作業の途中では度々起こりうる状況です。また、サードパーティーのライブラリを組み込みたい場合などにも発生しそうな状況です。

では、A, B, Cのメモリ確保の方法が統一されている場合、上記では(0,0,0), (1,1,1), (2,2,2)の3ケースでは、(1,1,1)の場合だけ結果が否となります。つまり、hipMallocManaged()で確保したメモリはMI300A(APU)環境ではrocBLASに渡してはいけない、というかなり衝撃的な結果です。

逆に本結果を逆手に取れば、MI300Aに関しては、hipMallocManaged()は使わず、CPUで触るメモリについてもhipMalloc()で確保してしまうのが、逃げ道かもしれません。そうすれば、環境変数でHSA_XNACK=0をセットすることで先日の記事のパフォーマンス劣化も回避できるという唯一のパスになり得ます。

結論

• APUであるMI300Aではメモリ確保の方法によってrocBLASの出力がおかしくなることがある。しかも、タイミングにより正否が変わる。
• hipMallocManaged()で確保したメモリはMI300A(APU)環境ではrocBLASに渡してはいけない。
• 環境変数でHSA_XNACK=0をセットした場合にもホストメモリ確保(mode=2)のアクセス不可であることを除いて、rocBLASの挙動は環境変数でHSA_XNACK=1の場合と共通。

Note

• AMDのマニュアルにはホストのメモリ確保としてmallocとhipHostMallocは違うというようなことも書かれているので、挙動が違うかもしれません。
• また、今回のテストではホストメモリの確保をnewで行っていますが、これもmallocと挙動が違ったりするかもしれません。

感想

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

また、APUではないAMD製GPUならどうなるのでしょうか？NvidiaのcuBLASならどうなるのでしょうか？興味(不安)は尽きません。