A Framework for High-throughput Sequence Alignment using Real Processing-in-Memory Systems

Fig. 1 compares the organization of conventional CPU processing systems to PIM systems. In conventional systems, illustrated in Fig. 1(a), data resides in dynamic random access memory (DRAM) which is typically organized into multiple DRAM banks. This data is transferred to the CPU cores where computations are performed on the data. If the same data is reused for many computations, the cost of moving that data from memory to the CPU cores is amortized. However, if the data is not reused, the cost of moving the data is much higher than the cost of the computations on that data. In this case, the CPU cores will be idle most of the time waiting for memory accesses to complete, and the movement of data between the CPU cores and the memory becomes a major performance bottleneck.

In a PIM system, illustrated in Fig. 1(b), small PIM cores are placed in the memory chip near the memory banks. These PIM cores are typically much less powerful than CPU cores at performing computations, however they are much faster at accessing data from memory because of their proximity. Hence, if only a few computations are performed on the data, it is cheaper to perform these computations in the PIM cores than to move the data all the way to the CPU cores. In this case, programmers would load their code onto the PIM cores and execute it there where the data can be fetched quickly, instead of executing it on the CPU cores. Sequence alignment algorithms are well-suited for such a system because there are usually a few computations performed for each data element accessed from the intermediate data structures (e.g., entries of the DP table in NW).

The UPMEM PIM architecture (devaux2019true, ) is the first publicly-available general-purpose programmable (https://sdk.upmem.com) processing-in-DRAM architecture. An UPMEM system consists of a set of UPMEM DIMM modules plugged alongside main memory (standard DDR4 DIMMs) and acting as parallel co-processors to the host CPU. An UPMEM module is a DDR4-2400 DIMM with 16 PIM-enabled chips, where each chip consists of eight general-purpose processing PIM cores called DRAM Processing Units (DPUs). Each DPU is coupled with a 64MB DRAM bank called main RAM (MRAM). A current UPMEM system supports up to 20 UPMEM modules, which is equivalent to 2560 DPUs and 160GB of memory. A DPU is a 32-bit RISC processor with a proprietary Instruction Set Architecture (ISA), which can potentially run at 500 Mhz. Each DPU has 24 hardware threads that share a 24KB instruction memory (IRAM) and a 64KB scratchpad memory which is also called a working RAM (WRAM). The threads also share the 64MB MRAM bank coupled with the DPU. DPUs cannot communicate with other DPUs or access data outside their own MRAM bank. The host CPU transfers data between the main memory and the MRAM banks, and coordinates communication and synchronization across DPUs if needed.

Fig. 2 illustrates our framework’s overall workflow. In Step (1), we read the sequence pairs from an input file on disk and store them to the main memory. In Step (2), we transfer the sequence pairs from the main memory to the UPMEM DIMMs, distributing them evenly across the MRAM banks of the different DPUs. We use parallel transfers so that the sequence pairs are written to multiple MRAM banks simultaneously, thereby optimizing the transfer latency.

Next, we launch the DPU kernels and have each thread in each DPU work independently to align a set of sequence pairs. This parallelization scheme avoids inter-DPU and inter-thread synchronization, which can be expensive in the UPMEM system (DBLP:journals/corr/abs-2105-03814, ). In Step (3), each DPU thread performs a DMA transfer to fetch a sequence pair from MRAM and store it in WRAM. In Step (4), the DPU thread aligns the sequence pair and extracts the alignment operations using traceback. Our framework supports five different alignment algorithms: NW, SWG, GenASM, WFA, and WFA-adaptive. In Step (5), the DPU thread performs a DMA transfer to write the alignment score and operations to MRAM. The thread then moves on to process the next sequence pair, repeating Steps (3), (4), and (5) until all sequence pairs have been processed.

Finally, once all DPUs finish execution, we retrieve and save the alignment results. In Step (6), we transfer the alignment score and operations of each sequence pair from the UPMEM DIMMs to main memory using parallel transfers. In Step (7), we load the results from main memory and write them to an output file on disk.

One important aspect of implementing alignment algorithms on the UPMEM PIM architecture efficiently is data management. Recall from Section 2.1 that an UPMEM DPU has access to two memory spaces for data: a 64KB WRAM and a 64MB MRAM. WRAM is fast and is accessed via loads and stores, whereas MRAM is slow and is accessed via DMA transfers to and from WRAM. Our framework always places the full sequence pair to be aligned and the full alignment result in WRAM for fast access because these data items are small. However, the intermediate data structures used by the alignment algorithms are relatively large. For this reason, it may not be possible to fit the entire intermediate data structure for each DPU thread in WRAM while supporting a large enough number of DPU threads to efficiently utilize the DPU pipeline. In this case, the constrained WRAM capacity can act as a limiting factor to parallelism.

To tackle this trade-off, we provide two alternative implementations of each alignment algorithm. This first implementation, illustrated in Fig. 3 (top), places the entire intermediate data structure for each DPU thread in WRAM, thereby prioritizing fast load/store access to the intermediate data structures. However, as the aligned sequences get larger, the WRAM capacity begins to constrain the number of DPU threads that can be launched, which causes the DPU pipeline to be underutilized. The second implementation, illustrated in Fig. 3 (bottom) places the intermediate data structure for each DPU thread in MRAM and transfers the parts of the data structure that need to be accessed to WRAM on-demand. With the WRAM capacity less of an issue, this approach enables using a larger number of DPU threads to better utilize the DPU pipeline. However, it incurs longer access latency to the intermediate data structures. The user of our framework can easily select which implementation they would like to use, and our framework automatically adapts the number of DPU threads launched depending on the implementation selected as well as the choice of alignment algorithm, read length, and error rate. We evaluate the trade-off between these two implementations in Section 3.5.

Another important aspect of data management is performing dynamic memory allocation efficiently. NW, SWG, and GenASM have fixed-size data structures that are allocated in WRAM or MRAM up-front. However, WFA and WFA-adaptive rely on dynamic memory allocation which is performed frequently as the algorithms run. Dynamic memory allocation is needed for allocating the wavefront components, which vary in size at run time depending on the read length and similarity. The UPMEM SDK has an incremental dynamic memory allocator which allocates memory incrementally from beginning to end and then frees it all at once. However, it is not suitable for our purpose because it is shared by all threads so it requires synchronization across threads to perform allocation, which degrades performance. To overcome this issue, we provide our own per-thread custom memory allocator to perform low-overhead dynamic memory allocation in the WFA and WFA-adaptive algorithms. The allocator also ensures that allocations are properly aligned such that they can be used in DMA transfers between WRAM and MRAM.

We compare the performance of our proposed framework to a mutli-threaded CPU baseline that uses OpenMP to align multiple sequence pairs in parallel. The baseline CPU implementations of NW, SWG, WFA, and WFA-adaptive are taken from the original WFA repository (marco2020fast, ), and the baseline CPU implementation of GenASM is taken from the GenASM repository (cali2020genasm, ).

We evaluate our PIM implementations on an UPMEM system with 2560 DPUs (20 UPMEM-DIMMs) running at 425MHz. We evaluate the CPU implementations on three different server-grade CPU systems shown in Table 1. We use execution time as the basis for comparison, which includes the wall clock time for performing data transfer, alignment, backtrace, and CIGAR string generation.

We use real and synthetic datasets to evaluate our proposed framework, as shown in Table 2. The real datasets are short read-reference pairs generated using minimap2 (10.1093/bioinformatics/bty191, ) by mapping the datasets (https://www.ebi.ac.uk/ena/browser/view) mentioned in Table 2 to the human reference genome GRCh37 (church2011modernizing, ). The synthetic datasets are long sequence pairs simulated using the synthetic data generator provided in the WFA repository (marco2020fast, ).

Fig. 4(a) shows how the execution times of the CPU implementations scale with the number of CPU threads while aligning five million sequence pairs using different alignment algorithms and read lengths, and an edit distance of 1%. The results are reported for the Xeon E5 CPU system, which is the best performing CPU system as we show in Section 3.3. Our key observation is that the CPU implementations face limited performance improvement as the number of CPU threads grows. To investigate the cause of this limited performance improvement, we profile the application using perf. Fig. 4(b) shows that as the number of CPU threads increases, the instructions executed per cycle by each thread decreases. Hence, the CPU is increasingly underutilized when more threads are added. To further understand what is causing the CPU to be underutilized, Fig. 4(c) shows that as the number of CPU threads increases, the cycles spent by the CPU stalling while waiting for memory requests increases. Hence, the limited performance improvement is caused by the inability of the memory to serve memory requests quickly enough. These results demonstrate the importance of using PIM to overcome the memory bandwidth bottleneck faced by sequence alignment applications.

Fig. 5 shows the execution time of our framework aligning five million sequence pairs using different alignment algorithms, read lengths, and edit distances. We report the execution time both with and without the data transfer time between main memory and the UPMEM DIMMs. Based on these results, we make four key observations.

The first observation is that among the three CPU baselines, the best performing baseline in the majority of cases is the Xeon E5 CPU system. Despite not having the largest number of threads, this baseline has the largest L3 cache. This result demonstrates the memory-boundedness of the sequence alignment problem, where having a larger L3 cache is favored over more having more threads.

The second observation is that our framework running on the UPMEM system outperforms the CPU baselines in the majority of cases, even when data transfer time is included. The speedup achieved over the best CPU baseline is up to 4.06$\times$ in the case of SWG. For the state-of-the-art algorithms, WFA and WFA-adaptive, the speedups achieved over the best CPU baseline are up to 1.83$\times$ and 2.56$\times$, respectively. These results demonstrate the effectiveness of PIM at accelerating memory-bound sequence alignment workloads.

Here, we would like to reiterate that our CPU baselines are powerful server-grade dual-socket systems running at full scale. CPU hardware and software have been developed and optimized for decades by large teams of engineers, whereas the UPMEM PIM system has been developed over a few years by a small team. We expect that the relative advantage of PIM systems will be more pronounced as PIM hardware and software matures. We also note that in our current system, the DPUs are running at 425 MHz, however they are expected to run at 500 MHz in future systems which would further improve performance.

The third observation is that, when the data transfer time is not included, our framework achieves a speedup over the best CPU baseline of up to 28.14$\times$ in the case of WFA-adaptive (25.93$\times$ for WFA). The results without transfer time are important for two reasons. The first reason is that in the current UPMEM system, the UPMEM DIMMs cannot be used as regular main memory. For this reason, the CPU must read the data from disk to main memory then transfer it from main memory to the UPMEM DIMMs. However, in future systems where a PIM module could also be used as main memory, the CPU could potentially read data from disk to the PIM module directly and then perform the alignment in the PIM module without the additional transfer. The second reason is that in the current UPMEM system, one cannot overlap writing to MRAM by the CPU and execution by the DPUs. However, in future systems where CPU access to a PIM module may be overlapped with execution in the PIM cores, such overlapping can hide some of the latency of writing to the PIM module. Therefore, in light of the two aforementioned reasons, the results without the transfer time demonstrate the potential that future PIM systems have for accelerating memory-bound sequence alignment workloads.

The fourth observation is that our framework does not outperform the best CPU baseline for NW and SWG at small read lengths. NW and SWG have more regular memory access patterns than the other alignment algorithms, and when the read lengths are small, the intermediate data structures fit into the CPU cache. The combination of these two factors makes the memory bandwidth bottleneck less pronounced, which diminishes the advantage of PIM. However, for other algorithms where the memory access pattern is less regular, and for larger read lengths where the sizes of the intermediate data structures outgrow the size of the CPU cache, the computation becomes more memory bound causing our framework to achieve large speedups over the CPU baseline.

To further study the scalability of our framework, we evaluate the state-of-the-art algorithm, WFA-adaptive, on large sequence lengths of 5Kbp and 10Kbp. NW and SWG have difficulty scaling to such large read lengths because their quadratic space complexity causes them to exceed the MRAM capacity even for a single alignment. However, WFA and WFA-adaptive are not limited by MRAM capacity due to their linear space complexity. The execution time of WFA-adaptive for large read lengths while aligning one million sequence pairs is shown in Fig. 6. We observe that our framework executing on the UPMEM system continues to outperform all CPU baselines. The only exception is for read length 10,000 with an edit distance threshold of 5%. In this case, the WRAM capacity is only sufficient to support a single DPU thread per DPU which underutilizes the DPU pipelines. Although the algorithm is not limited by the 64MB MRAM for storing all the wavefront components, it is limited by the 64KB WRAM for storing the select wavefront components needed for computing a new wavefront component. Hence, our current framework executing on the current hardware cannot scale to read lengths far beyond 10,000 for WFA and WFA-adaptive due to WRAM capacity constraints. In some bioinformatics applications (such as minimap2), sequence alignment is typically performed only between every two seed chains to avoid long execution time and peak memory that could result from performing sequence alignment on the complete read sequence. Thus, we believe that our framework is useful for a wide range of applications despite the sequence length limitation. However, if sequence lengths far beyond 10,000 are a concern, the current limitation can be mitigated by streaming partial wavefront components from MRAM to WRAM in order to reduce the WRAM footprint and support aligning more sequence pairs of larger length simultaneously, or by using multiple DPU threads to align a single sequence pair. These improvements are the subject of our future work. We also expect the supportable read length and the performance for large read lengths to improve in future systems as the hardware and software mature, as the clock frequency of the DPUs increases, and as future systems may have larger WRAM capacity.

Table 3 compares the throughput of our PIM implementations of WFA-adaptive (the fastest algorithm) to a recent GPU implementation of WFA-adaptive by (aguado2022wfa, ) using the results reported in that work. It is clear that our PIM implementation outperforms the GPU implementation in the majority of cases. This result shows that PIM is a promising technology for accelerating sequence alignment, even when compared to mature and widely used accelerators such as GPUs.

Recall from Section 2.3 that our framework provides two implementations of each algorithm: one that only uses WRAM for intermediate alignment data structures, and another that uses MRAM for these data structures and transfers data currently being accessed from these data structures to WRAM as needed. Fig. 7 compares the execution time and scalability of these two implementations for each algorithm and read length with edit distance 1% and 5%. Based on these results, we make three key observations.

The first observation is that for the algorithms that use a large amount of memory for the intermediate alignment data structures (i.e., NW and SWG), the implementations that use both WRAM and MRAM scale better with the read length than those that only use WRAM. In the case of NW for small read lengths, the implementation that uses WRAM only is up to 1.70$\times$ faster. However, for the remaining cases, this implementation is slower or cannot even execute. The reason is that the implementations that use WRAM only can only support a small number of DPU threads due to the constrained WRAM capacity, which prevents them from utilizing the DPU pipeline well. In contrast, the implementations that use both WRAM and MRAM can support a larger number of DPU threads, causing them to perform better for long read lengths despite incurring higher memory access latency. Note that the trends for these two algorithms are independent of the edit distance since their operations and memory consumption are independent of the alignment score.

The second observation is that for the algorithm that uses a small amount of memory for the intermediate alignment data structures (i.e., GenASM), the implementation that only uses WRAM is usually faster (up to 2.76$\times$ for edit distance 1%). The reason is that the WRAM only implementation can support a large enough number of DPU threads to utilize the DPU pipeline well. We note, however, that for larger edit distances that require more memory, the implementation that uses WRAM and MRAM becomes more favorable for large reads. For example, for read length 1,000 and edit distance 5%, the implementation that uses WRAM and MRAM together is 1.91$\times$ faster.

The third observation is that for the implementations that use a moderate amount of memory for the intermediate alignment data structures (i.e., WFA and WFA-adaptive), the implementations that use WRAM only are faster for shorter reads and low edit distance (up to 1.17$\times$ for WFA and 1.12$\times$ for WFA-adaptive at an edit distance of 1%). On the other hand, the implementations that use both WRAM and MRAM are faster for longer reads (up to 2.70$\times$ for WFA and 1.25$\times$ for WFA-adaptive at an edit distance of 1%). The reason is that the number of DPU threads that can be used decreases as the read length gets larger, which favors the implementation that can use more threads over the one with lower access latency. Note that the difference between the two implementations grows as the edit distance grows because the memory consumption of WFA and WFA-adaptive is highly sensitive to the edit distance. For example, for edit distance 5%, the implementation that uses WRAM and MRAM together is up to 6.15$\times$ faster.

The three observations presented in this section demonstrate the importance of our framework supporting both implementations of each algorithm, those that use WRAM only and those that use WRAM and MRAM for intermediate alignment data structures. In either case, our framework can automatically identify the maximum number of threads that can be used to alleviate this burden from the user.