AMARA hands MCMC sampling to the compiler, letting arbitrary discrete graphical models run on a programmable accelerator
Synopsis
AMARA splits MCMC sampling over discrete probabilistic graphical models into an intermediate representation the compiler can process, then maps it onto a programmable ASIC. Previously CPUs and GPUs raised throughput by replicating independent Markov chains, but replicating 4096 chains yielded only about a 270-fold gain, while specialized accelerators often assumed a fixed topology or update schedule. Evaluated on grid-structured and irregular energy-based models and causal inference networks, AMARA achieves an average (up to) higher sampling throughput than a JAX-GPU baseline on irregular workloads and supports a broader range of discrete models than state-of-the-art accelerators.
Interpretation
AMARA lets one compilation flow cover discrete graphical models such as grid-structured and irregular energy-based models and causal inference networks, so sampling programs no longer need hardware designed per topology. CPUs and GPUs stay flexible, but their common remedy is replicating independent Markov chains, and replicating 4096 chains yields only about a 270-fold throughput gain without shortening an individual chain's convergence path; specialized accelerators often assume a fixed topology, factor family or update schedule. Evaluated on 11 representative probabilistic graphical models, including Pegasus, Zephyr, Alarm, Water, Andes, Gset's G6, G22 and G49, Queen13, and uf100 and uf250; on irregular workloads AMARA achieves an average (up to) higher sampling throughput than the JAX-GPU baseline.
The compiler front end splits each variable update into local factorized computation and the dependency constraints between updates, and the back end derives a static placement and phase schedule from that split. MCMC-IR is a hardware-oriented domain-specific intermediate representation: a Factor Forest describes what each update kernel computes, and a Dependency Graph describes which kernels cannot be updated concurrently. The back end applies weighted METIS partitioning, then partition-aware DSATUR coloring to place independent kernels in the same color phase; on Gset instances, partition-aware iterative coloring improves throughput by 27.7% at 16 partitions.
The accelerator uses energy units to compute candidate-label log-likelihoods and sampler units to perform Gumbel-Max sampling directly, with label storage organized hierarchically into block-local and shared global memories. A conventional softmax sampler needs exponential evaluation, overflow prevention, normalization and cumulative-distribution construction, whose serial steps limit pipelining; fully crossbar-interconnected distributed register banks also carry higher communication complexity. The design is synthesized in TSMC 5FFP with 128 energy units and 32 sampler units, each energy unit having two SIMD lanes producing 16-bit fixed-point energy values; the hierarchical label memory lowers communication complexity from the full-crossbar level and uses a read-once policy that fetches each variable from global memory at most once per sweep.
Perspective
The evaluation covers 11 discrete probabilistic graphical models, including grid-structured and irregular energy-based models and causal inference networks; throughput is measured by cycle-accurate RTL simulation, and the hardware is synthesized in TSMC 5FFP. The performance results assume the model fits in on-chip SRAM, with off-chip DRAM access required beyond that capacity.
AMARA is currently limited to discrete, component-wise MCMC and recompiles its static mapping when model topology changes; the modeled off-chip bandwidth requirement reaches up to 573.44 GB/s assuming full hardware utilization, actual demand on the measured workloads is lower, and Water's large data table limits prefetch effectiveness.