大規模分子動力学における GPU 起動型通信の実践ガイド
NVIDIA は、大規模な分子動力学シミュレーションにおいて GPU 起動型通信の効率を劇的に向上させる実用的なガイドラインと手法を公開した。
キーポイント
GPU 起動型通信の実装ガイド
大規模な分子動力学シミュレーション向けに、CPU を介さずに GPU が直接通信を開始・管理する手法の詳細な実装手順を提示している。
計算リソースの最適化
従来の CPU 中心のアーキテクチャから脱却し、GPU の並列処理能力を最大限に活用することで、シミュレーション全体のレイテンシとスループットを改善する。
大規模スケールへの対応
数千〜数万の GPU を使用する超並列環境においても、通信オーバーヘッドを最小化し、安定した計算性能を維持するための戦略が解説されている。
影響分析・編集コメントを表示
影響分析
このガイドは、創薬や材料科学などの分野で不可欠な大規模分子動力学シミュレーションの計算効率を根本から変える可能性を秘めている。特に GPU の性能を最大限引き出すための通信プロトコルを標準化することで、研究開発期間の短縮とコスト削減に寄与する実用的な成果である。
編集コメント
AI 分野以外でも、GPU を活用した大規模計算の効率化において重要な示唆を与える実用的な技術記事です。
分子動力学（MD）シミュレーションは、計算科学において最も負荷の高いワークロードの一つです。これを用いることで、研究者はタンパク質の折りたたみから創薬や材料発見に至るまで、原子レベルの挙動を驚くほど詳細に観察することができます。
しかし、その代償は甚大です。シミュレーションでは数十万から数千万個の原子をモデル化し、1 歩あたりフェムト秒（femtosecond）単位で進み、数十億ステップにわたって計算が行われます。問題サイズが通常固定されているため、並列処理は利用可能なすべてのリソースに同時に作業を広げる必要があり、この状態は「強スケーリング（strong scaling）」と呼ばれます。
GROMACS は、世界で最も広く使用されている MD パッケージの一つです。最新のハードウェアと 非同期 CPU-GPU 並列化 の進展により、パフォーマンスはサブミリ秒領域に達し、複数の GPU にわたって 1 タイムステップあたり 100〜200 マイクロ秒を達成しています。このレベルでの強スケーリングを実現するには極めて低いカーネル遅延が求められますが、シミュレーションがより多くの GPU にスケールするにつれて、GPU から GPU への通信は依然として根本的なボトルネックとなっています。
GROMACS を含む大規模 HPC アプリケーションの多くは、プロセス間通信にメッセージ Passing インターフェース（MPI）を使用しています。しかし、MPI は CPU 中心の実行を想定して設計されています。
GROMACS が GPU で実行される場合、通信ワークフローでは CPU がデータ転送を調整して GPU に再開を合図する間、GPU は待機状態になります。GROMACS のハロー交換（隣接する GPU ドメイン間で境界原子データを共有するアルゴリズム）において、このハンドオフは 3 つの空間次元全体で繰り返され、ピーク反復レートでは総 CPU ウォール時間の 50% 以上を消費し、スケーラビリティに上限をもたらしています。
本稿では、NVSHMEM を用いた GPU ネイティブ通信へ置き換えることで、このボトルネックを解消する戦略を検討します。
デバイス開始型リモートメモリアクセスによる GROMACS の高速化
GROMACS を加速する鍵は、GPU ネイティブ通信によって CPU と GPU の間のハンドオフを排除することです。NVIDIA NVSHMEM は、OpenSHMEM パーティションドグローバルアドレス空間モデルに基づいてリモートメモリアクセスを実装するためのライブラリです。NVSHMEM により、GPU カーネルがデータ転送を直接開始できるようになり、CPU がクリティカルパスから外れることで、計算と通信のより良いオーバーラップが可能になります。
GROMACS のハロー交換は、段階的な転送メカニズムを使用します。データは、各空間次元内で 1 つ以上の通信ステップ（パルス）にわたって中間ランクを介して中継されます。各パルスは、境界原子を連続する送信バッファに集約するパックカーネルで始まります。これにより、従来の実装が粗いフェーズレベルの直列化によって処理していた、通信フェーズ間の依存関係（Z → Y → X）の連鎖が生じます。この再定式化では、これらの粗いバリアを、融合されたカーネル内で直接表現される微細粒度のパルスごとのシグナルに置き換えます。
自然な出発点はストリームトリガー API (nvshmemx_put_signal_nbi_on_stream) です。これは CPU-GPU の同期バリアを最小限のコード再構成で除去します。しかし、2 つの制限があります。1 つ目は、通信がカーネル境界でのみ開始できるため、パルス内でパック処理と通信のオーバーラップができないことです。2 つ目は、ストリームによるシーケンシングのため、依存する転送連鎖（Z → Y → X）を融合できないことです。
これら 2 つの問題は、ホスト側の通信の置き換え および パルス間の並列化の有効化 セクションで、微細粒度のパルスごとの依存関係を用いて解決します。
image*図 1. 2D DD（領域分割）における GROMACS の GPU 常駐 MPI スケジュール。X/F 通信パルスが 2 回実行される様子を示す。CPU タイムライン（上部）はカーネル起動、GPU 同期待ち、MPI 処理を示している。GPU ストリーム（下部）は重なる作業と力計算の通信を示している*
以下の図は、2 つの通信パルスを持つシミュレーションステップの実行ダイアグラムを示しており、カーネル開始型 RMA（リモートメモリアクセス）、カーネル融合、イベント駆動マルチキュー実行がどのようにして通信と計算を重畳させるかを示しています。
*図 2. 2D DD における GROMACS の NVSHMEM スケジュール。CPU カーネル起動と GPU 計算ストリームの重畳により、ハロー交換がクリティカルパスから除外されている*
MPI ベースラインから最適化された GPU ネイティブなハロー交換へ
既存の MPI ベースのハロー交換は、通信量を最小限に抑えるために隣接するランク間でデータを中継しますが、その代償として、Z → Y → X の 3 つのフェーズ全体で逐次的依存関係が生じます。各フェーズには 1 つ以上のパルスが含まれています。以下のコードは、GROMACS ソースコードから抽出された座標ハロー交換パックカーネルと、ホスト側からの呼び出しを示しています。
// GPU kernel: pack boundary atom coordinates into a contiguous send buffer
__global__ void packSendBufKernel(float3* dataPacked, // packed coordinate buffer
const float3* data, // full coordinate array
const int* map, // index map: which atoms to pack
const int mapSize)
{
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIndex < mapSize)
{
dataPacked[threadIndex] = data[map[threadIndex]];
}
}
// GPU kernel: unpack received halo forces, scatter back into the force array via index map
template<bool accumulate>
__global__ void unpackRecvBufKernel(float3* forces, // full force array
const float3* dataPacked, // received packed forces
const int* map,
const int mapSize)
{
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIndex < mapSize)
{
if (accumulate)
forces[map[threadIndex]] += dataPacked[threadIndex];
else
forces[map[threadIndex]] = dataPacked[threadIndex];
}
}
// CPU host code: coordinate halo exchange (MPI path)
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// 1. Launch GPU pack kernel on non-local stream
packSendBufKernel<<<grid, block, 0, nonLocalStream>>>(
sendBuf, coords, indexMap, mapSize);
// 2. CPU blocks until GPU pack completes
cudaStreamSynchronize(nonLocalStream); // *** BLOCKING ***
// 3. CPU orchestrates MPI transfer
MPI_Isend(sendBuf, sendSize, sendRank, ...);
MPI_Irecv(coords + atomOffset, recvSize, recvRank, ...);
MPI_Waitall(...); // *** BLOCKING ***
// Received data lands directly
// in the coordinate buffer at the correct offset and is consumed
// by the non-local non-bonded force kernel as-is.
これらのカーネル自体は単純ですが、ボトルネックとなるのはこれらを囲む CPU 側の制御フローです。各パルスごとに、CPU は GPU の完了をブロックし、MPI 転送を開始し、すべての 3 つの次元にわたる各パルスに対してシリアライズシーケンスを再開するよう GPU に信号を送信する必要があります。
力ハロー交換は、同じパターンに従いますが方向が逆（X → Y → Z）です。座標を送信したランクは、そのハロー原子上で計算された力を受信し、力は完全な力配列内の正しい位置に蓄積される必要があるため、unpackRecvBufKernel を通じて散乱・展開ステップが必要です。
各パルスには 2 つのブロッキング CPU–GPU 同期が必要となります。1 つは MPI 呼び出し前に、パックカーネルが完了したことを保証するためのものであり、もう 1 つは次のパルスが転送データを消費する前に強制されるものです。3D 分解を用いて次元ごとに 1 パルスとする場合、座標については時間ステップあたりブロッキング待機が 6 回、力についてもさらに 6 回必要となり、合計で 12 回の待機が発生します。それぞれがクリティカルパスにレイテンシを追加し、GROMACS のイテレーションレートに影響を与えます。
ホスト側の通信を NVSHMEM に置き換える
この考えに基づき、最初の設計では、パッキング、リモート PUT、完了待ちの 3 つを、各通信パルスごとに単一のカーネルに統合します。ホストは依然としてパルスの順序で起動を発行しますが、パルスの間に存在した CPU–GPU バリアを取り除きます。実行順序の正しさは、もはや NVIDIA CUDA ストリームのみによって保証されます。統合されたパック・PUT・待ちカーネルのホストからの呼び出しと実装は以下の通りです：
// CPU host code: one kernel launch per pulse, no CPU-GPU sync in between
int pulseOffset = 0;
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// Pack + send + wait are fused into a single GPU kernel
packSendBufAndPutNvshmemKernel<<<grid, block, 0, nonLocalStream>>>(
coords, dataPacked, indexMap, sendSize,
sendRank, atomOffsetInSendRank,
signalReceiverRank + pulseOffset, signalCounter, bar, recvSize);
pulseOffset++;
パッキング–PUT–待ちの一連の操作は、単一のデバイス側起動に統合されています。GROMACS の実装を基に改変されたこのカーネルは以下の通りです：
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Pack: grid-stride gather into contiguous send buffer
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dataPacked[i] = coords[map[i]];
// Device-scoped barrier: all CTAs must finish packing before put
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
// Put packed data into peer's coordinate array
nvshmemx_float_put_signal_nbi_block(
&coords[atomOffsetInSendRank], // peer destination
dataPacked, // local source
sendSize * 3, // float count
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
// Wait for incoming data from recvRank before exiting
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}
この実装により、クリティカルパスから CPU-GPU の同期をすべて排除しました。CPU はカーネルの起動のみを行い、これは GPU 実行と重畳させることができます。すべての GPU スレッド（CTA）はパッキング後にデバイススコープのバリアに到達し、ブロック 0 がそのバリアで待機して、nvshmemx_float_put_signal_nbi_block 呼び出しを発行する前に全スレッドが完了したことを確認します。この呼び出しはパッキングされたデータをピアへ転送し、ピア上の signalReceiverRank を設定して完了を確認します；受信側はその信号を待つことでデータを受け取ります。
ただし、このバージョンでは依然としてパルスを順次処理しており（ストリーム順序により強制されています）、各パルスごとに個別のカーネル起動が発生することに注意してください。
NVIDIA NVLink 接続 GPU 間での効率的な転送のための最適化
前述の設計では、データは常に NVSHMEM トランスポート（nvshmemx_float_put_signal_nbi_block）を介してルーティングされます。これはあらゆる相互接続で機能しますが、ピアが NVLink で接続されている場合、GPU は互いのメモリに直接ロード/ストアできます。ローカルバッファへのパッキングと put 発行の代わりに、直接ピアの座標配列へパッキングし、個別のグローバルメモリの往復を排除します。
nvshmem_ptr(remotePtr, peerRank) は、ピアが NVLink を介して到達可能であれば非 null のデバイスポインタを返し、それ以外の場合は null を返します。これをカーネル内で一度クエリし、以下に示すように分岐処理を行います：
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Probe for NVLink: non-null means direct store is possible
float3* remotePtr = (float3*)nvshmem_ptr(coords, sendRank);
bool isNVLink = (remotePtr != nullptr);
// NVLink: pack directly into peer's coordinate buffer
// IB: pack into local staging buffer for later put
float3* dest = isNVLink ? (remotePtr + atomOffsetInSendRank)
: dataPacked;
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dest[i] = coords[map[i]];
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
if (isNVLink)
{
// Data already in peer memory — just signal completion
uint64_t* peerSignal = (uint64_t*)nvshmem_ptr(
signalReceiverRank, sendRank);
storeReleaseSysAsm(peerSignal, signalCounter);
}
else
{
// IB: combined data transfer + signal notification
nvshmemx_float_put_signal_nbi_block(
(float*)&coords[atomOffsetInSendRank],
(float*)dataPacked,
sendSize * 3,
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
}
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}
NVLink パスでは、パッキングループはピアの座標配列内の正確な位置に直接書き込みます。ブロック 0 がデバイススコープのバリアを使用してすべての CTAs（並行実行単位）のパッキング完了を確認すると、システムスコープのリリースストア (st.release.sys) を介してピアにシグナルを送信します。これにより、ピアがシグナルを読み取る前に、先行するすべての直接書き込みがピアに対して可視化されることが保証されます。データはすでに配置されているため、追加のデータプット操作は不要です。
nvshmem_ptr が null を返す場合、動作は最初の設計と一致します。データはローカルのステージングバッファにパッキングされ、nvshmemx_float_put_signal_nbi_block に引き渡されます。これは NVSHMEM 転送層によって InfiniBand、Slingshot、またはその他の RDMA ファブリックを介して配信されます。
この最適化は、2 つのパス間の重要な非対称性を浮き彫りにしています。NVLink パスでは、NVLink バンド幅を飽和させるために多くの CUDA コア (SM) が同時にストアを発行する必要がありますが、InfiniBand パスでは NVSHMEM 転送層を使用してバルク転送を開始するために単一のスレッド (CTA) で十分です。残りのボトルネックは、パルスごとのカーネル起動オーバーヘッドと、一部のデータにクロスパルス依存関係がない場合でも存在する、パルス間の厳格な順序付けです。
カーネル融合によるパルス間並列性の有効化
重要な洞察は、パルス内のすべてのデータが過去の脈動に依存しているわけではないという点です。パルスの大部分の原子は即座にパックして送信可能であり、「転送」された原子、つまり過去の脈動からデータを受け取る必要があるものだけが待機します。これは、各パルスのインデックスマップを dependencyAtomOffset で分割することで実現されます。
- 独立した原子 (atomIndex < dependencyAtomOffset): パルス実行時に即座にパックして送信しても安全
- 依存する原子 (atomIndex ≥ dependencyAtomOffset): 関連する過去の脈動のシグナルを待つ必要がある
この分割により、単一のカーネルがすべてのパルスを並列処理します。依存する作業は、特定のフェーズ全体ではなく、必要な個別のパルスごとのシグナルのみを待機します。
依存関係認識型パック
packSubset デバイス関数は、しきい値に基づいて選択的に原子をパックします。まず独立した原子に対して実行され、その後シグナルを待機してから依存する原子に対して再度実行されます。
// Templated pack: selectively gather atoms based on threshold
// packAll=true: pack everything (no dependency check needed)
// packAll=false, packLessThan=true: pack only atomIndex < threshold (independent)
// packAll=false, packLessThan=false: pack only atomIndex >= threshold (dependent)
template<bool packAll, bool packLessThan>
__device__ void packSubset(float3* dest,
float3* data,
const int* map,
int sendSize,
int gridStride,
int threadIndex,
int threshold,
int& hasDependencyAtoms)
{
for (int idx = threadIndex; idx < sendSize; idx += gridStride)
{
int atomIndex = map[idx];
float3 srcVal = data[atomIndex];
if constexpr (packAll)
{
dest[idx] = srcVal;
}
else
{
bool packNow = packLessThan ? (atomIndex < threshold)
: (atomIndex >= threshold);
if (packNow)
dest[idx] = srcVal;
else if (packLessThan)
hasDependencyAtoms++; // track deferred atoms in first pass
}
}
}
パルスは Z→Y→X の順序でグローバルに番号付けされています。各パルスは、依存する原子をパックする前に、すべての先行するパルスのシグナルを待機します。1 つの次元ごとに 1 つのパルス [z0, y0, x0] を持つ場合、z0 は依存関係を持ちません。y0 は z0 を待ち、x0 は両方を待ちます。これは packHaloCoords で実装されている for (int i = currPulse; i > 0; i--) ループが実現するものです。
// Dependency-aware packing: pack independent data immediately,
// wait for signals from prior pulses, then pack dependent data
__device__ void packHaloCoords(float3* dest,
int sendSize,
float3* data,
const int* map,
int currPulse,
uint64_t signalCounter,
uint64_t* signalReceiverRankCurr,
int dependencyAtomOffset)
{
int gridStride = blockDim.x * gridDim.x;
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (currPulse > 0)
{
// Pass 1: pack atoms below dependency threshold (no waiting)
int hasDependencyAtoms = 0;
packSubset<false, true>(dest, data, map, sendSize, gridStride,
threadIndex, dependencyAtomOffset,
hasDependencyAtoms);
// Thread 0 waits for all prior pulses to deliver their data
if (threadIdx.x == 0)
{
for (int i = currPulse; i > 0; i--)
while (loadRelaxedSys(signalReceiverRankCurr - i) != signalCounter)
;
}
constexpr unsigned fullWarpMask = 0xFFFFFFFF;
int need_to_wait = __any_sync(fullWarpMask, hasDependencyAtoms > 0);
__syncthreads();
// Pass 2: pack the deferred atoms (>= threshold)
if (need_to_wait)
{
int unused = 0;
packSubset<false, false>(dest, data, map, sendSize, gridStride,
threadIndex, dependencyAtomOffset,
unused);
}
}
else
{
// Pulse 0: no dependencies, pack everything unconditionally
int unused = 0;
packSubset<true, true>(dest, data, map, sendSize, gridStride,
threadIndex, 0, unused);
}
__syncthreads();
}
実用的な注意点：loadRelaxedSys は、シグナルがリモート GPU によって書き込まれるためシステムスコープの PTX ld.relaxed.sys.global.u64 を使用し、スピンループであるため緩和された順序付けを使用します。待機後の __syncthreads() により、ガード付きデータがすべてのスレッドに対して可視化されます。本番環境のコードでは、libcu++ の cuda::atomic_ref<uint64_t, cuda::thread_scope_system> を cuda::memory_order_relaxed とともに使用することで同等のセマンティクスを提供できます。
fused kernel は、各パルスを blockIdx.y の行に割り当て、その行内のすべてのブロックがグリッド・ストライド・ループ（blockIdx.x でインデックス付け）を介して 1 つのパルスに対して協力します。事前に計算されたリモートポインタは、実行時に転送経路を選択します：非 null は NVLink 直接ストアを意味し、null の場合は NVSHMEM put パス（InfiniBand）にフォールバックします。
// 1 つのカーネル起動が、すべての個別のパルスに対する pack + send オペレーションを置き換えます
// グリッド：(blocksPerPulse, totalNumPulses, 1) — パルスごとに 1 つの blockIdx.y
__global__ void fusedPulsesPackAndSendKernel(
float3* data, // 完全な座標配列
HaloExchangeData* pulseData, // パルスごとのメタデータ配列
uint64_t* signalReceiverRank,
uint64_t signalCounter,
uint32_t* gridSync, // パルスごとのブロック完了カウンター
int totalNumPulses,
int dependencyAtomOffset)
{
int currPulse = blockIdx.y;
if (currPulse < totalNumPulses)
{
HaloExchangeData halo = pulseData[currPulse];
// Precomputed by nvshmem_ptr; null => IB/RDMA path
float3* remotePtr = ha
原文を表示
Molecular dynamics (MD) simulations are among the most demanding workloads in computational science. Using them, researchers can observe atomic behavior in extraordinary detail, from protein folding to drug and materials discovery.
But they come at a steep cost. Simulations model hundreds of thousands to tens of millions of atoms, advancing femtoseconds per step across billions of steps. Because problem size is typically fixed, parallelization must spread work across all available resources simultaneously, a regime known as strong scaling.
GROMACS is one of the world’s most widely used MD packages. Modern hardware and heterogeneous CPU-GPU parallelization has pushed performance into the sub-millisecond regime, reaching 100–200 microseconds per time-step across multiple GPUs. Achieving strong scaling at this level demands extremely low kernel latency, but GPU-to-GPU communication remains a fundamental bottleneck as simulations scale across more GPUs.
Most large-scale HPC applications, including GROMACS, use the Message Passing Interface (MPI) for inter-process communication. However, MPI was designed for CPU-centric execution.
When GROMACS runs on GPUs, the communication workflow forces the GPU to pause while the CPU orchestrates data transfers before signaling the GPU to resume. In GROMACS’s halo exchange—the algorithm that shares boundary atom data between neighboring GPU domains—this handoff repeats across all three spatial dimensions, consuming more than 50% of total CPU wall time at peak iteration rates and capping scalability.
In this post, we examine strategies to eliminate this bottleneck by replacing CPU-orchestrated MPI with GPU-native communication using NVSHMEM.
Accelerating GROMACS using Device-initiated Remote Memory Access
The key to accelerating GROMACS is eliminating CPU-GPU handoffs with GPU-native communication. NVIDIA NVSHMEM is a library for implementing remote memory access based on the OpenSHMEM partitioned global address space model. NVSHMEM enables GPU kernels to initiate data transfers directly, removing the CPU from the critical path and enabling better overlap of computation and communication.
The GROMACS halo exchange uses a staged forwarding mechanism: data is relayed through intermediate ranks within each spatial dimension across one or more communication steps, or pulses. Each pulse begins with a pack kernel that gathers boundary atoms into a contiguous send buffer. This creates a chain of dependencies between communication phases (Z → Y → X) that the traditional implementation handles through coarse phase-level serialization. The reformulation replaces these coarse barriers with fine-grained per-pulse signals expressed directly in a fused kernel.
A natural starting point is the stream-triggered API (nvshmemx_put_signal_nbi_on_stream), which removes CPU-GPU synchronization barriers with minimal code restructuring. It has two limitations: communication can only begin at a kernel boundary, preventing pack/communication overlap within a pulse, and stream-governed sequencing means the dependent forwarding chain (Z → Y → X) cannot be fused.
We address both with fine-grained per-pulse dependencies in the Replacing host-side communication and Enabling parallelism across pulsessections.
The following figure shows execution diagrams of a simulation timestep with two communication pulses, showing how kernel-initiated RMA, kernel fusion, and event-driven multi-queue execution overlap communication and computation.
*Figure 2. GROMACS NVSHMEM schedule for 2D DD. Overlapping CPU kernel launches and GPU compute streams remove halo exchange from the critical path*
From MPI baseline to optimized GPU-native halo exchange
The existing MPI-based halo exchange relays data through neighboring ranks to minimize communication volume, but at the cost of sequential dependencies across the three phases (Z → Y → X), each containing one or more pulses. The following code shows the coordinate halo-exchange pack kernel and its host-side invocation from the GROMACS source code.
`cuda</code></p><p><code>// GPU kernel: pack boundary atom coordinates into a contiguous send buffer </code></p><p><code>__global__ void packSendBufKernel(float3* dataPacked, // packed coordinate buffer </code></p><p><code> </code><code>const float3* data, // full coordinate array </code></p><p><code> </code><code>const int* map, // index map: which atoms to pack </code></p><p><code> </code><code>const int mapSize) </code></p><p><code>{ </code></p><p><code> </code><code>int threadIndex = blockIdx.x * blockDim.x + threadIdx.x; </code></p><p><code> </code><code>if (threadIndex < mapSize) </code></p><p><code> </code><code>{ </code></p><p><code> </code><code>dataPacked[threadIndex] = data[map[threadIndex]]; </code></p><p><code> </code><code>} </code></p><p><code>}</code></p><p><code>// GPU kernel: unpack received halo forces, scatter back into the force array via index map </code></p><p><code>template<bool accumulate> </code></p><p><code>__global__ void unpackRecvBufKernel(float3* forces, // full force array </code></p><p><code> </code><code>const float3* dataPacked, // received packed forces </code></p><p><code> </code><code>const int* map, </code></p><p><code> </code><code>const int mapSize) </code></p><p><code>{ </code></p><p><code> </code><code>int threadIndex = blockIdx.x * blockDim.x + threadIdx.x; </code></p><p><code> </code><code>if (threadIndex < mapSize) </code></p><p><code> </code><code>{ </code></p><p><code> </code><code>if (accumulate) </code></p><p><code> </code><code>forces[map[threadIndex]] += dataPacked[threadIndex]; </code></p><p><code> </code><code>else </code></p><p><code> </code><code>forces[map[threadIndex]] = dataPacked[threadIndex]; </code></p><p><code> </code><code>} </code></p><p><code>} </code></p><p><code>`
`c</code></p><p><code>// CPU host code: coordinate halo exchange (MPI path)</code></p><p><code>for each dimension d in [Z, Y, X]:</code></p><p><code> </code><code>for each pulse p in dimension d:</code></p><p><code> </code><code>// 1. Launch GPU pack kernel on non-local stream</code></p><p><code> </code><code>packSendBufKernel<<<grid, block, 0, nonLocalStream>>>(</code></p><p><code> </code><code>sendBuf, coords, indexMap, mapSize);</code></p><p><code> </code><code>// 2. CPU blocks until GPU pack completes</code></p><p><code> </code><code>cudaStreamSynchronize(nonLocalStream); // *** BLOCKING ***</code></p><p><code> </code><code>// 3. CPU orchestrates MPI transfer</code></p><p><code> </code><code>MPI_Isend(sendBuf, sendSize, sendRank, ...);</code></p><p><code> </code><code>MPI_Irecv(coords + atomOffset, recvSize, recvRank, ...);</code></p><p><code> </code><code>MPI_Waitall(...); // *** BLOCKING ***</code></p><p><code> </code><code>// Received data lands directly</code></p><p><code> </code><code>// in the coordinate buffer at the correct offset and is consumed</code></p><p><code> </code><code>// by the non-local non-bonded force kernel as-is.</code></p><p><code>`
While these kernels are simple, the bottleneck is the CPU-side control flow that wraps them. Every pulse requires the CPU to block on the GPU, start an MPI transfer, and signal the GPU to resume a serializing sequence for each pulse across all three dimensions.
The force halo exchange follows the same pattern but in the reverse direction (X → Y → Z). The rank that sent coordinates now receives the forces computed on those halo atoms, and forces require a scatter-unpack step through unpackRecvBufKernel because they must accumulate into the correct positions in the full-force array.
Every pulse requires two blocking CPU–GPU synchronizations: one before the MPI invocation to guarantee the pack kernel has completed, and one enforced before the next pulse can consume forwarded data. With a 3D decomposition and one pulse per dimension, that is six blocking waits per time-step for coordinates and six more for forces—12 in total. Each adds latency to the critical path, impacting GROMACS iteration rates.
Replacing host-side communication with NVSHMEM
Building on this idea, the first design folds packing, the remote put, and the completion-wait into a single kernel per communication pulse. The host still issues launches in pulse order, but it drops the CPU–GPU barriers that sat between pulses. Correct execution order is now guaranteed by the NVIDIA CUDA stream alone. The host invocation of the fused pack-put-wait kernel and its implementation are:
`c</code></p><p><code>// CPU host code: one kernel launch per pulse, no CPU-GPU sync in between</code></p><p><code>int pulseOffset = 0;</code></p><p><code>for each dimension d in [Z, Y, X]:</code></p><p><code> </code><code>for each pulse p in dimension d:</code></p><p><code> </code><code>// Pack + send + wait are fused into a single GPU kernel</code></p><p><code> </code><code>packSendBufAndPutNvshmemKernel<<<grid, block, 0, nonLocalStream>>>(</code></p><p><code> </code><code>coords, dataPacked, indexMap, sendSize,</code></p><p><code> </code><code>sendRank, atomOffsetInSendRank,</code></p><p><code> </code><code>signalReceiverRank + pulseOffset, signalCounter, bar, recvSize);</code></p><p><code> </code><code>pulseOffset++;</code></p><p><code>`
The entire pack–put–wait operation is fused into a single device-side launch. Adapted from the GROMACS implementation, the kernel is:
`cuda</code></p><p><code>__global__ void packSendBufAndPutNvshmemKernel(</code></p><p><code> </code><code>float3* coords, float3* dataPacked, int* map,</code></p><p><code> </code><code>int sendSize, int sendRank, int atomOffsetInSendRank,</code></p><p><code> </code><code>uint64_t* signalReceiverRank, uint64_t signalCounter,</code></p><p><code> </code><code>cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)</code></p><p><code>{</code></p><p><code> </code><code>int tid = blockIdx.x * blockDim.x + threadIdx.x;</code></p><p><code> </code><code>if (sendSize > 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// Pack: grid-stride gather into contiguous send buffer</code></p><p><code> </code><code>for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)</code></p><p><code> </code><code>dataPacked[i] = coords[map[i]];</code></p><p><code> </code><code>// Device-scoped barrier: all CTAs must finish packing before put</code></p><p><code> </code><code>auto token = bar->arrive();</code></p><p><code> </code><code>if (blockIdx.x == 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>bar->wait(std::move(token));</code></p><p><code> </code><code>// Put packed data into peer's coordinate array</code></p><p><code> </code><code>nvshmemx_float_put_signal_nbi_block(</code></p><p><code> </code><code>&coords[atomOffsetInSendRank], // peer destination</code></p><p><code> </code><code>dataPacked, // local source</code></p><p><code> </code><code>sendSize * 3, // float count</code></p><p><code> </code><code>signalReceiverRank, signalCounter,</code></p><p><code> </code><code>NVSHMEM_SIGNAL_SET, sendRank);</code></p><p><code> </code><code>}</code></p><p><code> </code><code>}</code></p><p><code> </code><code>// Wait for incoming data from recvRank before exiting</code></p><p><code> </code><code>if (tid == 0 && recvSize > 0)</code></p><p><code> </code><code>nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);</code></p><p><code>}</code></p><p><code>`
This implementation eliminates all CPU-GPU synchronization from the critical path. The CPU only performs kernel launches, which can overlap with GPU execution. All GPU threads (CTAs) arrive at a device-scoped barrier after packing, and block 0 waits on that barrier to ensure all threads have finished before issuing the nvshmemx_float_put_signal_nbi_block call. That call transfers the packed data to the peer and sets signalReceiverRank on the peer to confirm completion; the receiver waits on that signal before consuming the data.
However, note that this version still processes pulses sequentially (stream ordering enforces it), and each pulse still incurs a separate kernel launch.
Optimizing for efficient transfers across NVIDIA NVLink-connected GPUs
In the design described, data always routes through the NVSHMEM transport (nvshmemx_float_put_signal_nbi_block), which works for any interconnect, but when the peer is NVLink-connected, GPUs can load/store directly into each other’s memory. Instead of packing into a local buffer and issuing a put, pack directly into the peer’s coordinate array and eliminate the separate global memory round-trip.
nvshmem_ptr(remotePtr, peerRank) returns a non-null device pointer when the peer is reachable via NVLink, and null otherwise. We query this once in the kernel and branch accordingly, as shown:
`cuda</code></p><p><code>__global__ void packSendBufAndPutNvshmemKernel(</code></p><p><code> </code><code>float3* coords, float3* dataPacked, int* map,</code></p><p><code> </code><code>int sendSize, int sendRank, int atomOffsetInSendRank,</code></p><p><code> </code><code>uint64_t* signalReceiverRank, uint64_t signalCounter,</code></p><p><code> </code><code>cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)</code></p><p><code>{</code></p><p><code> </code><code>int tid = blockIdx.x * blockDim.x + threadIdx.x;</code></p><p><code> </code><code>if (sendSize > 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// Probe for NVLink: non-null means direct store is possible</code></p><p><code> </code><code>float3* remotePtr = (float3*)nvshmem_ptr(coords, sendRank);</code></p><p><code> </code><code>bool isNVLink = (remotePtr != nullptr);</code></p><p><code> </code><code>// NVLink: pack directly into peer's coordinate buffer</code></p><p><code> </code><code>// IB: pack into local staging buffer for later put</code></p><p><code> </code><code>float3* dest = isNVLink ? (remotePtr + atomOffsetInSendRank)</code></p><p><code> </code><code>: dataPacked;</code></p><p><code> </code><code>for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)</code></p><p><code> </code><code>dest[i] = coords[map[i]];</code></p><p><code> </code><code>auto token = bar->arrive();</code></p><p><code> </code><code>if (blockIdx.x == 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>bar->wait(std::move(token));</code></p><p><code> </code><code>if (isNVLink)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// Data already in peer memory — just signal completion</code></p><p><code> </code><code>uint64_t* peerSignal = (uint64_t*)nvshmem_ptr(</code></p><p><code> </code><code>signalReceiverRank, sendRank);</code></p><p><code> </code><code>storeReleaseSysAsm(peerSignal, signalCounter);</code></p><p><code> </code><code>}</code></p><p><code> </code><code>else</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// IB: combined data transfer + signal notification</code></p><p><code> </code><code>nvshmemx_float_put_signal_nbi_block(</code></p><p><code> </code><code>(float*)&coords[atomOffsetInSendRank],</code></p><p><code> </code><code>(float*)dataPacked,</code></p><p><code> </code><code>sendSize * 3,</code></p><p><code> </code><code>signalReceiverRank, signalCounter,</code></p><p><code> </code><code>NVSHMEM_SIGNAL_SET, sendRank);</code></p><p><code> </code><code>}</code></p><p><code> </code><code>}</code></p><p><code> </code><code>}</code></p><p><code> </code><code>if (tid == 0 && recvSize > 0)</code></p><p><code> </code><code>nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);</code></p><p><code>}</code></p><p><code>`
On the NVLink path, the packing loop writes directly into the exact location in the peer’s coordinate array. Once block 0 confirms all CTAs have finished packing using the device-scoped barrier, it signals the peer with a system-scope release store (st.release.sys), which ensures all the preceding direct writes are visible to the peer before it reads the signal. No additional data put operation is needed because the data is already in place.
When nvshmem_ptr returns null, the behavior matches the first design. Data is packed into a local staging buffer and handed off to nvshmemx_float_put_signal_nbi_block, which the NVSHMEM transport delivers over InfiniBand, Slingshot, or other RDMA fabrics.
This optimization highlights an important asymmetry between the two paths. The NVLink path requires many CUDA cores (SMs) issuing stores concurrently to saturate the NVLink bandwidth, whereas the InfiniBand path only needs a single thread (CTA) to initiate a bulk transfer using the NVSHMEM transport. The remaining bottlenecks are per-pulse kernel launch overhead and rigid sequential ordering between pulses, even when some data has no cross-pulse dependencies.
Enabling parallelism across pulses through kernel fusion
The key insight is that not all data in a pulse depends on prior pulses. Most atoms in a pulse can be packed and sent immediately; only the “forwarded” atoms that receive data from a previous pulse must wait. This is achieved by partitioning each pulse’s index map at dependencyAtomOffset:
- Independent atoms (atomIndex < dependencyAtomOffset): Safe to pack and send as soon as the pulse runs
- Dependent atoms (atomIndex ≥ dependencyAtomOffset): Must wait for the relevant previous-pulse signal
With this partition, a single kernel processes all pulses concurrently. Dependent work waits only on the specific per-pulse signal it needs, not on an entire phase.
Dependency-aware packing
The packSubset device function selectively packs atoms by threshold—first for independent atoms, then a second time for dependent atoms after waiting on signals.
`cuda</code></p><p><code>// Templated pack: selectively gather atoms based on threshold</code></p><p><code>// packAll=true: pack everything (no dependency check needed)</code></p><p><code>// packAll=false, packLessThan=true: pack only atomIndex < threshold (independent)</code></p><p><code>// packAll=false, packLessThan=false: pack only atomIndex >= threshold (dependent)</code></p><p><code>template<bool packAll, bool packLessThan></code></p><p><code>__device__ void packSubset(float3* dest,</code></p><p><code> </code><code>float3* data,</code></p><p><code> </code><code>const int* map,</code></p><p><code> </code><code>int sendSize,</code></p><p><code> </code><code>int gridStride,</code></p><p><code> </code><code>int threadIndex,</code></p><p><code> </code><code>int threshold,</code></p><p><code> </code><code>int& hasDependencyAtoms)</code></p><p><code>{</code></p><p><code> </code><code>for (int idx = threadIndex; idx < sendSize; idx += gridStride)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>int atomIndex = map[idx];</code></p><p><code> </code><code>float3 srcVal = data[atomIndex];</code></p><p><code> </code><code>if constexpr (packAll)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>dest[idx] = srcVal;</code></p><p><code> </code><code>}</code></p><p><code> </code><code>else</code></p><p><code> </code><code>{</code></p><p><code> </code><code>bool packNow = packLessThan ? (atomIndex < threshold)</code></p><p><code> </code><code>: (atomIndex >= threshold);</code></p><p><code> </code><code>if (packNow)</code></p><p><code> </code><code>dest[idx] = srcVal;</code></p><p><code> </code><code>else if (packLessThan)</code></p><p><code> </code><code>hasDependencyAtoms++; // track deferred atoms in first pass</code></p><p><code> </code><code>}</code></p><p><code> </code><code>}</code></p><p><code>}</code></p><p><code>`
Pulses are numbered globally in Z→Y→X order. Each waits for all preceding pulse signals before packing its dependent atoms. With one pulse per dimension [z0, y0, x0]: z0 has no dependencies; y0 waits for z0; x0 waits for both. This is what the loop for (int i = currPulse; i > 0; i--) in packHaloCoords implements:
`cuda</code></p><p><code>// Dependency-aware packing: pack independent data immediately,</code></p><p><code>// wait for signals from prior pulses, then pack dependent data</code></p><p><code>__device__ void packHaloCoords(float3* dest,</code></p><p><code> </code><code>int sendSize,</code></p><p><code> </code><code>float3* data,</code></p><p><code> </code><code>const int* map,</code></p><p><code> </code><code>int currPulse,</code></p><p><code> </code><code>uint64_t signalCounter,</code></p><p><code> </code><code>uint64_t* signalReceiverRankCurr,</code></p><p><code> </code><code>int dependencyAtomOffset)</code></p><p><code>{</code></p><p><code> </code><code>int gridStride = blockDim.x * gridDim.x;</code></p><p><code> </code><code>int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;</code></p><p><code> </code><code>if (currPulse > 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// Pass 1: pack atoms below dependency threshold (no waiting)</code></p><p><code> </code><code>int hasDependencyAtoms = 0;</code></p><p><code> </code><code>packSubset<false, true>(dest, data, map, sendSize, gridStride,</code></p><p><code> </code><code>threadIndex, dependencyAtomOffset,</code></p><p><code> </code><code>hasDependencyAtoms);</code></p><p><code> </code><code>// Thread 0 waits for all prior pulses to deliver their data</code></p><p><code> </code><code>if (threadIdx.x == 0)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>for (int i = currPulse; i > 0; i--)</code></p><p><code> </code><code>while (loadRelaxedSys(signalReceiverRankCurr - i) != signalCounter)</code></p><p><code> </code><code>;</code></p><p><code> </code><code>}</code></p><p><code> </code><code>constexpr unsigned fullWarpMask = 0xFFFFFFFF;</code></p><p><code> </code><code>int need_to_wait = __any_sync(fullWarpMask, hasDependencyAtoms > 0);</code></p><p><code> </code><code>__syncthreads();</code></p><p><code> </code><code>// Pass 2: pack the deferred atoms (>= threshold)</code></p><p><code> </code><code>if (need_to_wait)</code></p><p><code> </code><code>{</code></p><p><code> </code><code>int unused = 0;</code></p><p><code> </code><code>packSubset<false, false>(dest, data, map, sendSize, gridStride,</code></p><p><code> </code><code>threadIndex, dependencyAtomOffset,</code></p><p><code> </code><code>unused);</code></p><p><code> </code><code>}</code></p><p><code> </code><code>}</code></p><p><code> </code><code>else</code></p><p><code> </code><code>{</code></p><p><code> </code><code>// Pulse 0: no dependencies, pack everything unconditionally</code></p><p><code> </code><code>int unused = 0;</code></p><p><code> </code><code>packSubset<true, true>(dest, data, map, sendSize, gridStride,</code></p><p><code> </code><code>threadIndex, 0, unused);</code></p><p><code> </code><code>}</code></p><p><code> </code><code>__syncthreads();</code></p><p><code>}</code></p><p><code>`
A practical note: loadRelaxedSys uses PTX ld.relaxed.sys.global.u64—system scope because the signal is written by a remote GPU, relaxed ordering because it’s a spin loop. The __syncthreads() after the wait makes the guarded data visible to all threads. In production code, cuda::atomic_ref<uint64_t, cuda::thread_scope_system> with cuda::memory_order_relaxed from libcu++ provides equivalent semantics.
The fused kernel assigns each pulse to a blockIdx.y row, with all blocks in that row cooperating on one pulse via a grid-stride loop (indexed by blockIdx.x). The pre-computed remote pointer selects the transport at runtime: non-null means NVLink direct stores, null falls back to the NVSHMEM put path (InfiniBand):
`cuda
// One kernel launch replaces all per-pulse pack + send operations
// Grid: (blocksPerPulse, totalNumPulses, 1) — one blockIdx.y per pulse
__global__ void fusedPulsesPackAndSendKernel(
float3* data, // full coordinate array
HaloExchangeData* pulseData, // per-pulse metadata array
uint64_t* signalReceiverRank,
uint64_t signalCounter,
uint32_t* gridSync, // per-pulse block-completion counters
int totalNumPulses,
int dependencyAtomOffset)
{
int currPulse = blockIdx.y;
if (currPulse < totalNumPulses)
{
HaloExchangeData halo = pulseData[currPulse];
// Precomputed by nvshmem_ptr; null => IB/RDMA path
float3* remotePtr = ha
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み