跳到主要内容
返回时间线
arXiv来源发表:

AMARA把MCMC采样交给编译器:任意拓扑离散图模型也能上专用加速器

核心概要

AMARA把离散概率图模型上的MCMC采样拆成编译器可处理的中间表示,再映射到可编程专用芯片上执行。此前CPU和GPU靠复制多条马尔可夫链提高吞吐,但复制4096条链只带来约270倍增益,专用加速器又常绑定固定拓扑或更新调度。AMARA在网格结构、不规则能量模型和因果推断网络上评测,在不规则负载上相对JAX-GPU基线取得平均(最高)更高的采样吞吐,并比现有加速器支持更广的离散模型范围。

Source-provided article image: AMARA: Compiler-Driven Accelerator Architecture Design for Graphical Probabilistic Sampling
arXiv

深度剖析

AMARA让同一套编译流程覆盖网格结构、不规则能量模型和因果推断网络等离散概率图模型,采样程序不再需要为每种拓扑单独设计硬件。 CPU和GPU保留灵活性,但常用手段是复制相互独立的马尔可夫链,复制4096条链只换来约270倍吞吐增益,单条链的收敛路径并未缩短;专用加速器则往往假定固定拓扑、因子族或更新调度。 在11个代表性概率图模型上评测,包括Pegasus、Zephyr、Alarm、Water、Andes、Gset的G6、G22、G49、Queen13以及uf100、uf250;在不规则负载上,AMARA相对JAX-GPU基线取得平均(最高)更高的采样吞吐。

编译前端把每次变量更新拆成局部因子计算和更新之间的依赖约束两部分,后端据此生成静态的放置与分阶段调度。 MCMC-IR是面向硬件的领域专用中间表示,因子森林描述每个更新核计算什么,依赖图描述哪些核不能同时更新。 后端用METIS做加权划分,再用考虑划分的DSATUR着色把独立核分到同一颜色阶段;在Gset实例上,16个划分时划分感知的迭代着色把吞吐提高27.7%。

加速器用能量单元计算候选标签的对数似然,用采样单元直接完成Gumbel-Max采样,标签存储按块本地加全局共享分层组织。 传统softmax采样需要指数运算、防溢出、归一化和累积分布函数构造,串行环节限制流水;全交叉互连的分布式寄存器组通信复杂度也更高。 设计在TSMC 5FFP工艺下综合,含128个能量单元和32个采样单元,每单元两条SIMD通道输出16位定点能量值;分层标签存储把通信复杂度从全互连的水平降下来,并采用每次扫描每个变量最多从全局存储读取一次的读一次策略。

启示与展望

评测覆盖11个离散概率图模型,包括网格结构、不规则能量模型和因果推断网络,吞吐由周期精确的RTL仿真测得,硬件在TSMC 5FFP工艺下综合;性能结论假定模型能放进片上SRAM,超出容量时需访问片外DRAM。

AMARA目前限于离散、逐分量的MCMC,模型拓扑变化时需要重新编译其静态映射;片外带宽需求按满硬件利用率建模最高可达573.44 GB/s,实测负载的实际需求更低,而Water的大数据表会限制预取效果。