Hardware-Software Co-design for
Broad Acceleration on Commercial PIM Architectures

The solution of partial differential equations (PDEs) is critical to many large-scale problems in HPC systems. One such use case, wave simulation, requires solving the wave equation to model the propagation of waves through different media and is used extensively in domains including medical imaging, earthquake modeling, oil and gas exploration, and antenna and radar modeling.

The Discontinuous Galerkin Method (DGM) is a popular algorithm for wave simulation due to its scalability [49]. Like many PDE solvers, DGM discretizes the wave space into a mesh of elements which are distributed among processors in the system (Figure 2a). It then iteratively executes a volume computation, a flux computation, along with communication and support computations to model wave propagation. The volume computation (termed wavesim-volume primitive) performs computations local to each mesh element; the flux computation (termed wavesim-flux primitive) propagates conditions at the boundaries (faces) of each mesh element. These two computations dominate execution time for most simulation tasks and as such we focus on these in our work.

Machine learning (ML) continues to become ever more pervasive. At the heart of many ML [44] workloads is General Matrix-Matrix multiplication (GEMM). Unlike prior commercial PIM evaluations which focus on skinny GEMMs (the non-shared dimension is small) which are dense, in this work, we focus on GEMMs where one of the matrix inputs is also sparse (has many zeros, as in Figure 2b). We term these sparse skinny gemms (ss-gemm) and they manifest in many ML inference scenarios (e.g., Deep Learning Recommendation Model (DLRM) [43] with small batch sizes).

Graph analytics attempts to derive insights by analyzing the connectivity of graphs and data associated with its edges and vertices. Graph analytics is regularly used for navigation, chemical and biological modeling, social network monitoring and analysis, and many more applications.

Many common graph analytics workloads operate by iteratively propagating vertex properties (pull or push) across graph edges to neighboring vertices. In pull-based algorithms, a local vertex is processed by reading properties from each of its neighbor vertices and updating the local vertex based on what was read. In push-based algorithms (Figure 2c), a local vertex is processed by reading its properties and updating neighbor vertices based on what was read (using atomic RMWs to avoid race conditions). Push implementations have been found to offer attractive performance properties for many graph algorithms and inputs [12, 45], and they are widely used in GPU graph frameworks [48, 23] and benchmark suites [41, 20]. As such, we focus on a push-based algorithm primitive (termed push-primitive).

PIM’s primary performance benefit comes from increasing effective memory bandwidth.¹¹1Some forms of PIM may offer improved memory latency as well, but this work focuses on the bandwidth benefits. Therefore, it will not improve performance for workloads that do not stress available memory bandwidth. Memory bandwidth sensitivity will depend on both the workload and the target architecture. This property can be tested analytically by calculating the algorithmic op/byte ratio (Figure 3a) and determining whether it falls in the compute-limited or memory-limited range on a roofline model for the target architecture.

Even for workloads that are limited by memory bandwidth, residency and on-chip reuse of computation’s inputs/outputs can preclude PIM amenability. Should a computation’s inputs/outputs be resident in on-chip structures or manifest enough reuse when moved from memory to on-chip structures, the computation is less likely to attain acceleration with PIM. For the former, employing PIM would necessitate flushing data to memory, resulting in added data movement overhead. For the latter, with enough reuse, moving data closer to the processor to take advantage of the fast and low energy caches and compute available at the processor is a better solution than using slower compute available in memory.

We can test for both memory residency and low on-chip reuse by determining the proportion of computations that require accessing physical memory versus those that only access on-chip structures (Figure 3b). Comparing this heuristic to a PIM memory bandwidth multiplier (Section 2.2) can be used to test if an application manifests this PIM-amenable characteristic. Application dataflow can indicate if a computation’s inputs/outputs are likely to be resident in on-chip structures. Computations benefiting from cache-aware optimizations (e.g., tiling, kernel fusion) will also manifest good on-chip reuse.

As discussed in Section 2.2, compute units in commercial PIM designs are associated with specific memory bank(s). As such, interacting operands in a computation should map to the same bank to effectively harness PIM acceleration (Figure 3c). We term this property operand locality in our discussion. In the absence of operand locality, costly GPU-orchestrated data transfers between banks will be necessary, which will eat into PIM acceleration.

We propose an operand interaction-centric testing of operand locality. Single operand scenarios (e.g., in-place updates) trivially manifest operand locality. Commutative interactions between multiple elements in a single data structure (e.g., reductions) are also trivial, as interactions between elements in the same bank can simply be performed first. For multi-operand, multi-data structure cases, we observe that localized operand interactions, wherein small subsets of operands within multiple data structures interact with each other, are generally PIM-amenable. A good example of this behavior is element-wise computations (e.g., in vector sum, a single element in each data structure interacts with a single element in another data structure) which can achieve operand locality via careful co-alignment at allocation  [7]. In some cases, localized operand interaction can be induced via data map** (e.g., packing a matrix in matrix-vector multiplication). In such cases, the costs of doing so have to be factored in when assessing PIM impact.

The bandwidth boost attained in PIM is possible via execution of the same operation in parallel across multiple banks. As discussed in Section 2.2, as memory operations have associated row and column addresses, this bank-parallel execution can be employed when operands in different banks in a computation are at the same row/column locations (e.g., see accessing operand a across banks vs accessing operand c in Figure 3d). Note that, within a single DRAM word (256bit single DRAM column), interacting operands therein (e.g., 32bit operands) also have to align (depicted as SIMD alignment in Figure 3d). In absence of this, costly shift operations will be necessary²²2Shifter implementations can be costly in DRAM technology due to the limited number of metal layers.. We term these properties together as aligned data parallelism. As processors often spread a contiguous physical address chunk across multiple channels/banks, ensuring interacting PIM operands are interleaved similarly at allocation time can help attain this characteristic.

While assessing the offloading of a primitive to PIM, as discussed in Section 3, a programmer can first assesses the primitive using our proposed PIM-amenability-test. Assuming PIM amenability exists, we discuss here a methodology which can serve as a general template to offload a primitive to PIM. We depict it in Figure 4a.

As discussed in Section 3.1, ensuring operand locality and aligned data parallelism is critical to attaining PIM acceleration. To that end, identifying operands or data structures ( ) and, more crucially, identifying interactions between them ( ) are the first steps in offloading to PIM. Subsequently, operands are placed in DRAM banks ( ) such that costly inter-bank communication, cross SIMD operations, and (where possible) inter-row interactions are avoided. Finally, ( ), a stream of pim-instructions is deduced to orchestrate the computation over PIM units.

In many cases, compute orchestration follows naturally once data-placement is deduced. However, we also observe that allocation/management of pim-registers to minimize row activation overheads (i.e., staging data from open rows into pim-registers) is often necessary. Note that GPU implementations similarly optimize for effective use of registers.

Data placement: As vector-sum manifests elementwise interaction between two input and one output array, allocating structures such that elements at a given offset are mapped to the same bank [6] ensures PIM-amenable data placement. Note that this is a common operand interaction scenario for PIM-amenable computations.

Command orchestration: Command orchestration for vector-sum is done via broadcast pim-commands, which sequentially read input values, perform adds, and write output values. Further, effective use of pim-registers to stage data from DRAM to minimize row activation overheads is needed.

Data placement: Wave simulation largely operates over arrays and employs three types of operand interaction: elementwise, reduction, and neighboring mesh element. Of these, while elementwise interaction can be tackled via data placement as was done for vector-sum, reduction and neighbor interaction warrant more care. For the former, blocked data placement is employed (discussed below in the context of ss-gemm primitive which also employs this). For the latter, array (grid) elements are placed such that, to the extent possible, neighboring faces reside in the same bank as depicted in Figure 4b.

Command orchestration: Despite their PIM amenable properties, wavesim primitives exhibit complex interaction patterns between operands which complicate orchestration. Considerable care is necessary to effectively utilize available registers while avoiding memory spills and lowering row activation overheads. While we hand schedule the computation in our analysis, existing compiler methods for register allocation [18, 16] can be adapted for PIM-specific cost awareness.

Data placement: To harness broadcast pim-commands and concomitant performance, we place the input dense matrix in a blocked format as depicted in Figure 5, which is tailored to both the dimensions of the matrix and the address interleaving of the system. This layout has multiple attractive properties: skinny matrix values can be broadcast as an immediate PIM operand on the data bus, and accumulation of partial products avoids inter-bank, intra-SIMD, and (to the extent possible) inter-row operations.

Command orchestration: Compute orchestration follows from data placement: an element of the skinny matrix is broadcast to banks and partial results are accumulated in pim-registers before being written to memory.

Data placement: Variation in graph connectivity precludes the use of broadcast commands and co-location of interacting neighbors (source and destination nodes in Figure 4d). Instead, single-bank pim-commands execute in-place destination updates, avoiding operand locality layout constraints.

Command orchestration: Compute orchestration for push-primitive also follows from its data placement. The source node’s value is read from memory, then updates to neighboring nodes are calculated and applied using single-bank pim-commands (namely, a pim-ADD command loads the current value and adds an operand supplied on the data bus, placing the result in a pim-register, and a pim-store command stores this result to memory).

In our work, we use analytical models to evaluate performance. Our choice is guided by the following reasons. First, commercial PIM designs are still only available as functional prototypes. Second, we aim to study primitives considering realistic problem sizes, where PIM is likely to be beneficial. This renders GPU simulators difficult to use due to long simulation times. Finally, we assume a strong GPU baseline, even considering future software optimizations (see below).

GPU Performance Model: For our GPU baseline, we assume the execution time is primarily a function of memory bandwidth (assumed to be 90% of peak) and data accessed. Further, we assume perfect reuse with two exceptions: inter-timestep reuse is not modeled for wavesim (we assume polynomial degree $p=2$, $729$ data points per element, and $65K$ elements per GPU, which is too large to fit in cache), and cache locality for push-primitive is based on actual cache hit rates measured using rocprof [3] with push-based workloads from graphBIG [41] (specifically, hit rates of 44%, 20%, and 57% are observed for roadnet-usa [20], a synthetic power-law graph with 1M nodes and 10M edges, and a synthetic power-law graph with 10M nodes and 100M edges, respectively).

We believe this to be a fair assumption for memory-limited workloads which manifest low op/byte ratios as discussed in Section 3.2. Further, we assume HBM3 memory [2] in our analysis (Table 2). We do so to be both forward-looking and avail our baseline GPU with the best possible memory bandwidth. As such, this provisions a compelling baseline against which to compare PIM acceleration benefits.

Further, for ss-gemm, we assume an optimized GPU baseline which can exploit row-sparsity to both avoid loading the zero rows and computing on them. We estimate this sparsity by analyzing row-sparsity occurrence for computations in MLPerf DLRM-based recommendation model [43] using the Terabyte Click Logs testing dataset [19].

PIM Performance Model: For PIM, we first deduce detailed pim-commands (Section 4.2) and subsequently take into account DRAM timings (Table 2) and operations (row activation, etc.) to determine PIM execution time. Multi-bank pim-commands are issued in-order at half the rate³³3As dictated by the $t_{CCDL}$ timing parameter for back-to-back requests to the same bank group, as opposed to the minimum possible time between reads/writes: $t_{CCDS}$. of regular reads/writes as is the case with the HBM-PIM design [34]. Single-bank pim-commands can be freely reordered and can be issued at the same rate as regular reads/writes. Further, push-primitive updates are also assumed to occur atomically, which can be guaranteed by existing per-address ordering assumptions in the memory controller.

We depict PIM speedups for studied primitives in Figure 6. For our PIM strawman design, the upper bound for performance is about 4x assuming the baseline GPU can utilize 100% of peak memory bandwidth (optimistic). As depicted, while vector-sum attains over 2.6x speedup, primitives under study in this work fare far worse, delivering between 0.23x-1.66x speedups vis-a-vis GPU.

For the ss-gemm primitive, except very skinny matrices ($N=2$), PIM incurs increasing slowdown (between 25-77%) as $N$ increases compared to the GPU. This is expected because, as data reuse improves, moving data to the GPU and exploiting the reuse on chip is more beneficial. Similar trends are observed for push-primitive as cache hit rate improves. Overall, we observe that, even with considerable care to place data appropriately and orchestrate computation efficiently, commercial PIM designs do not attain broad acceleration.

We discuss in this section further analysis of our PIM speedups, focusing on unique challenges when computations are offloaded to PIM and also a unique opportunity. We also discuss how the identified challenges are more a function of commercial PIM designs and not specific to primitives under study. That is, addressing these challenges will likely enable broad acceleration with commercial PIM.

Challenge - Row-activations: A key impediment to PIM acceleration for all primitives (including vector-sum) is row activation overheads - that is, latency costs to open DRAM rows. In fact, this accounts for 27% and 50% of total latency for wavesim-volume and wavesim-flux primitives. In the baseline GPU, bank-level parallelism, wherein row activation latency in one bank is overlapped with data access from another bank, can be employed. However, this is not possible for multi-bank pim-commands. Further, while pim-registers can be used to stage data from rows to lower row activations, this is difficult to realize for primitives with complex operand interactions and high intermediate results which also consume registers (e.g., wavesim primitives).

Challenge - Cache reuse: Primitives whose access patterns exhibit high cache locality are a poor fit for PIM, since on-chip reuse exploited by the GPU can outweigh the bandwidth benefit of PIM. In some cases, statically determining reuse potential and/or which accesses will exhibit cache reuse is difficult because locality and reuse potential is dependent on the input data (e.g., for ss-gemm and push-primitive). With a binary PIM-offload policy (all or nothing), while the GPU can harness the benefits of locality, baseline PIM does not, leading to poor performance.

Challenge - Registers/command bandwidth: We also observe in this work that PIM architecture decisions, specifically the number of registers and command bandwidth available to single-bank pim-commands, can limit PIM acceleration. While we discussed register space implications above, command bandwidth is a unique bottleneck for push-primitive which relies on single-bank PIM commands. While the regular GPU access rate is limited by data bandwidth, and multi-bank PIM command rate is limited by ALU resources, single-bank PIM command rate, in the case of push-primitive, is exclusively limited by command bandwidth especially since this primitive relies on commands which do not send operands on the data bus (pim-store, Section 4.2.5).

Opportunity - Sparsity: For our ss-gemm GPU implementation, we assume a GPU baseline that can exploit sparsity. If, similarly, PIM can exploit sparsity, perhaps even finer-grain sparsity than is possible at the GPU, PIM acceleration can be improved further. To this end, we observe that commercial PIM designs orchestrate computation using pim-commands issued subject to fixed timing constraints (similar to how traditional memory operations are issued). These commands are issued at fine granularity (a pim-command accesses at most one 32B DRAM word per bank). This, as we show below, can be used to exploit fine-grain data sparsity with PIM.

We first tackle row activation overhead evident in PIM executions. To this end, we design a PIM architecture-aware optimization, which we term architecture-aware row activation. In the baseline PIM design, row activation for all banks is incurred on the critical path and is followed by compute commands, each of which are issued to a subset (even or odd) of the available banks (see Baseline schedule in Figure 7a), as the PIM unit/ALU is shared by two DRAM banks. Note that kee** PIM compute commands in program order is necessary to honor register dependencies.

Instead of this schedule, we propose to first split all-bank row activations into separate even and odd bank activations. This allows for a decoupled parallel schedule depicted as PIM arch-aware schedule in Figure 7a. This schedule hides activation latency behind useful work by eagerly activating the next row in one subset of banks while compute commands are being performed in the opposite subset. It also does not impact the order of compute commands or serialized activation latencies and dependencies within odd and even banks, which ensures functionally correct execution. Memory controllers can be augmented to generate/issue such architecture-aware row activations. Further, a compiler pass can also aid in the generation of architecture-aware row activations.

A unique opportunity that commercial PIM designs avail (Section 4.3.3) is thanks to their reliance on pim-commands issued at fine-granularity subject to fixed DRAM timing constraints (similar to how traditional memory operations are issued). This design decision can be harnessed at the software level to perform additional checks at the processor to opportunistically skip issuing commands when certain conditions (e.g., sparsity) are met. Figure 7b illustrates how this works for ss-gemm; before issuing a pim-command to multiply a value of the sparse skinny matrix with a column in the dense matrix, the sparse matrix value is inspected; if this value is zero, then the command is skipped, improving performance.

Our above sparsity-aware PIM optimization has several notable benefits. First, this functionality enables PIM to harness dynamic sparsity (execution time sparsity) which is typically hard to exploit in GPUs [54]. Further, our proposed functionality allows harnessing of sparsity without reliance on specialized sparsity formats which incur data-transformation and metadata overheads. Finally, sparsity-aware PIM allows exploiting of finer-grain sparsity than possible for GPU. The PIM design can easily exploit element-level sparsity whereas GPU harnesses row-level sparsity: skip loading rows with all zeros for the skinny matrix (note, dynamic element-level sparsity at the GPU will require creating specialized sparsity-aware data formats at runtime).

A binary PIM offload decision hurts PIM performance (push-primitive) by obviating input-dependent caching benefits for PIM executions (Section 4.3.3). Instead, we believe that selective offloads to PIM that consider cache reuse is a more superior strategy. To this end, we make a case for cache-aware PIM optimization, where existing fine-grain pim-instructions are offloaded to PIM under the guidance of a locality predictor. Note that such techniques have been studied by prior works (e.g., dynamic schemes [9], offline schemes [5]) and can be augmented to work with commercial PIM designs.

While so far we have discussed targeted optimizations which can broaden commercial PIM acceleration reach, we also believe that PIM architecture decisions, specifically the command bandwidth available to single-bank pim-commands and the number of PIM registers, are important determinants of PIM acceleration. To that end, we also study in this work the effects of varying these on PIM performance. We believe that careful attention to these decisions will be necessary to avail broad PIM acceleration.

Wavesim acceleration improvements with architecture-aware row activation and increased register resources⁴⁴4Baseline PIM in line with commercial PIM designs assumes sixteen registers. are depicted in Figure 8. For the wavesim-volume primitive, architecture-aware optimization improves PIM speedup from 1.5x to 2.04x. Further, this optimization entirely eliminates row activation overheads for this primitive such that more registers do not improve performance. This is in contrast to the wavesim-flux primitive which exhibits higher register pressure and row activation overheads. At lower register counts (16), architecture-aware activation does not improve performance because there are not enough commands per row activation to hide parallel activation latency or to amortize serial activation latency. However, more resources reduce register pressure, enabling this optimization to better hide activation latency and achieve up to $2.63\times$ speedup over the GPU baseline.

For ss-gemm, we focus on implications of our sparsity-aware PIM optimization (depicted in Figure 9). We observe here that sparsity-aware PIM significantly improves the PIM speedup (more than $3\times$) with expected tapering in benefits with increased reuse at the GPU (increasing $N$). Further, it also allows PIM to manifest acceleration in scenarios where baseline PIM manifested a slowdown (speedup of $1.07\times$ for $N=8$, while baseline PIM suffers from 57% slowdown).

As PIM acceleration for push-primitive is most hindered from cache unawareness and lower command bandwidth for single-bank pim-commands, we preferentially study the effects of both of these for this primitive and depict the resulting acceleration in Figure 10. For this study, we model the effects of a locality-based predictor, using a cache model (16-way, 4MB, LRU replacement) which classifies updates to graph nodes in push-primitive as either likely manifesting reuse (performed in cache) or not (performed in PIM). As before, we evaluate these workloads on three graph inputs with varying degrees of locality, each of which is labeled with its observed L2 cache hit rate. Overall, cache-aware PIM prevents performance degradation related to the cache reuse observed in the baseline PIM, leading to an average speedup of $1.20\times$ (max $1.39\times$).

Further, we also model an optimized GPU baseline wherein the GPU can also leverage the locality predictor to reduce access granularity (i.e. use 32B rather than 64B accesses) for updates that do not benefit from caching. We term this cache-aware GPU, and it achieves up to $1.68\times$ speedup relative to the baseline GPU. Although cache-aware PIM reduces data transferred across the memory interface relative to cache-aware GPU, it does not reduce the command bandwidth demand, and the strictly modeled DRAM latency requirements actually lead to worse performance for cache-aware PIM.

Additional command bandwidth only benefits single-bank PIM commands that do not carry data (i.e., push-primitive pim-store commands). With $4\times$ as much command bandwidth,⁵⁵5A 4x multiplier roughly corresponds to how much additional command bandwidth is possible if underutilized data bandwidth could carry command information. PIM performance improves further to exceed cache-aware GPU performance for all inputs and provide up to $2.02\times$ speedup relative to the baseline GPU.

Overall, while the introduction of commercial PIM designs reveals new bottlenecks for many primitives, there clearly is opportunity to address these bottlenecks to support a broader range of primitives.