Compiler speeds up memory aware processing for dynamic programming
Compiling Semi-Ring Dynamic Programming to Tier-Aware 3D-DRAM Processing-in-Memory
Hardware Architecture
Summary
Dynamic programming tasks like finding shortest paths or aligning sequences need lots of memory and computing. Using special 3D memory chips with built-in processors can speed this up, but programming them is hard and slow. The authors created GenMLIR, a compiler that automatically converts high-level math formulas into efficient programs that run well on these memory chips. Their tool makes programs much faster and much shorter, taking away the need for hand-tuning.
What this means in practice
- •For chip designers: Generate optimized code automatically for 3D DRAM processors to accelerate dynamic programming workloads with memory-aware scheduling.
- •For bioinformatics developers: Build faster sequence alignment tools by compiling high-level algorithms to run efficiently on tiered memory processing hardware.
Authors
Mahbod Afarin, Tsung-Han Lu, Tajana Rosing
Abstract
Processing-in-Memory (PIM) on monolithic 3D (M3D) DRAM is a promising answer to the memory wall for data-intensive dynamic programming (DP), yet extracting its performance today demands hand-written kernels: the programmer must pick a tile size, place data across non-uniform-latency memory tiers, partition work between heterogeneous processing units, and insert the right broadcasts and queues by hand. We present GenMLIR, an MLIR compiler that automates this mapping for tier-aware 3D PIM. GenMLIR encodes the semi-ring generalized grid update, the algebraic form shared by all-pairs shortest path (APSP) and sequence alignment, as a first-class IR abstraction, and lowers it through four GenDRAM-aware pass groups that derive blocked tiling, tier-aware placement, tile-to-PU assignment, and explicit communication. We also characterize precisely which DP recurrences the abstraction admits and which it does not. On the GenDRAM architecture, GenMLIR-generated code runs up to 5.8 (APSP) and 16.7 (alignment) faster than a lowering without compiler support, and 1.2-1.4 faster than a standard affine-tiling PIM compiler we also implement, reaching 90-100% of an achievable-performance bound. It expresses these workloads in 4-9 lines instead of 300-500 and compiles in a negligible fraction of runtime.