AMARA: Compiler-Driven Accelerator Architecture Design for Graphical Probabilistic Sampling

Abstract

Sampling methods are fundamental primitives for training and inference on probabilistic models. However, practical deployment of sampling algorithms such as Markov Chain Monte Carlo (MCMC) remains difficult on general-purpose processors. Existing hardware accelerators improve sampling efficiency, but they are typically specialized for a narrow class of models, graph topologies, or update schedules, making it difficult to support the diverse workloads used in real applications. This paper presents AMARA, an end-to-end compiler-driven accelerator framework for probabilistic sampling on arbitrary probabilistic graphical models. AMARA introduces MCMC-IR, a domain-specific intermediate representation that separates local factorized computation from global update dependencies. Based on this representation, the compiler performs architecture-aware partitioning, scheduling, and memory-access optimization to expose parallelism and reduce irregular graph communication. The optimized program is lowered to a compact ISA and executed on a programmable ASIC accelerator with specialized compute, sampling, and memory support. By co-designing the compiler and the architecture, AMARA bridges the gap between model generality and accelerator efficiency for practical MCMC-based probabilistic computing. We evaluate AMARA across representative workloads with different structures, including grid-structured and irregular energy-based models, and causal inference networks. The results show that AMARA supports a broad spectrum of discrete probabilistic models while achieving high sampling throughput and efficient hardware utilization. AMARA achieves an average of 6.1$\times$ (up to 16.3$\times$) higher sampling throughput than a JAX-GPU baseline on irregular workloads, and offers broader programmability than state-of-the-art accelerators.