SoMa: Identifying, Exploring, and Understanding the DRAM Communication Scheduling Space for DNN Accelerators

TL;DR

The paper tackles the DRAM bandwidth bottleneck in DNN accelerators by moving beyond single-layer dataflow optimization and introducing SoMa, a framework that jointly explores Layer Fusion and Prefetching/Delayed Storing within a Tensor-centric DRAM Communication Scheduling Space. It defines a six-attribute encoding (LFA and DLSA) and uses a two-stage Simulated Annealing-based exploration with a Buffer Allocator to coordinate global scheduling decisions, supported by an accurate evaluator and a compiler flow. Empirically, SoMa achieves an average 2.11x performance improvement and 37.3% energy reduction over the SOTA Cocco, with Stage 2 approaching the theoretical optimum and revealing nuanced insights about decode-stage limits and batch-size effects. The work demonstrates that smart buffer usage, fusion flexibility, and timing-aware DRAM scheduling can substantially reduce DRAM traffic and stall, offering practical guidance for designing DRAM-efficient DNN accelerators and enabling broader exploration of the DRAM scheduling space in both CNNs and LLMs.

Abstract

Modern Deep Neural Network (DNN) accelerators are equipped with increasingly larger on-chip buffers to provide more opportunities to alleviate the increasingly severe DRAM bandwidth pressure. However, most existing research on buffer utilization still primarily focuses on single-layer dataflow scheduling optimization. As buffers grow large enough to accommodate most single-layer weights in most networks, the impact of single-layer dataflow optimization on DRAM communication diminishes significantly. Therefore, developing new paradigms that fuse multiple layers to fully leverage the increasingly abundant on-chip buffer resources to reduce DRAM accesses has become particularly important, yet remains an open challenge. To address this challenge, we first identify the optimization opportunities in DRAM communication scheduling by analyzing the drawbacks of existing works on the layer fusion paradigm and recognizing the vast optimization potential in scheduling the timing of data prefetching from and storing to DRAM. To fully exploit these optimization opportunities, we develop a Tensor-centric Notation and its corresponding parsing method to represent different DRAM communication scheduling schemes and depict the overall space of DRAM communication scheduling. Then, to thoroughly and efficiently explore the space of DRAM communication scheduling for diverse accelerators and workloads, we develop an end-to-end scheduling framework, SoMa, which has already been developed into a compiler for our commercial accelerator product. Compared with the state-of-the-art (SOTA) Cocco framework, SoMa achieves, on average, a 2.11x performance improvement and a 37.3% reduction in energy cost simultaneously. Then, we leverage SoMa to study optimizations for LLM, perform design space exploration (DSE), and analyze the DRAM communication scheduling space through a practical example, yielding some..(more)

Figures (8)

• Figure 1: DNN Accelerator Template
• Figure 2: A Practical Layer-fusion Group ($LG$) Example
• Figure 3: (a) and (b) show the normalized DRAM access and the normalized operation number for each layer in ResNet-50 and Transformer-Large, respectively (each point represents a layer). (c) and (d) show the normalized DRAM access and the normalized operation number for each smallest computing unit (Tile) of ResNet-50 and Transformer-Large, respectively, scheduled using the SOTA Cocco Framework (each point represents a Tile). The darker the color, the more identical overlapped points there are. The normalization method involves dividing the value of each point by the maximum value among all points (DRAM access and operations are independently normalized). We use the default edge accelerator and batch size 1, as introduced in Sec. \ref{['sec:setup']}
• Figure 4: Parsing an Example Encode of a Five-Layer Network into Actual Scheduling Schemes. The DRAM, COMPUTE, and BUFFER in the right part show the DRAM access, workload computation, and buffer usage, respectively. We use $M_{i}$ to represent the $i$th tile of layer $M$, and $I/O_{Mi}$ to represent the ifmaps and ofmaps of $M_{i}$. $W_{M}$ represent the weights of layer $M$. In the BUFFER, blocks with the same background color represent the same data.
• Figure 5: SoMa Framework