Samsung LPDDR5X-PIM: Processing-in-Memory Architecture and Implementation

SamsungのLPDDR5X-PIM実装は、LPDDR5Xチップ内に演算処理を直接行うことを可能にし、外部バスのボトルネックを回避して614 GB/sの内部帯域幅を実現します。このアーキテクチャは、DRAMバンク内に積算演算（MAC）ユニットを配置することでAIおよび機械学習のワークロードを加速するように設計されており、標準的なメモリコントローラとの互換性を維持しながら、チップを制約付きのSIMDプロセッサとして機能させることができます。

Hardware Architecture and Throughput

SamsungのLPDDR5X-PIMは、LPDDR5X-9600チップの16個の各バンクにPIMブロックを統合しています。DRAMバンクに内部的にアクセスすることで、これらのブロックは、通常最大76.8 GB/sに制限されるチップの外部インターフェースの制限を回避します。

MAC Unit Specifications

各PIMブロックには、MACツリー、レジスタファイル、および制御ロジックが含まれています：

Instruction Register File: 1024-bit、最大64個の16-bit命令を保持。
Source Register File: アクティベーションベクトルのための4 kbit。
Scale Register: 計算前にモデルの重みをスケーリングするための2 kbit。
Data Flow: モデルの重みは接続されたDRAMブロックに保存され、アクティベーションベクトルはソースレジスタを介して供給されます。

Compute Performance

MACアレイは、INT8およびFP8を含む低精度フォーマットをサポートしています。単一のPIMブロックは、データクロックあたり4つのINT8またはFP8 MAC演算を維持できます（1サイクルあたり8つ）。4-bit入力重みを使用する場合、スループットは倍増し、パッケージ全体で2.4 TOPSの演算スループットを実現します。IntelのMeteor Lakeのような一般的なNPUの性能に匹敵させるには、システムは約8個のLPDDR5X-PIMチップを必要とし、合計128 GBのシステムメモリが必要となります。

Protocol Integration and Control

LPDDR5X-PIMは、特定の行アドレスをMemory-Mapped I/O (MMIO) トリガーとして再利用し、動作モードを切り替えることで、標準的なLPDDR5Xプロトコルとの互換性を維持しています。

Mode Switching

Single-Bank Mode: 通常のDRAMアクセス用の標準的な動作モード。
Multi-Bank Mode: 内部帯域幅を活用するために、全16バンクに対して同時にコマンドを適用します。
PIM Registers Activated Mode: 特殊なバンクごとの行によってトリガーされ、このモードでは読み取りおよび書き込みコマンドをPIMレジスタにリダイレクトします（DRAMバンクの内容ではなく）。

SIMD Execution Flow

マルチバンクモードでは、チップはSIMDプロセッサとして動作します。単一の書き込みコマンドが全16バンクにブロードキャストされ、チップ全体に同じ演算、スケール係数、およびソースオペランドが適用されることを保証します。メモリコントローラの並べ替えを防止するため、Samsungは**Address Align Mode (AAM)**を利用しており、これにより命令はアクセスされる列アドレスからソースレジスタのインデックスを推論できます。

Critical Software and System Challenges

ハードウェアの効率性にもかかわらず、PIMを現代のコンピューティング環境に統合することは、特にCPUがメモリをどのように扱うかに関して、深刻なアーキテクチャ上の衝突を引き起こします。

Cache Coherency and Speculation

PIM演算はメモリの内容とレジスタの状態を内部的に変更するため、CPUキャッシュ階層の根本的な仮定を崩します：

Uncacheable Memory: SamsungはPIMメモリをuncacheableとしてマッピングすることを推奨しています。これにより、CPUキャッシュの利点が消失し、レイテンシが大幅に増加し、コアがストールします。 n- Speculative Execution: 現代のCPUはプリフェッチャと分岐予測を使用してデータを投機的にロードします。PIMシステムにおいて、投機的な読み取りは受動的な操作ではなく、PIM Vector Register File (VRF) を変更する演算をトリガーし、実際のプログラムの状態を破壊する可能性があります。

Multitasking and OS Integration

PIMモードの切り替えは、メモリチャネルのグローバルな状態変化です。これは、マルチタスクOSにとって大きな障害となります：

Thread Isolation: あるスレッドがPIMを使用している間に別のスレッドが通常のメモリ操作を行う場合、非PIMスレッドが意図せずPIM演算をトリガーしたり、VRFデータを誤ったアドレスに書き込んだりする可能性があります。
Context Switching: PIMスレッドのプリエンプション（中断）には、OSがチャネルをPIMモードから切り替え、全バンクにわたって命令、ソース、スケール、およびベクトルレジスタファイルのステートを、手動で保存する必要があります。
Memory Interleaving: PIM領域を分離するため、システムはチャネル間のアドレス・インターリービングを無効にする必要がある場合があります。これは、非PIMアプリケーションの用の全般的なメモリ帯域幅を減少させます。

Technical Perspectives and Alternatives

業界の議論では、PIMが「Von Neumann bottleneck」に対処する一方で、データ配置に厳格な制約を課すことが強調されています。コミュニティの貢献者によって指摘されているように：

"The tradeoff with putting the compute in the memory is that you have to know exactly where the dependent information will be at all times. Most problems do not fit this pattern very well."

これらの問題を解決するために、著者は、DRAMインターフェースを専用の演算コマンドで拡張したり、メモリコントローラをキャッシュコヒーレンシのためにCPUコアのピアとして動作させる（Read-For-Ownershipリクエストを使用）メモリコントローラを実装したり、ソフトウェア層からPIMハードウェアを抽象化するために透明なCPU命令（例：仮説的な rep macb）を導入したりするなど、より実行可能な道筋にはシステム的なハードウェア変更が必要であると示唆しています。