Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures

We illustrate our PARLOOPER/TPP framework with a running General Matrix Multiplication (GEMM) example (see Figure 1 and Listing 1). The zero_tpp sets the input 2D tensor to zero values. The brgemm_tpp corresponds to the Batch-Reduce GEMM (BRGEMM) TPP which is the main tensor contraction tool in the TPP collection [21, 12]. BRGEMM materializes $C=\beta\cdot C+\sum_{i=0}^{brcount-1}A_{i}\times B_{i}$. This kernel multiplies the specified blocks $A_{i}^{bm\times bk}$ and $B_{i}^{bk\times bn}$ and reduces the partial results to a block $C^{bm\times bn}$. We use the stride-based variant of BRGEMM, where the addresses of $A_{i}$ and $B_{i}$ are: $address\_A_{i}=address\_A_{i-1}+stride\_A$ and $address\_B_{i}=address\_B_{i-1}+stride\_B$ [12]. The Matrix Multiplication input tensors are logically 2D matrices $A^{M\times K}$ and $B^{K\times N}$ that need to be multiplied and added onto $C^{M\times N}$. We follow the approach of previous work [21] and we block the dimensions $M$, $K$, and $N$ by factors $bm$, $bk$, and $bn$ respectively (lines 1-3 of Listing 1). We employ the BRGEMM TPP to perform the tensor contraction with $A$ and $B$ across their dimensions $Kb$ and $bk$. The sub-blocks $A_{i}$ and $B_{i}$ that have to be multiplied and reduced are apart by fixed strides $stride\_A=bk\cdot bm$ and $stride\_B=bn\cdot bk$.

The Matrix Multiplication algorithm is comprised of three logical nested loops which are declared in lines 5-9 of Listing 1/Figure 1-Box A1. The computation involves three logical loops, thus we declare: ThreadedLoop$<3>$ (line 5).

The first loop (line 6) with mnemonic a, corresponds to a loop with start $0$, upper bound $Kb$ and step $k\_step$, and it reflects the “$K$” loop of the GEMM ($K$ is the inner-product dimension). This loop has an optional list of step/blocking parameters $\{l1\_k\_step,\ l0\_k\_step\}$. The second loop (line 7) with the mnemonic b, corresponds to a loop with start $0$, upper bound $Mb$ and step $m\_step$, and reflects the “$M$” loop of the GEMM. Similarly, this loop has an optional list of step/blocking parameters $\{l1\_m\_step,\ l0\_m\_step\}$. Finally, the third loop (line 8) with mnemonic c, corresponds to a loop with start $0$, upper bound $Nb$ and step $n\_step$, and reflects the “$N$” loop of the GEMM. This loop has an optional list of step/blocking parameters $\{l1\_n\_step,\ l0\_n\_step\}$.

The specific instantiation of these loops, i.e. the loop order with which they appear, the number of times each one is blocked and also the way they are parallelized are controlled by the run-time knob loop_spec_string (line 9). Given such a string during run-time, the constructor ThreadedLoop invokes our custom loop generator and emits a C++ function for the target loop instantiation (Figure 1-Box C1). Listing 2 illustrates an exemplary target-loop instantiation function. Subsequently, a C++ compiler is invoked (currently we support icc, clang, and gcc), and the target-loop instantiation function is compiled Just-In-Time (JIT), which can be further used by the code (line 11 in Listing 1). To avoid JIT overheads whenever possible, we cache the JITed target loops: if we request a loop nest with the same loop_spec_string, we merely return the function pointer of the already compiled and cached loop-nest.

In regard to the legality/validity of the loop_spec_string (i.e. what loop permutations are allowed), this is responsibility of the user entity and depends on the computation at hand. The two rules for constructing a valid loop_spec_string are:

• •

  RULE 1: Loop ordering and blockings

Each character (from a to z, depending on the number of the logical loops) can appear in any order and any number of times. In our case, since we have 3 logical loops the characters range from $a$ to $c$. The order with which the loop characters appear in the string determine the nesting loop order, and the times each character appears determines how many times the corresponding logical loop is blocked. For example, a loop_spec_string bcabcb corresponds to a loop where logical loop $b$ is blocked twice (the character $b$ appears 3 times), logical loop $c$ is blocked once (the character $c$ appears 2 times) and the logical loop $a$ is not blocked (it appears only once). The blocking/tiling sizes for each logical loop level are extracted from the corresponding list of step/blocking parameters in order they appear in the list. Our Proof-Of-Concept (POC) implementation allows only perfectly nested blocking/tiling sizes, i.e. in the example above it should hold: $l1\_m\_step\ \%\ l0\_m\_step=0$, $l0\_m\_step\ \%\ m\_step=0$ and $l1\_n\_step\ \%\ n\_step=0$. All these blocking/tiling sizes lists may be provided at runtime (e.g., programmatically determine the blocking sizes given a problem) and do not have to be statically determined.

• •

  RULE 2: Parallelization of loops

If a loop character in the loop_spec_string appears in its upper-case form, it dictates the intention to parallelize this loop at the specific nesting level it appears. In the previous example, the loop_spec_string bcaBcb, corresponds to a loop nest where the 2nd occurrence of loop $b$ is parallelized. Our implementation supports 2 modes of parallelization:

PAR-MODE 1: Relying on OpenMP [23] runtime for parallelizing the loops. In this method, the loop parallelization strategy corresponds to the directive#pragma omp for nowait. If the user intends to parallelize multiple loops, the corresponding capitalized characters should appear consecutively in the loop_spec_string, and it would result in parallelization using collapse semantics. E.g., for the loop_spec_string $bcaBCb$, PARLOOPER generates the loop nest in Listing 2. We allow optionally adding directives at the end of the loop_spec_string by using the special character $@$ as separator. E.g., the loop string “$bcaBCb$ $@$ $schedule(dynamic,1)$” yields the directive: #pragma omp for collapse(2) schedule(dynamic,1) nowait. This method allows PARLOOPER to leverage features from the OpenMP runtime like dynamic thread scheduling. The constructed loop nest is embraced by a #pragma omp parallel region (line 4 in Listing 2). A barrier at the end of a specific loop-level may be requested using the special character “$|$”.

 PAR-MODE 2: Using explicit multi-dimensional thread decompositions. The user can specify explicit 1D, 2D or 3D loop parallelization schemes. For the 1D decomposition, the threads form a logical $R\times 1$ 1D grid and are assigned the corresponding parallelized loop iterations in a block fashion. The user merely has to append after the desired upper-case loop character the substring {R:#threads}, where $\#threads$ is a number dictating in how many ways to parallelize that loop. For a 2D decomposition, the threads are forming a logical $R\times C$ 2D grid that effectively parallelizes the requested two loops. E.g., consider the loop_spec_string bC{R:16}aB{C:4}cb. The loop $c0$ is parallelized 16-ways and the loop $b1$ is parallelized 4-ways using a logical thread grid of 16$\times$4 ($R=16$, $C=4$) – see Listing 3. Each loop that is parallelized is done so in a block fashion using the requested number of “ways”. The 3D decomposition works in an analogous manner.

The current POC supports OpenMP for concurrency purposes, however our back-end loop generator can be extended to support other runtimes (e.g. TBB [24] or pthreads [25]).

After declaring the nested loops, we obtain a ThreadedLoop object (gemm_loop in Line 5 of Listing 1) which can be passed at runtime (up to) three parameters:

• •

  (Optional) A pointer to a function: void init_func()

This function is called before the generated loop-nest and is used for “initialization” purposes (e.g., code that initializes data structures etc.).

• •

  (Optional) A pointer to a function: void term_func()

This function is called after the generated loop-nest and is used for “termination ” purposes (e.g., code that cleans up data structures etc).

• •

  A pointer to a function: void body_func(int *ind)

The function body_func is called at the inner-most level of the generated loop-nest (e.g. line 16 in Listing 2), and it performs the desired computation. The function body_func gets as input an array of integer values which contains the values of the logical indices used in the nested loop in alphabetical order: ind[0] corresponds to the value of the logical index $a$ in the current nested-loop iteration, ind[1] corresponds to the logical index $b$ in the current iteration etc. This index array is allocated and initialized by PARLOOPER (e.g. line 15 in Listing 2). By leveraging the values of the logical indices, the user can express the desired computation as a function of those. We use as body_func an in-place C++ lambda expression, but this is not a requirement in our framework. Considering the GEMM example, and using lambda expression for body_func we express the GEMM computation using the zero_tpp, the brgemm_tpp, and the logical indices with lines 12-17 of Listing 1 and Figure 1-Box A2. In line 13, we extract the logical indices $ik$, $im$ and $in$ from the input array $ind$. We set batch-reduce count in Line 14 to $k\_step$ since the tensor contraction is performed on blocked matrices. In Line 15 we set the current $C_{in,im}$ block to zero if this is the first time we encounter it. Finally, lines 16-17 execute the actual tensor contraction via the BRGEMM TPP.

We highlight the separation of concerns in the heart of PARLOOPER’s design: The logical loops are declared at a high-level in as little as 3 lines of code (lines 6-8 in Listing 1 and Figure 1-Box A1). The main computational core is orthogonal to the loop declaration and is expressed via TPPs in 5 lines of code by making use of the logical loop indices (Figure 1-Box A2). Both the loop-nest instantiation and the TPP code generation are abstracted from the user and are determined at run-time considering runtime parameters (loop_spec_string for the loop generation and the target platform/problem shapes for the TPP code generation – Figure 1-Box C). E.g., on a CPU platform with 3 levels of cache one may block each of the logical loops up to 3 times to maximize the data re-use out of each level of cache [26]. Moreover, skewed tensor sizes may favor loop orders where the “smaller” tensor is accessed repeatedly in the innermost loop. The parallelization of the relevant $M$ and $N$ loops also affects the data locality and the degree of data-sharing among cores. All these loop-instantiation decisions are completely determined via the runtime knob loop_spec_string. Similarly, the JITed code for the BRGEMM TPP (incl. loop unrolling, vectorization, register blocking, instruction selection) depends on the target platform and the TPP shapes [12]. For example, on Intel Sapphire Rapids the BF16 TPP backend can emit Advanced Matrix eXtension (AMX) instructions to accelerate the tensor contractions, whereas on an AMD Zen4 platform the TPP backend can emit AVX512 BF16 FMA instructions. All these low-level details are completely abstracted from the user-level code Figure 1-Box A.

Also, the GEMM in Listing 1 is data-type agnostic. All tensors have templated datatype, and the TPPs are precision-aware per design: depending on how they are setup, they operate with any supported datatype internally. Therefore, the same code works for all precisions without any change.

As mentioned earlier, the legality/validity check of the loop_spec_string is responsibility of the user entity and depends on the computation at hand. An advanced framework, where the legality check is not user’s responsibility, requires two steps/components: i) extracting the semantics of the loop body (written via TPPs) in a way that the data dependencies of the core computation are exposed, and ii) perform a traditional loop-dependence analysis on the loop permutation at hand to check its legality. These steps essentially would bring the current PARLOOPER/TPP framework closer to a “Tensor compiler stack with TPP abstractions”. In the current form of the PARLOOPER/TPP framework, the correctness burden from the user’s point of view is equivalent to writing parallel code in an imperative language (e.g. C with OpenMP): for any nested loop with OpenMP parallelization directives, it is the user’s responsibility to dictate which loops have to be parallelized without introducing race conditions, otherwise the program is incorrect (i.e. the compiler does not check legality of the parallelization).

A proper value of loop_spec_string is essential in order to obtain the best performance on a given combination of problem size and CPU architecture. One way to select loop_spec_string is to hand-craft specific loop strings that are expected to be performant given custom performance modeling of the given platform/problem (Boxes B1 and B3 in Figure 1). Alternatively, one may rely on off-line auto-tuning (Box B2 in Figure 1). Given a logical loop declaration, along with its specifications, the tunable parameters are: (i) How many times to block each logical loops, (ii) What are the blocking sizes, (iii) Which loops to parallelize, and (iv) In which order to arrange the loops. A key observation is that all these decisions can be mapped in 1-on-1 fashion to a specific loop_spec_string along with a list of block sizes.

We created an infrastructure to auto-generate an exhaustive list of strings that observe a set of constraints (see Figure 1-Box B2). Considering the GEMM in Listing 1 and the decisions (i)-(iv) to be made, one can specify the following set of constraints:

1. 1.

   Block loop a up to 2 times, and loops b and c up to 3 times. This captures multi-level caches on modern CPUs.

2. 2.

   For each logical loop $i$, consider its logical trip count $T_{i}$ and find the prime factorization of $Ti=p_{0}\cdot...\cdot p_{n}$. Then pick as block factors the prefix products of the prime factors with the loop’s steps. E.g., for the $M$ loop we can select as $l0\_m\_step=m\_step\cdot p_{0}$ and $l1\_m\_step=m\_step\cdot p_{0}\cdot p_{1}$. This is one example of programmatically picking the blocking factors.

3. 3.

   We may decide to parallelize the $M$ (b) and the $N$ (c) logical loops as those define independent tasks in GEMM. We can parallelize any of the blocked occurrence of the $M$/$N$ loops since they also constitute independent iterations.

4. 4.

   Consider all permutations of a string subject to rules 1-3.

With such a set of constraints, we generate a set of loop_spec_strings configurations to be benchmarked off-line. We implemented this infrastructure using bash scripts, but one could use other tools like Open-tuner [27]. This auto-tuning process involves 0 lines of code change in the GEMM user-code which remains the same as the one in Listing 1/Figure 1-Box A; all different loop variants dictated by the values of loop_spec_strings are auto-generated and JITed by the PARLOOPER back-end.

To assist the selection of the loop-instantiation loop_spec_string knob, instead of relying exclusively on manual settings/exhaustive auto-tuning, we designed a lightweight, high-level performance modeling tool (Figure 1-Box B3). This tool takes as inputs the logical loop order along with a candidate loop_spec_string, the computational core expressed via tensor-contraction BRGEMM TPPs, and few parameters modeling the target CPU, and it simulates/predicts the performance of such a setup. With this tool, one can do off-line, cross-architecture loop-tunings, and obtain a set of good loop-instantiation knobs.

Currently the POC implementation of this tool works only for PARLOOPER-generated loops with BRGEMM TPPs within body_func but may be extended to cover the rest TPPs. The idea is rather simple: If one were to focus on the execution of a specific thread in a specific loop-instantiation with BRGEMM TPP as main computation, all the data that are accessed are logical Tensor slices identified by the tensor-block indices. In Listing 1, the corresponding slices of $A_{im,ik}$, $B_{in,ik}$ and $C_{in,im}$ are characterized by the relevant indices. Consequently, each thread can create a trace of its $A$, $B$ and $C$ accesses that arise in chronological order as the thread proceeds with the execution of the specific loop-istantiation. These traces are compact since they register accesses of full tensor slices instead of individual cache-lines [28]. Therefore, running a data reuse algorithm that simulates multilevel caches on a per-thread basis is a low-overhead operation.

We simulate up to 3 levels of caches, each having its own size and bandwidth. The replacement policy for each cache is Least Recently Used (LRU). When processing a thread’s trace access $t_{i}$ in chronological order, we query if the corresponding tensor slice is in any level of its cache or if it resides in memory. Then, given a BRGEMM execution at iteration $it$, and knowing the memory/cache-locations of the current slices $A_{it}$, $B_{it}$ and $C_{it}$ we can predict the execution cycles taken by BRGEMM by accounting for the relative cache bandwidths and the compute-peak of the platform. Each level of cache is represented as set and is updated based on the LRU policy as the execution progresses. For simplicity we ignore data-sharing, but the traces could be processed in lock-step fashion to account for common sub-tensors in shared levels of cache. Despite these simplifications, our performance modelling is able to single out loop instantiations with poor temporal/spatial locality, and identifies parallel schedules with poor concurrency. Such inefficient loop_spec_strings are assigned a low score and can be considered sub-optimal.

We employ the direct convolution method that uses the BRGEMM TPP for contractions [20, 21]. The input tensor $I$ is convoluted with the weight tensor $W$ to yield the output $O$. The activation tensor conceptually consist of 4 dimensions: the minibatch $N$, the number of feature maps $C$ and the spatial dims $H$ and $W$. We denote the dimensions of $I$ with $N$, $C$, $H$ and $W$ while the corresponding dimensions of $O$ are $N$, $K$, $P$ and $Q$. The weight tensor is logically consisting of 4 dimensions: the feature map dimensions $C$, $K$ and the spatial dimensions $R$ and $S$. We block the dimension $C$ by $bc$ and the dimension $K$ by $bk$, and we get the blocked tensors in lines 1-3 of Listing 4. We declare with PARLOOPER the 7 logical loops that traverse the iteration space in lines 5-13 of Listing 4. Finally, lines 16-25 express the convolution compute kernel in a lambda function by means of the 7 logical indices and the brgemm_tpp. For convolutions with $R=S=1$ we can setup a stride-based BRGEMM whereas for the other cases we can setup offset-based BRGEMM to further increase the kernel’s performance (i.e. the $R$ and $S$ loops would be folded in the BRGEMM by using offsets arrays [12]).

The Sparse $\times$ Dense Matrix Multiplication (SpMM) is a kernel widely used in HPC, including linear solvers and graph analytics [33]. The SpMM $C=A\times B$ involves one sparse matrix $A$ and two dense matrices $B$ and $C$. By exploiting the sparsity of $A$ one can reduce both the bandwidth and the compute requirements of the computation. In addition, by enforcing block sparsity on matrix $A$ one can increase the performance of SpMM by having better data locality in the accesses of the sparse $A$ and by leveraging low-precision/hardware acceleration available on modern CPU and GPU platforms. With the recent explosion in model sizes, and the intrinsic redundancy/sparsity of large models, a promising research vector is exploring block sparsity in DL [34, 35].

We extend the GEMM TPP to support $A$ being block-sparse, while $B$ and $C$ matrices remain dense. We implement block sparse $\times$ dense kernel with BF16 datatype to leverage Fused Multiple-ADD (FMA) hardware-acceleration on x86 and AArch64 CPU platforms. For our implementation, we extend the TPP implementation in LIBXSMM. We enable support for $A$ in Block Compressed Sparse Columns (BCSC) format where the block-size $bm\times bk$ is parameterized. The code-generation of the microkernel works as follows: We iterate over a block row of $A$ and for each non-empty block $bm\times bk$, we multiply it with the corresponding dense block $bk\times bn$ of $B$. The dense multiplication of $bm\times bk$ with $bk\times bn$ sub-blocks of $A$ and $B$ is using 2D register blocking [21] whenever possible (i.e. large $bn$ and $bm$) to hide FMA latency and maximize data-reuse from registers. Aiming to deploy low-precision instructions, we opt to use a packed format for the dense $B$ which is pre-formatted in VNNI layout (see lines 3-4 in Listing 5 where $v$ is the vnni blocking-factor). We implement this Block-SpMM TPP for x86 with both avx512-bf16 and AMX-bf16 acceleration, and for AArch64 with both BF16 dot-product and BF16-MMLA instructions.

The PARLOOPER kernel in Listing 5 for block-SpMM is similar to the regular GEMM. Lines 7-9 declare the loops of the sparse GEMM, while the main compute kernel is the bcsc_spmm_tpp sketched in the previous paragraph. This TPP gets as input arguments the sparse matrix A in BCSC format (line 16) and matrices $B$ and $C$ are regular dense tensors.

The BERT model is a transformer pre-trained via a mixture of masked language modeling objective, and next-sentence prediction [36]. We implemented the end-to-end workload using the architecture from Hugging Face [39]. We implemented four fused layers as PyTorch extensions in C++ using PARLOOPER and TPP: Bert-Embeddings, Bert-Self-Attention, Bert-Output/Bert-SelfOutput and Bert-Intermediate layers. In Listing 6 we show an implementation of the Bert-Output/Bert-SelfOutput module with PARLOOPER/TPPs. We invoke a BRGEMM TPP over blocked tensor layouts, and fuse bias TPP, dropout TPP, residual add TPPs and layernorm-equation TPPs [12] on a small 2D-block granularity to maximize the out-of-cache-reuse of tensors among subsequent operators [40, 41]. Unlike previous work [12] where the nested loops have a fixed ordering and parallelization, we abstract the specific instantiation with PARLOOPER and auto-tune them for the problem sizes of interest. In a similar fashion, the Bert-Self-Attention layer consists of blocked tensor contractions, fused with scale, add, dropout and softmax TPP blocks. The Bert-Embeddings layer is comprised of embedded lookups, followed by layernorm and dropout TPPs. Finally, Bert-Intermediate layer is composed of BRGEMM TPP, cascaded by bias add and Gaussian Error Linear Unit (GELU) TPP. All modules are implemented via the PARLOOPER/TPP paradigm and are auto-tuned for specific shapes.

Further, by composing the aforementioned Transformer building-blocks in different ways we can build inference LLM architectures/pipelines like GPT-J [37] and Llama2 [38].

We extend the BERT PyTorch extension described in Section IV-A to leverage the Block-SpMM kernel for the tensor contractions instead of dense BRGEMM TPPs. To get a block-sparse BERT model, we follow the methodology of the recent work [42]. In short, we obtain an unstructured block-sparse BERT model from a densely trained checkpoint, by applying knowledge distillation and block-wise weight pruning. We fine-tuned the block-sparse model for 40 epochs, and the final sparsity target was achieved in incremental fashion [42].

For ResNet-50 [22] training we integrated our PARLOOPER CNN kernels into PyTorch as C++ extensions. We followed the architecture from the original paper [22], where the convolution layers (see Listing 4) are followed by batch-normalization layers (see prior work [12] for batchnorm with TPPs). For the Fully Connected Layer (Listing 1) and the Pooling layers we simply used their respective TPP implementation (also through the PyTorch C++ extensions).