Samsung LPDDR5X-PIM: Processing-in-Memory 架构与实现

Samsung 的 Processing-in-Memory (PIM) 在 LPDDR5X 芯片上的实现允许计算操作直接在内存中进行，从而绕过外部总线瓶颈，实现 614 GB/s 的内部带宽。该架构旨在通过在 DRAM bank 中放置乘累加 (MAC) 单元，加速 AI 和机器学习工作负载，使芯片能够在保持与标准内存控制器兼容的同时，充当受限的 SIMD 处理器。

硬件架构与吞吐量

Samsung 的 LPDDR5X-PIM 将 PIM 模块集成到 LPDDR5X-9600 芯片的 16 个 bank 中。通过在内部访问 DRAM bank，这些模块避开了芯片外部接口的限制，该接口通常最高仅为 76.8 GB/s。

MAC 单元规格

每个 PIM 模块包含一个 MAC 树、寄存器文件和控制逻辑：

Instruction Register File: 1024-bit，可容纳多达 64 个 16-bit 指令。
Source Register File: 4 kbit，用于激活向量。
Scale Register: 2 kbit，用于在计算前对模型权重进行缩放。
Data Flow: 模型权重存储在连接的 DRAM 模块中，而激活向量通过源寄存器提供。

计算性能

MAC 阵列支持低精度格式，包括 INT8 和 FP8。单个 PIM 模块在每个数据时钟下可以维持四个 INT8 或 FP8 MAC 操作（每个周期八个）。使用 4-bit 输入权重时，吞吐量会翻倍，从而实现 2.4 TOPS 的全封装计算吞吐量。为了匹配典型 NPU (例如 Intel 的 Meteor Lake) 的性能，系统将需要大约八个 LPDDR5X-PIM 芯片，总计 128 GB 的系统内存。

协议集成与控制

LPDDR5X-PIM 通过将特定的行地址重新用作内存映射 I/O (MMIO) 触发器，在操作模式之间进行切换，从而保持与标准 LPDDR5X 协议的兼容性。

模式切换

Single-Bank Mode: 标准操作模式，用于常规 DRAM 访问。
Multi-Bank Mode: 同时对所有 16 个 bank 应用命令，以利用内部带宽。
PIM Registers Activated Mode: 由特殊的每 bank 行触发，此模式将读写命令重定向到 PIM 寄存器，而不是 DRAM bank 内容。

SIMD 执行流

在 multi-bank 模式下，芯片作为一个 SIMD 处理器运行。单个写命令会被广播到所有 16 个 bank，确保整个芯片应用相同的操作、缩放因子和源操作数。为了防止内存控制器重排序带来的问题，Samsung 使用了 Address Align Mode (AAM)，它允许指令从正在访问的列地址中推断其源寄存器索引。

关键软件与系统挑战

尽管硬件效率很高，但将 PIM 集成到现代计算环境中会引入严重的架构冲突，特别是关于 CPU 如何处理内存。

Cache Coherency 与 Speculation

由于 PIM 操作会修改内存内容和寄存器状态，它们破坏了 CPU 缓存层级的基本假设：

Uncacheable Memory: Samsung 建议将 PIM 内存映射为 uncacheable。这消除了 CPU 缓存的益处，显著增加了延迟并导致核心停顿。
Speculative Execution: 现代 CPU 使用预取器和分支预测器来加载数据。在 PIM 系统中，投机性读取并非被动操作——它会触发计算，从而修改 PIM Vector Register File (VRF)，可能导致实际程序的程序状态被破坏。

多任务处理与 OS 集成

PIM 模式切换是内存通道的全局状态变更。这为多任务操作系统带来了重大障碍：

Thread Isolation: 如果一个线程使用 PIM while 另一个线程执行常规内存访问，则非 PIM 线程可能会无意中触发 PIM 计算，或将 VRF 数据写入错误的地址。
Context Switching: 抢占 PIM 线程需要操作系统将通道切换出 PIM 模式，并手动保存所有 bank 中指令、源、缩放和向量寄存器文件的状态。
Memory Interleaving: 为了隔离 PIM 区域，系统可能需要禁用跨通道的地址交织，这会降低非 PIM 应用的整体内存带宽。

技术视角与替代方案

行业讨论强调，虽然 PIM 解决了 "Von Neumann 瓶颈"，但它对数据放置提出了严格的约束。正如社区贡献者所言：

"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 请求），以及引入透明的 CPU 指令（例如，假设的 rep macb）来抽象 PIM 硬件与软件层。