科学・技術
サムスンのプロセッシング・イン・メモリ(PIM)
Samsung's Processing-in-Memory (PIM) (chipsandcheese.com)
要約
サムスンは、LPDDR5Xメモリチップ内にMAC(Multiply-Accumulate)演算ユニットを組み込んだプロセッシング・イン・メモリ(PIM)技術を開発しました。これにより、従来のDRAMと演算コア間のデータ転送遅延を回避し、メモリチップ内部の広帯域幅を最大限に活用することで、AI/MLワークロードにおける大幅な性能向上が期待されます。
全文翻訳
サムスンのプロセッシング・イン・メモリ(PIM)
インメモリコンピューティングは、メモリチップ内の演算がその高い内部帯域幅を利用できるため、長年にわたり魅力的な提案となってきました。さらに、インメモリコンピューティングはDRAMと従来の演算コア間の長い遅延パスを回避します。Hot Chips 2026で、サムスンはPIM(Processing-in-Memory)の推進によるインメモリコンピューティングの継続的な追求について議論しました。彼らは、チップの標準的なメモリコントローラーとのインターフェース能力を維持しながら、LPDDR5Xチップ内にMACユニットを実装しています。
DRAMチップは内部的にバンクに分割されており、それぞれが独自の読み書きロジックを持っています。通常のDRAMアクセス中、メモリコントローラーはバンクを選択し、その中の行をアクティブ化し、カラムアクセスストローブ(CAS)コマンド経由でデータにアクセスします。帯域幅は、チップの外部DRAMインターフェースによって制限されます。メモリコントローラーがすべてのバンクを同時にアクティブ化できたとしても、すべてのバンクにわたる完全な帯域幅をすべて利用することはできません。
サムスンのLPDDR5X-PIMは、16個のバンクを持つ通常のLPDDR5X-9600チップに似ていますが、各バンクにPIM(Processing-in-Memory)ブロックを配置しています。これらのPIMブロックは、チップの外部バスに制約されることなく、接続されたDRAMバンクにアクセスします。これらを合わせることで、16個のバンクすべてにわたるチップの内部帯域幅を利用でき、これは614 GB/sに達します。比較として、通常のDRAMアクセスは2つのバンクを並列にヒットでき、最大76.8 GB/sに達します。
PIMブロックは内部的に、レジスタファイルと制御ロジックに囲まれたMACツリーで構成されています。1024ビットの命令レジスタファイルには、最大64個の16ビット命令が格納されます。4 kbitのソースレジスタファイルはアクティベーションベクトル用であり、MACアレイに1つのソースオペランドを供給します。サムスンは、ソフトウェアがモデルウェイトをDRAMにロードすることを想定しているため、接続されたDRAMブロックが2番目のオペランドを供給します。モデルウェイトは、MAC演算の前にスケーリングでき、スケールファクターは2 kbitのスケールレジスタから取得されます。
PIMブロックのMACアレイは、さまざまな低精度フォーマットをサポートします。サムスンのプレゼンテーションからの数値は、各PIMブロックのMACアレイがデータクロックあたり4つのINT8またはFP8 MAC演算を維持できること、またはダブルデータレートをカウントしない場合はサイクルあたり8つを維持できることを示唆しています。4ビット入力ウェイトの場合、スループットは倍増し、パッケージ全体の演算スループットは2.4 TOPSに達します。
これは非常に高い数値ではありませんが、多くのLPDDR5Xチップを使用した実装では、より高い集約スループットが得られます。例えば、8個のLPDDR5Xチップを組み合わせると、9.6 INT8 TOPSとなり、これはIntelのMeteor LakeのNPUとほぼ同等です。しかし、8個の16 GB LPDDR5Xチップは128 GBのシステムメモリに相当するため、高価なセットアップになります。
標準DDRコマンドによる演算へのアクセス
LPDDR5X-PIMのハイライトの1つは、メモリ標準の一部ではない演算能力を公開しながら、標準のLPDDR5Xプロトコル内に留まっていることです。サムスンは、MMIOアドレスのような役割を果たす特別な行アドレスを確保することでこれを実現しています。各チャネルには、モード制御用の2つの定義済み行があります。これらの行のいずれかをアクティブ化すると、チップはシングルバンクモードに設定され、もう一方のアクティブ化でマルチバンクモードに設定されます。シングルバンクは通常のモードであり、マルチバンクは16個のバンクすべてにコマンドを適用してチップの内部帯域幅を利用します。
バンクごとの特別な行は、読み書きコマンドの動作を変更します。これらの特別な行のいずれかをアクティブ化すると、読み書きコマンドは通常のDRAMバンクの内容ではなく、PIMレジスタにアクセスします(PIMレジスタアクティブ化モード)。サムスンは、ソフトウェアが通常のシングルバンクモードでモデルウェイトをDRAMにロードし、その後マルチバンクモードに切り替えてPIMレジスタアクティブ化モードに入るMLユースケースを想定しています。これにより、コードはアクティベーション値をPIMソースレジスタに書き込み、スケールファクターをPIMスケールレジスタに設定し、PIM命令レジスタに記入される演算を指定できます。
チップはマルチバンクモードであるため、各PIMレジスタへの書き込みは16個のバンクすべてにブロードキャストされます。したがって、PIM演算は非常に制約されたSIMDプロセッサのように機能し、演算、スケールファクター、および1つのソースオペランドはすべてのバンクで同じです。サムスンはシングルバンクモードでのPIMレジスタへの書き込みも許可していますが、その機能はデバッグ目的です。各DRAMパケットは256ビット(BL=16)です。各ソースレジスタを埋めるには16回の書き込みコマンドが必要です。これを16個のバンクすべてで1バンクずつ行うと、256回の書き込みコマンドが必要になり、ホストからPIMレジスタへの書き込み帯域幅がボトルネックになります。サムスンは実際にはシングルバンクモードでのPIMレジスタアクセスを許可していますが、それはデバッグ機能として意図されています。
PIMレジスタを準備した後、ソフトウェアはマルチバンクモードに戻り、読み込みコマンドを発行します。DRAMの内容を読み取る代わりに、これらの読み込みコマンドは演算を開始し、結果をPIMベクトルレジスタファイルに蓄積します。その後、書き込みコマンドはPIMブロックにVRFの内容をDRAMバンクに書き込むように指示します。
PIMは、通常のメモリコントローラーが行う可能性のある並べ替えを処理する必要があります。コードがバンクをアクティブ化してPIMをセットアップする際、PIMは通常、命令がソースレジスタ要素に順番にアクセスするように命令レジスタファイルをセットアップします。例えば、最初の命令は最初のソースレジスタ要素を参照し、2番目の命令は2番目のソースレジスタ要素を参照します。しかし、メモリコントローラーがアクセスを並べ替えると、これはうまくいきません。サムスンは、各命令がアクセスされているカラムアドレスからソースレジスタインデックスを推測するアドレスアラインモード(AAM)でこれを回避します。
ホストがインメモリコンピューティングの使用を終えて結果を読み取りたい場合、DRAMチップをシングルバンクモードに戻します。その後、通常のDRAMの読み書きは通常通りDRAMの内容にアクセスし始めます。
ソフトウェアの頭痛の種?
サムスンは、標準LPDDR5Xを使用する場合と比較して、LPDDR5X-PIMを利用することで、社内で大幅なパフォーマンス向上を達成しました。標準メモリコントローラーで動作するチップの能力は印象的であり、サムスンはこの問題へのアプローチにおいて非常に創造的でした。
標準DRAMコマンドを再利用することはハードウェアを単純化するはずですが、ソフトウェアの課題は急峻に見えます。PIMモードはDRAMアクセスコマンドの意味を変更するため、ソフトウェアはPIMを使用しながら通常のメモリアクセスを実行することはできません。これはスレッド間でも適用されます。なぜなら、メモリコントローラーとDRAMチップは、アクセスがどのスレッド用であるかを認識しないからです。PIMを使用している間に別のスレッドがメモリから読み取ると、最初のスレッドが意図しない演算を引き起こし、PIM VRFに誤った結果をもたらす可能性があります。非PIMスレッドからの書き込みは、PIMブロックがVRFデータを誤ったアドレスに書き戻す原因となる可能性があります。
サムスンは、ホストがメモリ内のPIM領域を分離することでこれに対処しています。メモリ帯域幅とPIMパフォーマンスを損なうことなく、典型的なシステムでこれを容易に行う方法を私は思いつきません。ハードウェアは通常、チャネル間でアドレスをインターリーブしますが、これにより一般的なアクセスパターンがそれらのチャネル全体で帯域幅を自然に利用できます。PIMはシングル/マルチバンクモードの変更を制御するためにチャネルごとの行を使用するため、インターリーブをドロップし、メモリチャネルをPIM専用として指定することが、PIM領域を作成する唯一の合理的な方法でしょう。その場合、非PIMアプリケーションは、PIM専用に予約されたチャネルからの帯域幅を利用できなくなります。PIMコードは、非PIMチャネルからの帯域幅と演算を逃すことになります。後者は、チップごとの演算スループットがあまり高くないため、重大な問題となる可能性があります。
PIM領域を分離した後でも、マルチタスクの問題は残る可能性があります。アプリケーションがPIMを使用し、マルチスレッドを活用したい場合、あるスレッドがPIM演算を実行しようとしている間に別のスレッドが通常のメモリアクセスを試みるケースを防ぐために、PIM領域へのアクセスをロックで保護する必要があります。現代のマルチタスクオペレーティングシステムでは、複数のプロセスがメモリにアクセスするため、状況はさらに悪化します。