Reconfigurable Low-latency Memory System for Sparse Matricized Tensor Times Khatri-Rao Product on FPGA

This paper follows the definitions and symbols used in Tensor Algebra Compiler (TACO)[18]. The only difference is, we use capital calligraphic letters (e.g., $\mathscr{B}$) to represent tensors with N dimensions.

Table I summarizes the symbols commonly used in the paper.

MTTKRP is the core computation kernel in the alternating least square (ALS) method, which is a popular method for finding the factor matrices in CPD. The CPD, with R as the decomposition rank, approximates tensor ($\mathscr{B}$) as the sum of R rank-one tensors: $\mathscr{B}\approx\sum_{r=0}^{R-1}a_{r}\circ d_{r}\circ c_{r}$, where $\circ$ denotes the outer product, and a, c, d are the components of the factor matrices A, C, and D. ALS method is a tensor approximation method computes using the above equation. In ALS, one of the factor matrices is updated at a time while fixing the rest. Algorithm 1 shows the steps of CP-ALS algorithm. The MTTKRP operation, which is the focus of this paper, takes place in lines 2, 3, and 4 of Algorithm 1 to compute each of the factor matrices of all modes.

MTTKRP consists of N-dimensional tensor multiply with N-1 factor matrices. Here, N is the order of the tensor. Mathematically, mode-1 MTTKRP (for three-dimensional tensors) can be expressed as Equation 1. In this equation, A, C, and D are dense matrices while $\mathscr{B}$ is a three-dimensional tensor (matricizied along with the first mode), and $\odot$ denotes the Khatri-Rao product.

This operation can also express in index notation as Equation 2.

In the following sections, we target 3D tensors since they are easy to understand. Even though we focus on 3-dimensional tensors as examples in this paper, we can expand our work into higher dimensions without any extra work.

The sparsity of the tensors creates irregular data access patterns, which makes the external memory access more challenging than the dense MTTKRP. All the real-world tensors are sparse in nature. Algorithm 2 shows the sequential sparse MTTKRP (spMTTKRP) approach for third-order tensors in COO format[19][20].

We use the term fibers throughout this paper. Fibers are the building blocks of tensors. Fiber is a one-dimensional fragment of a tensor obtained by fixing all indices but one. Tensor fibers are a higher-order extension of matrix rows and columns. The fibers of a three-dimensional tensor can be represented as $\mathscr{B}_{:,j,k}$, $\mathscr{B}_{i,:,k}$, and $\mathscr{B}_{i,j,:}$. Similarly, for a matrix $C$, its rows $C_{i,:}$ and columns $C_{:,j}$ are its fibers. In Equation 3 and 4, the term fiber corresponds to the rows of the factor matrices.

The DMA engine is in charge of communicating the fibers of the matrices between PEs and the external memory. We store the matrices in row-major order because the MTTKRP algorithm encourages row-wise matrix accesses. The matrices in row-major order make the fiber accesses into a bulk of sequential memory transfers. Figure 2 shows the DMA engine and its internals. It has several DMA buffers inside. Therefore, it can support multiple fiber reads and writes simultaneously. The number of DMA buffers is proportional to the number of PEs connected to the same LMB. A large number of DMA buffers in a LMB can reduce the maximum operating clock frequency due to increased hardware routing complexities.

The cache focuses on satisfying a single memory request with minimum latency. Within the cache engine, we explore the spatial and temporal locality. It loads the scalars of the tensor from the external memory.

Our non-blocking cache uses a 3-stage pipeline to achieve high frequency. We keep the cache-line width similar to the data width of DRAM Interface IP to avoid implementation complexities. Instead of forwarding a single element from the cache to PEs, the cache passes the complete cache-line to the Request Reductor (RR). Then, the RR advances the requested portion to each PE while storing the incoming cache-line inside its temporary buffer.

RR converts element-wise cache reads to cache-line accesses. As shown in Figure 3, the RR is a 2-stage pipeline. In the first step, a temporary buffer stores the most recent memory reads. It is a CAM-based memory implementation that keeps the most recent external memory reads inside. Since CAMs are hardware expensive, we keep the number of elements in the buffer small.

If requested data is not in the temporary buffer, the read request advances to the Recent Request Status Holder (RRSH). RRSH keeps the status of recently forwarded requests to the cache. If the incoming read request belongs to one of the pending cache-line requests, the PE id and address are kept in the RRSH. When a cache-reply from cache reaches the RRSH, the pending requests corresponding to that cache line are satisfied by sending the corresponding data elements to the requested PEs. RRSH helps to convert element-wise cache access to cache-line-wise accesses. It drastically reduces the memory traffic to the cache. RRSH logic can be implemented using a hash table. In our work, we use XOR-base hash table[26] considering its high throughput and scalability.

The request router has 2 responsibilities. They are (a) receive memory requests from different LMB units and forward them to the DRAM interface IP, (b) forward the data coming from external memory to the LMB units.

Users can configure our design during the synthesis step. The configuration heavily depends on the characteristics of the compute fabric and data layout. In the current literature, all the FPGA or CGRA based implementations use a variation of COO Format (e.g., CISS) to store the high dimensional tensors. Meanwhile, the dense matrices use a row-major format. We assume all the compute fabrics that use our memory system use the same kind of memory layout. Under these conditions: (a) the sparse tensors use the cache due to the spatial and temporal locality of the data access pattern, (b) The fibers of the matrices are accessed through DMAs due to their high spatial locality and poor temporal locality.

In Section V-C, we dive deeper into configuring the number of LMBs depending on the type of the compute fabric. Our experiments show that the performance improvement due to the total number of DMAs in an LMB saturates after 4 DMAs. Increasing the number of DMAs also negatively impacts the maximum operating frequency due to increased place and route complexity. We further observed that the cache size also influences the maximum operating frequency of the overall design. We keep cache size including, the total number of cache lines and cache line width, as configurable parameters. We maintain the cache line width equal to the memory interface IP’s data width since it avoids design complexities of the cache.

We implement different configurations of the memory controller on Xilinx Alveo U250 FPGA[28] using Verilog HDL. This device has 1,728 K LUTs, 3,456 K Flip-flops, and 327 MB of URAM memory. Simulations, synthesis, and place-and-route implementations are performed using Xilinx Vivado Design Suite 2020.2. In our work, we focus on accessing a single external DRAM component efficiently. We use the Xilinx memory interface IP[29] to connect our design with the external memory. For the Alveo U250 board, Xilinx provides Memory Interface IP with 31-bit address width and 512-bit data width (with ECC turned on).

We compare our proposed memory system with 3 approaches.

There are 2 types of spMTTKRP kernel accelerators based on the communication between external memory and compute fabric.
Type-1: The systolic array-based spMTTKRP compute fabrics[24][25] in the current literature have a single point of access to the external memory for each type of data structure. After loading the data to the compute fabric, they are shared among PEs using the PEs near the memory system. For instance,[25] use shared Loading and Storing Units for each type of data structure (i.e., Matrix Loading Unit (MLU), Tensor Loading Unit (TLU), and Matrix Store Unit (MSU)) and share the data among PEs using the edge PEs.
Type-2: The compute fabrics with multiple points of access to external memory fall into this category. For instance, Algorithm 3 can be mapped into a compute fabric with multiple PEs with independent memory accesses. Here, the computations executed by each PE are independent of others. The memory system should be sophisticated enough to handle this type of compute fabric.

In the experiments, we use 2 types of configurations depending on the above communication types between external memory and compute fabric.
Configuration-A: For Type-1, having multiple LMBs does not help because their data structures only have a single point of access to the external memory. Therefore, with Type-1 compute fabrics, we use a single Large LMB. Table II shows the resource utilization of each module in the memory system that we used in our experiments.
Configuration-B: For Type-2, multiple LMBs can increase the performance. Their compute fabric has several points of access for each of the data structures. In our experiments, we use 4 LMBs, with each LMB connected to a PE. Table II summarizes the resource utilization of each module in the design we used in our experiments.

As we discussed previously, we ran the experiments on different datasets with different memory system configurations. Figure 4 shows the improvement after using our proposed memory systems over other alternatives. We compute the memory access speedup by comparing the total memory access time with the commercial memory controller IP-only setting. Our proposed system achieved around 3.5x speedup compared to the commercial memory controller IP-only setting. Also, compared with cache-only and DMA-only memory systems, our approach achieved around 2x and 1.26x speedup, respectively.

Our approach always achieves higher speedup compared to the cache-only setting. In MTTKRP, the consecutive secondary cache misses to the same cache line occur while accessing the fibers. In modern caches, secondary misses are avoided using Miss Status Holding Registers (MSHR). But conventional MSHR can not handle a large number of secondary cache misses without losing the performance. Further, the memory traffic between the cache and compute fabric can also reduce the performance in the cache-only setting. In our method, we reduce this traffic while avoiding secondary misses using the Request Reductor.

DMAs do not exploit the temporal locality of data. Further, there can be garbage data in DMA transactions when the length of the data requests is shorter than the width of the memory interface IP. Our proposed system outperforms the DMA-only setting simply because it avoids the above issues with DMAs.