DGX SparkでCUDAストリーム依存とバッファ寿命を検証する

2つのCUDAストリームで同じデバイスバッファを共有していると、何度か続けて正しく動作していても、入力サイズを変えた途端に古い値が現れることがあります。よくある原因は演算式ではなく、生成側ストリームの書き込みが終わる前に、消費側ストリームが読み取りを始めていることです。本記事ではDGX Spark上で2つのストリームを使う最小実験を構築し、投入順序、実行依存関係、バッファを再利用できるタイミングを個別に検証します。1回の実行が成功しただけでは、正しいとは判断しません。

1. ストリームを作る前にデータ依存関係を整理する

DGX SparkはNVIDIA GB10 Grace Blackwellを搭載し、128GBのLPDDR5x統合メモリを備えています。CPUとGPUは同じメモリプールを共有しますが、物理的にメモリが統合されていても、アプリケーションのストリーム間の実行順序まで自動的に保証されるわけではありません。明示的なコピー、カーネル起動、ホストからのアクセスには、それぞれの完了条件を守る必要があります。

この実験では、処理の境界を確認しやすいよう、デバイスメモリを明示的に確保し、ホスト側にはページロックメモリを使います。生成側ストリームでは「入力をデバイスにコピーし、各要素を2倍にする」処理を行い、消費側ストリームでは「各要素に1を加え、結果をホストにコピーする」処理を行います。依存関係の全体像は次のとおりです。

入力の準備完了 → コピー完了 → 2倍する処理の完了 → 1を加える処理の開始 → ホストへの転送完了 → ホストでの検査。

同じストリーム内の操作はストリームの順序に従って実行されますが、2つのストリームの間は明示的につなぐ必要があります。イベントは生成側ストリームの2倍するカーネルの後に記録し、消費側ストリームではそのイベントを待ってから、1を加えるカーネルを投入します。

ホストが生成側ストリームを先に、消費側ストリームを後に呼び出しても、それが示すのは投入順序だけです。ストリーム間のデータ依存関係には明確な同期の根拠が必要であり、負荷が軽いときにたまたま成立する実行順序に頼ってはいけません。

2. デフォルトストリームに依存しない最小プログラムを作る

TcpNodeでの支払い後、コンソールで接続情報を確認し、SSHでNVIDIA DGX OSにログインします。まず uname -m、nvidia-smi、nvcc --version を実行し、ホストのアーキテクチャ、デバイスとドライバーの情報、コンパイル用ツールチェーンの情報をそれぞれ保存します。アーキテクチャはARM64に対応する aarch64 であることを確認してください。コンパイラーのバージョンをドライバーの出力から推測してはいけません。

以下を stream_contract.cu として保存します。両方のストリームにノンブロッキングフラグを指定し、従来のデフォルトストリームによる暗黙の同期が、依存関係の欠落を覆い隠さないようにします。イベントは順序制御だけに使い、時間は計測しません。

#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>

static void check(cudaError_t result) {
    if (result != cudaSuccess) {
        std::fprintf(stderr, "%s\n", cudaGetErrorString(result));
        std::exit(1);
    }
}

__global__ void scale(float* data, int count) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < count) data[i] *= 2.0f;
}

__global__ void add(float* data, int count) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < count) data[i] += 1.0f;
}

int main() {
    constexpr int count = 4096;
    const size_t bytes = count * sizeof(float);
    float *input = nullptr, *output = nullptr, *device = nullptr;
    cudaStream_t producer, consumer;
    cudaEvent_t ready;

    check(cudaMallocHost(reinterpret_cast<void**>(&input), bytes));
    check(cudaMallocHost(reinterpret_cast<void**>(&output), bytes));
    check(cudaMalloc(reinterpret_cast<void**>(&device), bytes));
    check(cudaStreamCreateWithFlags(&producer, cudaStreamNonBlocking));
    check(cudaStreamCreateWithFlags(&consumer, cudaStreamNonBlocking));
    check(cudaEventCreateWithFlags(&ready, cudaEventDisableTiming));

    for (int i = 0; i < count; ++i) {
        input[i] = static_cast<float>(i);
        output[i] = -1.0f;
    }

    check(cudaMemcpyAsync(device, input, bytes,
                          cudaMemcpyHostToDevice, producer));
    scale<<<(count + 255) / 256, 256, 0, producer>>>(device, count);
    check(cudaGetLastError());
    check(cudaEventRecord(ready, producer));

    check(cudaStreamWaitEvent(consumer, ready, 0));
    add<<<(count + 255) / 256, 256, 0, consumer>>>(device, count);
    check(cudaGetLastError());
    check(cudaMemcpyAsync(output, device, bytes,
                          cudaMemcpyDeviceToHost, consumer));
    check(cudaStreamSynchronize(consumer));

    int failures = 0;
    for (int i = 0; i < count; ++i) {
        float expected = 2.0f * static_cast<float>(i) + 1.0f;
        if (output[i] != expected) ++failures;
    }
    std::printf("elements=%d mismatches=%d\n", count, failures);

    check(cudaEventDestroy(ready));
    check(cudaStreamDestroy(consumer));
    check(cudaStreamDestroy(producer));
    check(cudaFree(device));
    check(cudaFreeHost(output));
    check(cudaFreeHost(input));
    return failures == 0 ? 0 : 2;
}

3. 出力と終了コードの両方で検証する

同じディレクトリでコンパイルして実行します。

nvcc -O2 stream_contract.cu -o stream_contract
./stream_contract
printf 'exit=%s\n' "$?"

合格条件は、出力が elements=4096 mismatches=0 で、終了コードが0になることです。API呼び出しまたは同期が失敗した場合は1、計算結果が一致しない場合は2を返します。プロセスがクラッシュしなかったことだけを確認すると、処理は完了していても結果が誤っているケースを見逃します。

この例の値はすべて、単精度浮動小数点数で正確に表現できる整数の範囲内にあるため、等値で直接比較しています。実際の浮動小数点モデルにこの判定をそのまま適用することはできません。アルゴリズムに応じて、絶対誤差、相対誤差、非有限値の扱いを定める必要があります。

続いて要素数をそれぞれ1、257、4097に変更し、再コンパイルして実行します。これらのサイズで、単一要素、スレッドブロックの末尾、複数ブロックにまたがる境界を確認します。これはテスト入力であり、性能データではありません。成功時のスクリーンショットだけを残すのではなく、毎回ソースファイル、コンパイルコマンド、ツールチェーン情報、実行出力を保存してください。

失敗時の検出経路が機能しているか確認する

期待値を求める式の末尾にある 1.0f を一時的に 2.0f に変更します。プログラムが不一致を報告して2を返すことを確認し、その後で元の式に戻します。これにより、結果の検査が機能していることを確かめられます。「待機を削除すれば必ずエラーになる」ことを負のテストにしてはいけません。競合状態が表面化するかどうかはスケジューリング次第であり、依存関係が欠けているプログラムでも一時的にテストを通過することがあります。

4. イベントの境界をメモリ寿命に結び付ける

通常、cudaMemcpyAsync が戻ったことから分かるのは、呼び出しが投入されたということだけです。それを根拠に入力を書き換えてはいけません。この例では、コピーを投入してから消費側ストリームの同期が完了するまで、入力を変更しません。出力も、ホストへの転送を行うストリームの同期後にだけ読み取ります。この方法は保守的ですが境界が明確であり、正しさを確認するための最初の基準に適しています。

ページロックメモリは非同期転送の条件を満たすうえで役立ちますが、完了通知の代わりにはなりません。デバイスバッファも、別のストリームがまだアクセスしている間に解放してはいけません。また、128GBの統合メモリがあるからといって、大量のホストメモリをすべてページロックすべきではありません。バッチ容量は実際のワーキングセットに合わせて制限する必要があります。

生成側のイベントが対象とするのは、そのイベントより前の操作だけです。消費側ストリームで行う加算とホストへの転送は含まれません。したがって、生成側イベントの完了によって生成段階の終了は確認できますが、デバイスバッファを次のバッチで上書きできる状態になったとは判断できません。パイプラインに拡張する場合は、消費側ストリームでホストへの転送を行った後に完了イベントを記録し、それを使ってバッファ一式の再利用を制御します。

並行処理する各バッチスロットには、専用の入力、出力、デバイスバッファ、完了状態を持たせます。まず1つのスロットで検証し、その後でスロットを増やしてください。ストリーム数だけを増やし、同じポインターを保護なしで共有し続けてはいけません。

5. 切り分けでは同期範囲を段階的に絞り込む

結果がときどき誤る場合は、まず待機が消費側カーネルより前にあること、イベントの記録が生成側カーネルより後にあること、記録と待機に同じイベントを使っていることを確認します。イベントを記録する位置を1行早めただけで、「データの書き込みが完了した」という条件が、誤って「コピーが完了した」だけに変わってしまうことがあります。

次に、各カーネル起動後のエラーと、最後の同期処理の戻り値を確認します。前者は起動時の問題の一部を検出し、後者は非同期実行中のエラーを受け取るためのものです。互いに代用はできません。起動時の戻り値だけを調べ、完了を待たなければ、エラーが後続のAPI呼び出しで初めて表面化することがあります。

診断用のコピーでは、生成段階に一時的にストリーム同期を追加し、問題がストリーム間の順序に関係するかを調べられます。エラーが消えた場合は、デバイス全体の同期を入れて済ませるのではなく、依存関係とバッファの所有権を引き続き確認してください。全体同期はスケジューリングを変え、元のメモリ寿命の問題を隠してしまう可能性があります。

最終的に残すべきものは、明示的な生成側イベント、消費側の待機、再利用を制御する完了イベント、そして繰り返し実行できる出力検査です。依存関係のつながりとメモリの所有権をどちらも明確に説明できて初めて、マルチストリームプログラムの並行性をさらに拡張する土台が整います。

よくある質問