UNIT: Unifying Tensorized Instruction Compilation

Deep learning is computationally expensive, requiring substantial compute and memory resources. As deep learning becomes more pervasive, researchers are designing both software and hardware techniques to reduce the compute and memory burden. A widely adopted approach in this context is using mixed precision for expensive operations, e.g., convolution or dense operations [24, 18]. In practice, this means representing 32-bit floating point (fp32) operands with a lower bitwidth datatype - 16-bit floating point numbers (fp16) or 8/16-bit integer numbers (int8, int16). To keep the accuracy in check, it is helpful to accumulate the results in higher precision (fp32 or int32). This type of mixed precision computation is often called quantization for integer values [18]. In this paper, we will always use mixed precision for brevity.

While using mixed precision data types reduces memory footprint, it might not necessarily lead to performance improvement. To investigate this, we conducted an experiment to compare the performance of Nvidia cuDNN performance for fp16 and fp32 in the absence of Nvidia mixed precision tensorized instructions (Tensor Core). As shown in Figure 1, we observe that blindly using mixed precision leads to substantial slowdown because of the overhead of casting between two data types.

Therefore, mainstream hardware vendors (Intel, ARM and Nvidia) have introduced mixed precision tensorized instructions to achieve better performance. These instructions add mixed precision arithmetic support where operands are of lower precision while the accumulation happens in higher precision, potentially leading to 2$\times$ - 4$\times$ speedup. The most popular examples of these tensorized instructions are Intel VNNI, ARM DOT and Nvidia Tensor Core. We will discuss the semantics of these operations in Section III.

Hardware vendors have a long history of adding new instructions to accelerate important applications. However, the mixed precision tensorized instructions introduce a unique idiom - horizontal accumulation. These tensorized instructions typically conduct a sequence of elementwise multiplications governed by a memory access pattern, followed by a horizontal accumulation. The accumulation is termed horizontal because all values to be accumulated are present in the same vector register. For example, as it is shown in Figure 2(a), Intel VNNI executes a dot product of two vectors, each having 4 int8 elements, while performing the accumulation in int32. We observe a similar pattern, though with different numbers of entries and data types, for Nvidia Tensor Core (in Figure 2(b)) and ARM DOT instructions (this is omitted, because it is similar to VNNI).

Though tensorized instructions seem promising, their adoption pace is limited because of the absence of an automatic technique that can detect and use these instructions seamlessly. Currently, their usage in the deep learning domain is limited to hardware vendor libraries like Intel oneDNN and Nvidia cuDNN, which may provide high performance for the pre-defined operations but are inflexible as discussed in Section I.

Similarly, conventional loop vectorizers find it hard to exploit the profitability of these tensorized instructions, as they are not designed to work with the horizontal reduction idiom. Conventional loop vectorizers in general-purpose compilers like GCC and LLVM mainly focus on either analyzing the innermost loop body or combining instructions in the unrolled loop bodies. When it comes to the horizontal reduction idiom, these compilers often reorder the computation and generate epilogue reduction, preventing us from using the tensorized instructions.

There have been some recent works in compiling programs to leverage tensorized instructions. PolyDL [36] generates CPU programs for convolution kernels in neural networks that call a GEMM micro-kernel using Intel VNNI instructions. Bhaskaracharya et al. [35] generate CUDA programs for matrix computation leveraging Nvidia Tensor Core. However, these works are limited to one platform and its specific instruction, which lacks generalizability. A generic solution to handle tensorized instructions from multiple platforms together is still missing.

Compilers often have multiple levels of intermediate representation (IR) to express the program; each level is designed to enable different analyses and transformations. In this section, we describe the background of a tensor domain specific language (DSL) and the multi-level IR.

In order to unify the compilation of tensorized instructions from different platforms and keep the system open to integrate new instructions, the first question to answer is how to have a unified description of the semantics of tensorized instructions. As explained in Section II, we employ ubiquitous tensor DSL and tensor IR to solve the abstraction problem. All mixed precision tensorized instructions perform some elementwise operations for vectors, followed by a horizontal reduction. Each tensorized instruction, therefore, can be regarded as a small tensor operation program written in the tensor DSL.

Figure 4(a) shows how an Intel VNNI instruction is described in the tensor DSL. Three source operands of Intel VNNI are 512-bit registers. Two of them are 64 lanes of unsigned 8-bit integers (uint8) and signed 8-bit integers (int8), and the other one is 16 lanes of signed 32-bit integers (int32), which correspond to the tensors a, b, c we defined. The arithmetic behavior is defined by the loop variables and the expression of d[i]. Here we annotate that loop i is data parallel, since these 16 elements are independent from each other; loop j is reduction since for every independent element it sums up 4 elements along with this loop. A similar loop pattern appears in the other tensor operations shown in Figure 5. The description of ARM DOT, shown in Figure 4(b), is similar to Intel VNNI, with a different number of lanes and data types.

Nvidia Tensor Core, on the other hand, performs a $16^{3}$ square matrix multiplication as shown in Figure 4(c). Comparing with (a) and (b), a key difference is that it requires the accumulator register to be the same as the addition register (note the +=). This is due to the data type opaqueness of the Tensor Core instruction, which prevents us from giving arbitrary initial values for the accumulators.

We describe the semantics of each tensorized instruction in tensor DSL. The deep learning compiler pipeline parses the operation into tensor Op, which preserves tensor information like the expression tree, the loop trip count, and the array buffers. This information is essential for the analysis and transformation passes in Inspector and Rewriter.

To determine if a tensorized instruction can be applied to a tensor operation, the Inspector pass uses a two-step approach. It first determines if (part of) the tensor operation program and the instruction can be arithmetically equivalent by checking a form of isomorphism between their associated expression trees. After that, it inspects the data access pattern to confirm the assembly operands can be prepared so as to guide the Rewriter transformation.

There are three phases in the code transformation: loop reorganization, tensorized instruction replacement, and tuning.

The inspector pass is implemented by analyzing TVM’s ComputeOp data structure. This matches the expression tree of both the instruction and program and enumerates mappings between the loop variables. We enumerate the loops from the tensor’s innermost dimension to outermost dimension, and greedily return the first eligible one because of the better potential data locality for inner dimensions. The enumerated mapping provides us with the correspondence of loop variables between the instructions and the tensor operations.

These following steps will be performed by the rewriter:

1. 1.

   According to the loop correspondence analyzed by the inspector, we reorganize the loops to be tensorized by tiling these loops by the trip counts of the corresponding loops in the instruction, and reorder them to be the innermost loops. These loops will be annotated by a tensorize pragma to hint the instruction injection.

2. 2.

   Based on the strategies discussed in Section III-C, we reorganize the loops above not involved in instruction rewriting to tune the performance.

3. 3.

   We lower the manipulated loop nest to the tensor IR, and replace the loop body annotated with the tensorize pragma with the target instructions, as shown in Figure 5(c).

Steps 1 and 2 are achieved by invoking TVM scheduling primitives on the tensor DSL level, and step 3 is a tensor IR transformation pass.

Next, we discuss the implementation of the tuning strategies discussed in the last section.

On the other hand, greedily applying each specialized optimization does not always improve the performance. Though dimension fusion may save the memory traffic, it also introduces software overhead on data rearrangement. Similarly, though splitting the reduction loop introduces more parallelism, it also introduces thread synchronization overhead and register pressure. We enumerate each parameter, including the degree of reduction parallelization and whether to fuse the width and height dimensions, and then apply these transformations to the program and profile the performance to determine which transformation leads to the best performance.

We assess UNIT on three hardware platforms:

In the evaluation, we target the N=1 cases, because it is hard to optimize but critical for inference use cases. Comparing with batched cases where N>1, we cannot reuse the kernel tensor across samples, or exploit the parallelism brought by the data-parallel batching dimension.

In this subsection, we show the UNIT end-to-end effectiveness on Intel x86 and Nvidia GPU processors for tensorizing mixed precision instructions. For Intel x86 experiments, we use MXNet integrated with Intel oneDNN (referred to as MXNet-oneDNN) as the baseline. Another comparison of ours is TVM with manually written schedules using Intel’s VNNI instruction. The findings of this experiment are shown in Figure 9.

We observe that UNIT achieves significant speedup compared to MXNet-oneDNN. Note that Intel oneDNN has access to manually written schedules that have been aggressively optimized and tuned by domain experts. We also observe that TVM overall achieves better performance than MXNet-oneDNN, but has suboptimal performance on resnet50 and resnet50b, which were heavily tuned by oneDNN engineers. On the other hand, UNIT outperforms both baselines, by 1.3$\times$ over MXNet-oneDNN and by 1.18$\times$ over TVM.

Next, we test the efficacy of UNIT on utilizing Nvidia Tensor Core instructions for Nvidia GPUs. For the baseline, we integrate TVM with cuDNN, which has access to manually written aggressively tuned Tensor Core schedules. The findings of this experiment are shown in Figure 9. We observe that UNIT consistently achieves better performance than cuDNN with a mean speedup of 1.75$\times$ and up to 2.2$\times$.

In this subsection, we focus on the convolution operators of the DNN models to perform an in-depth analysis of the impact of different optimization techniques used by UNIT’s Rewriter. This is essentially an ablation study, showing how important different parts of UNIT are. There are 148 different convolution workloads (i.e., convolution with different feature map sizes, kernel sizes, strides, etc.) in the models, out of which we choose 16 representative convolution layers. These kernels cover diverse input shapes and strides. Other workloads behave similarly in the ablation study. We summarize the characteristics, namely, convolution attributes, like shapes, strides, etc., of the selected workloads in Table I.

First we fuse outer loop nests such that the loop bound of the fused loop nest is $<$ 3000, and measure the latency of the resulting kernel (shown by Parallel). Then, we take the remaining loop nests, and tile and unroll them such the unrolling factor is $<$ 8, and measure this performance (shown by +Unroll). Finally, instead of setting the limits as 3000 and 8, we tune the search space and measure performance (shown by +Tune), getting the final latency UNIT achieves. We observe that Parallel and Unroll together is responsible for most of the speedup. The additional speedup introduced by Tuning is quite small. It turns out that more than half of the kernels get the optimal performance on the first tuning pair (i.e. 3000 and 8), and more than 95% of the kernels get the optimal performance within the first 8 tuning pairs.

CPU does poorly on workloads #1 and #4, because their output shapes (OH/OW) can neither be perfectly tiled nor fully unrolled. Inherited from TVM, loop residues are handled the by guarding it with a likely clause, which results in an if-branch that harms the performance.

According to our evaluation, any unrolling degree (p in Figure 6) larger than 2 may overwhelm the registers, so we use p=2 to apply the generic optimization. The generic optimization already beat cuDNN in most cases (shown by Generic). Then, depending on the height and width values, Rewriter fuses the height and width dimensions to save memory bandwidth (shown by +FuseDim). Then, we split the reduction dimension K by 64 and measure the performance (+SplitK). Finally, we let Rewriter to choose the sizes for these 3 optimizations and measure performance (shown by +Tune).

We observe that SplitK leads to the maximal speedup, as it leads to significant parallelism and keeps the Tensor Cores busy. More than 70% of the kernels can get high performance by employing fusion and parallelizing the reduction dimension. Similar to CPUs, the additional speedup by tuning is small.

UNIT cannot outperform cuDNN on #1 and #15, because the strided data accesses lead to less data locality. However, since these adversarial cases (both CPU and GPU) only occupy a very small portion among all these models, we can still outperform vendor-provided libraries because of the generality of our optimization.

We evaluate the extensibility of UNIT in two aspects: to new hardware platforms and to new deep learning tensor operations. We observe that by just representing the semantics of the new tensorized instruction in tensor DSL, UNIT can easily extend to new tensorized instructions and tensor operations.