Enhancing the Performance of Learned Index on Processing-in-Memory Architecture

Learned indexes serve as essential index structures within storage systems, yet their efficiency is constrained by the memory bandwidth and latency limitations inherent in traditional computer architectures. Processing-in-memory (PIM) technology presents a promising approach by embedding processing units directly into memory components. In this paper, we introduce PIMLex, a well-designed learned index utilizing PIM to address the challenges posed by memory-bound operations. PIMLex tackles the capacity constraints found in current PIM hardware by adopting a decoupled two-level architecture, which harnesses both the robust data processing power of PIM and the expansive capacity of conventional DRAM. Furthermore, the design incorporates a PIM-optimized model structure to reduce computational demands that are inefficient for PIM. Coupled with a hotness-aware replication strategy that promotes load balancing across multiple PIM modules, PIMLex achieves superior performance under diverse workload scenarios. We have implemented PIMLex on UPMEM, a commercially available PIM platform, demonstrating a throughput improvement of 2.2 × compared to the DRAM-based ALEX.