Technical Summary: DCC – Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architectures

1. Problem Statement

Processing-In-Memory (PIM) architectures offer a promising solution to the memory bottleneck in Machine Learning (ML) by placing processing units near DRAM banks. However, a fundamental mismatch exists between the data layout requirements of high-performance Host processors (e.g., GPUs) and PIM cores.

• Host Requirement: To exploit bank-level parallelism and high bandwidth, Host processors require consecutive data elements to be distributed across multiple DRAM banks.
• PIM Requirement: PIM cores can only access data from their local banks. To maximize local bandwidth and enable efficient multi-element access, consecutive elements must be placed within the same bank.

This necessitates data rearrangements (input and output) performed via the Host memory bus, which incurs significant data movement costs that can dominate end-to-end execution time. Current compilation approaches for PIM systems suffer from two main limitations:

1. Lack of Systematic Optimization: Existing tools often target specific backends (e.g., UPMEM) or specific kernels, lacking support for diverse ML workloads and multiple PIM devices (e.g., Samsung HBM-PIM, SK Hynix GDDR6-AiM).
2. Compute-Centric Tuning: Existing compilers (e.g., ATiM) optimize compute schedules independently of data rearrangement costs. The authors demonstrate that compute transformations and data rearrangements are interdependent; optimizing them in isolation yields suboptimal performance because a compute-efficient schedule may require expensive data rearrangements, while a less efficient compute schedule might enable cheaper data movement.

2. Methodology: DCC

The authors propose DCC, the first data-centric ML compiler for PIM systems. DCC jointly co-optimizes data rearrangement strategies and loop-level compute schedules within a unified tuning process. The system comprises four key components:

2.1 Multi-Layer Abstraction

DCC introduces a generic abstraction that maps the PIM memory hierarchy into a compute hierarchy, decoupling the compiler from backend-specific semantics. This abstraction operates at three levels:

• System Level: The Host xPU (e.g., GPU), Host memory space, and PIM devices.
• PIM Group Level: Groups of PIM cores (and their local banks) belonging to the same memory channel or rank. This allows for coarse-grained optimizations and group-level instruction broadcasting.
• PIM Core Level: Individual processing units and their local banks, enabling fine-grained optimizations for memory access and local data layout.

2.2 Data-Centric Schedule Generator

Unlike traditional compilers that start with loop transformations, DCC adopts a data-first strategy:

1. Dimension Set Construction: Tensor dimensions are correlated with loop variables to form "complex dimension sets."
2. Data-Tile Construction: The system exhaustively generates candidate data partitions (tiles) across PIM groups and cores.
3. Draft Pruning: Three rules eliminate unpromising drafts:
   • Redundancy: Removing drafts with identical data partitions and mappings.
   • SIMD Alignment: Pruning drafts where partition sizes do not align with the backend's SIMD width (e.g., 16-way).
   • Performance Equivalence: Keeping only one draft if multiple drafts yield the same worst-case per-core performance.
4. Mapping: The remaining drafts are mapped to compute tiles, determining loop ranges and tensor indices for each core.

2.3 PIM-Specific Code Optimizer

This component applies optimizations tailored to PIM constraints:

• Data Rearrangement Optimization: Exploits controllable on-chip memory (e.g., GPU shared memory) to reorganize data. It performs serial reads from Host memory channels followed by parallel writes to PIM channels, maximizing channel parallelism for both read and write operations.
• Compute Code Optimization: Utilizes hierarchical command generation. It prioritizes group-level SIMD-style commands (broadcasting a single instruction to all cores in a group) to reduce command bus overhead, falling back to bank-level commands only when necessary.

2.4 Coupled Performance Predictor

To avoid the prohibitive cost of profiling thousands of candidate drafts, DCC employs a learning-based predictor (XGBoost).

• Function: It jointly estimates data rearrangement costs and compute execution times to select the best-performing end-to-end configuration.
• Training: The model is trained offline on a diverse subset of drafts.
• Runtime: For static tensor sizes, the best schedule is pre-selected. For dynamic sizes (e.g., in LLMs), the predictor estimates performance on-the-fly.

2.5 Programming Interface

DCC provides a Python interface compatible with PyTorch. Users annotate kernels with @DCC_kernel. The compiler parses these high-level definitions, generates the Intermediate Representation (IR) involving data and compute tiles, and emits low-level backend instructions.

3. Key Results

The authors evaluated DCC on two state-of-the-art PIM backends: Samsung HBM-PIM and AttAcc (a research prototype with specialized GEMV/Softmax units), comparing against GPU-only execution and baseline PIM implementations.

• Individual ML Kernels:
  • HBM-PIM: DCC achieves up to 7.68× speedup (2.21× average) over GPU execution.
  • AttAcc: DCC achieves up to 13.17× speedup (3.92× average) over GPU execution.
  • Comparison to Baseline PIM: DCC improves upon the baseline AttAcc implementation by up to 2.74× (1.75× average) and HBM-PIM by up to 2.09× (1.53× average) for specific kernels like GEMV and RED.
• End-to-End LLM Inference:
  • On AttAcc, DCC accelerates GPT-3 and LLaMA-2 inference by an average of 4.52× (up to 7.71× for LLaMA-2) over GPU-only execution.
  • DCC significantly outperforms the baseline AttAcc implementation (which uses fixed tiling) by enabling efficient execution of QKV generation and projection layers, which were previously bottlenecked on the GPU.
• Energy Efficiency: DCC reduces energy consumption by an average of 1.24× compared to the baseline AttAcc system due to reduced data movement.
• Compilation Overhead: The offline compilation time is low (seconds), and the pruning scheme reduces the search space by ~9×. The coupled predictor achieves high accuracy (97.37% on LLMs) and reduces compilation time by ~3.8× compared to exhaustive search.

4. Significance and Claims

The paper claims DCC is the first data-centric ML compiler for PIM systems that:

1. Jointly Optimizes data rearrangement and compute schedules, demonstrating that these two factors are interdependent and must be tuned together to avoid suboptimal performance.
2. Supports Multiple Backends through a multi-layer abstraction, enabling portability across diverse PIM architectures (e.g., HBM-PIM, AttAcc) without requiring manual, backend-specific tuning.
3. Enables High Performance for memory-intensive ML kernels and large-scale LLM inference, significantly outperforming both GPU-only execution and existing PIM compilation approaches.
4. Improves Programmability by providing a high-level Python interface that abstracts away the complexity of low-level data layout management and instruction scheduling for PIM devices.

The authors conclude that DCC provides a practical and effective path for heterogeneous ML acceleration, balancing high performance with programmability in emerging PIM ecosystems. The code is open-sourced to facilitate further research.