Weight Block Sparsity: Training, Compilation, and AI Engine Accelerators

Consider $\Gamma$ and $\Omega$ two matrices in $\{0,1\}^{N\times N}$.

More precisely, consider non-zero blocks of size $k\times k$ so that

Thanks to the sparsity and if we store only non-zeros, then $\gamma_{i*N+k}$ and $\dot{\omega}_{j*N+k}$ are contiguous. There will be a meaningful product to compute if and only if $\gamma_{i*N+k}=1$ and $\dot{\omega}_{j*N+k}=1$. We merge-sort these vectors. See how the sparse sparse matrix multiplication using Coordinate Block Structure (COO) is applied in Figure 3. We provide software to reproduce this [7].

In the case of achieving a fixed sparsity (i.e., density) of 50% for a square matrix of size $N$ and we choose the block size $k\times k$. The larger $k$ is, the smaller the overhead will be. The relative performance of the $k^{3}$ multiplication is better as $k$ get larger because spatial, temporal locality, and further optimized code for a constant/parameterized such as $k$.

In Figure 3, we present two scatter plots: on the abscissa the effective multiplication-and-addition number, on the ordinate the performance in GFLOPS, when the sparse matrix with dense block is $1\times 1$ (above) and $8\times 8$ (below). Given the same task, we deploy more threads and thus the vertical effect (AMD 16 cores Threadripper). With the same number of effective operations, the block permits and exploits higher GFLOPS per effective operation (Float is 2x faster than Double precision and this can be emphasized further [11, 19] and [18], for AIE2 we can do 256 multiply accumulate per cycle using 8 bit operands, 128 per 16 bits, 64 per 32bits).

We target convolutions first and without quantization. We take a model and we create a copy where we enhance the convolution with a (non-trainable) $\Gamma$. $\Gamma$ is the sparsity ratio and mask. This experiment is conducted using Keras[5] and the code is available [6]. A convolution will have three parameters (saving the model into a different format). The forward computation is modified so that the weights used for convolution are $\Gamma{W}$. We assume the backward computation (i.e., gradient) is done automatically from the forward definition. There is no need to change the bias. For example, we take Resnet50 from the Keras application repository, we start with a $\Gamma=1$, and we trained one epoch using Imagenet dataset [9]. The goal is to choose $\Gamma$ in such a way we achieve the required sparsity and the minimum loss in accuracy. We start from $W_{0}$, the current optimal weight, then choose values in $W_{0}$ to make zeros using different approaches such as incremental, Fisher measure, Hessian, diagonal Hessian, and custom penalty losses.

For example, for every convolution, $n_{i}=\sum\gamma_{i}$, which is the number of blocks, reduce this number to $\frac{n_{i}}{2}$. For each epoch (say every two training epochs), we consider the current non-set-yet mask elements $\sum(1-\gamma_{i})=k<\frac{n_{i}}{2}$. We compute our importance measure for all in ascending order. This time, we zero the first $\min(\frac{5}{100}k,1)$. We keep this process until we reach 50% sparsity. At each iteration at least one block will be set to zero. We trace a solution path as much as geodesic as possible.

If we want to make $\Gamma$ part of the optimization process as trainable variable we could introduce a penalty function into the loss ${\bf\ell}(x,w)+\lambda(w)$. First let us introduce an approximation for the $\max(x)$, so when in this section you will read $\max,\min$, this is a log sum exponential, which is continuous, derivable everywhere, and convex:

With $T$ we represent the number of non-zero block in $\Gamma$

This is a simplified loss so that we minimize the value of $\Gamma$ but also try to maximize the difference of the elements.

This last penalty function represents our attempt to state that the $\gamma_{i}$ should be positive and in the interval $[0,1]$ and in a such a way that we maximize the distance of the elements between 0s and 1s.

If we have $N$ parameters/weights, the Hessian $H\in R^{N\times N}$ has quite the space complexity (consider even small models could have million parameters). When we are already close to the optimal solution or we are trying to move from the optimal solution without using information from the gradient, the Hessian provides the most information close by an already established solution point. There are also ways to compute the Hessian and the effects of the Hessian by using Fisher information [26, 27, 28]. This will reduce to the computation of the gradient of the loss function.

We applied a Fisher measure and computed $\nabla^{2}{\bf\ell}$, that is computed just the diagonal of the Hessian. Again, we use the $L_{2}$ over the normalized weight and went through the process of training. The elements of the diagonal are not representative in general, but they are a good approximation of the contribution of a single weight. The Fisher and $\nabla^{2}{\bf\ell}$ did not provide any main advantages. But this information is very useful in quantization and optimizations within the same field and application.

Take a convolution with $\Gamma=1$ and weights $W$. Once again for each $\gamma_{i}$, this will be representative of a block $W_{i}\in R^{8\times h\times w\times 8}\sim R^{8\times 8}$. We can choose the $W_{i}$ using a measure of importance:

• •

  $L_{2}(W_{i})=\sqrt{\sum_{k}w_{k}^{2}}$ with $w_{k}\in W_{i}$,

• •

  $L_{1}(W_{i})=\sum_{k}|w_{k}|$ as above,

• •

  Variance $\sigma^{2}=\frac{1}{64}\sum_{k}(w_{k}-\mu)^{2}$ with $\mu=\frac{1}{64}\sum w_{k}$, with $w_{k}\in W_{i}$ oder $\frac{1}{N}\sum w_{k}$, with $w_{k}\in W$.

We can then sort them in ascending order. We set the first half to zero. Then we start re-training. We do this for the entire network or for one convolution at a time.

In Table I, we show the results by using one-time mask and full training: VGG-16, ResNet-50, Inceptionv3 on ImageNet20 (20 classes) and ImageNet1k (1000 classes). We use three samples per class for the validation accuracy for ImageNet1k data set; instead, we use 50 samples per class for ImageNet20. Fine-tuning sparse networks on the original ImageNet data-set [9] is expensive. To reduce the training time, we chose 20 classes (from the original 1000 classes) with the least number of images per class in the training data-set and this choice will affect the accuracy because there are fewer samples for re-training.

Classification accuracy on ImageNet1k drops by only 1-2% after applying 50% sparsity with a $8\times 8$ block (this is without any quantization). We experiment with different block shapes such as $16\times 4$ and $4\times 16$ on ResNet-50, but the accuracy is slightly worse. The different shape has the same volume and either more parallel vector operations or longer vector computations. This last process based on PyTorch has been the most accurate while the other approaches had a drop of 7 points in accuracy at least. It is not clear if it is a methodology weakness or just bad execution. It is important to us to present all the tools attempted.

Fine-grained sparsity ($1\times 1$ block or unstructured) does not sacrifice any accuracy (i.e., almost any). This is not equivalent to 2 over 4 (or 4 over 8) sparsity now available in GPUs.

As shown in Figure 4 as an example, we work with a mesh of 4x4 tensor cores connected by 4 horizontal and 4 vertical interconnections. We present estimates for square 2x2, .. $i\times i$ .. 8x8 and rectangular shapes are in the works ($4\times 1$, $4\times 2$, and $8\times 2$ into a $8\times 8$ with 2 Memtiles per column). Each core has 8 banks memories for a total 64 KB. About six banks are used as input/output/weight buffers and two banks are used as temporary space for kernels. Each core can request and send data to its direct neighbors (if aware of connection and control but this utility is not used here). Double buffering using ping/pong is used for inputs and outputs.

There are four Memtiles: each 512 KB and each is connected to one columns and its direct neighbor column, or it is connected to a row and its neighbor. The total amount of space is 2 MB. Memtile is a circular buffer to exploit more flexible allocation. Note a $2\times 2$ architecture will have one Memtile per column and a total of two Memtiles (1 MB).

A Memtile can broadcast data per column or per row; it is a design choice. We can dedicate one Memtile for weights, one for activations, or we can share it. In this work, we present results for shared Memtiles. To maximize the computation parallelism, every core will write data per column into Memtile.

The computation is split by Memtile and thus by column (cores columns). The output tensor is computed and split evenly by width. Thus one Memtile will store one partial tensor by width, each core will compute different output channels, and the computation streams the output tensor by rows and using ping/pong double buffering. We prioritize to reuse weights in core. The cores set is a cohort and we always choose symmetric computations.

If we have the inputs, output, and weights in Memtile, the minimum computation is one output channel and one row (i.e, by height). If this is not possible, we try to reduce the size of the width (e.g., shaping the tensor in Memtile by using DDR spills) and we can manage to split the input channels and to split the weights accordingly and prepare for accumulation. We call W-Split the distribution of tensor by columns in the Tensor mesh. We can $C_{OUT}$-split, this requires the partial transfer of weights. We can $C_{IN}$-split when we need to split by input channel, this is the last resort because it is also the most expensive accumulation of the outputs. $C_{IN}$-split can be implemented as a graph optimization by splitting the convolution into two and then use an element wise operation to combine the results, which can be done recursively.

The subvolume describes the smallest shape of the weights that we need to manage and the largest computation in the core. We compress the weight accordingly. Any data movement will always be a multiple of the subvolume, a multiple of $8\times 8$, and it is a single load. Such a compressed data will have the same properties whether it is sparse or dense.

During the scheduling of each layer, we evaluate what tensors can fit in Memtile. Here, activation and weight tensors share the space. At each step, the memory allocation will check if we can allocate inputs, weights, and outputs. If we cannot, we evict all tensors into DDR and then split the time of the computation.

At the end of this stage, every tensor will have an address in Memtile or DDR (or both). If there are only DDR addresses, the compiler will take the basic layer computation and, by heuristics, will split the computation and the output tensor by width, output channel, height, and input channel (no necessarily in this order). The heuristics have a single objective to find the largest problem fitting the (each) memory level. We deploy a recursive approach of tiling. Formally, $\dot{\sum}$ is a parallel loop and a W-split can be written as follows:

The split is a function of the footprint. Before and after each convolution, there will be an explicit data movement as option. At this stage each input, output, and weights have addresses associated with each sub-computation. Then the code generation of each $Conv(X_{w},W)$ is independent and recursive as needed. This is a tree-like structure. If the convolution has strides or a large kernel, each sub-convolution has overlap data; however, the sub-convolution has defined addresses and data movements. For a W-split such as this, we are computing the output by rows and the weights are reused (read once).

The compiler creates a list of operations. These operations become smaller and smaller and then can be executed from Memtile to Memtile. There is a further decomposition using only the Tensor cores and it is completely determined by the subvolume. Here, we show how we generate code at this level and estimate time as in Figure 5. This is the computation of a convolution with top/bottom padding by height:

An important feature of the current system is the concept of iteration between Memtile and core. Using locks and chaining the locks, we write a single instruction from the prospective of a single core (as a SIMD instruction) and driving all cores at once for multiple iterations $\dot{\sum}_{h=1}^{i}Conv(X_{w})$ in Equation 14. The ASM-like code follows:

  LOADFM Lock k_0 Memtile addr core addr iter i
  CONV iteration      i
  WRITEFM Lock k_1 Memtile addr core addr iter i

There is an implicit lock (say k_x) that is used for the pong and the system cycles between locks (k_x and k_0). These three operations execute a number of iterations i and, using a ping/pong, they will load different slices of data and compute different slices of data.

Equation 14 is encoded as follows:

  ## Head top pad < 50 us First comp block
  LOADFM Lock k_0 Memtile addr_0 core addr iter 1
  CONV iteration 1
  WRITEFM Lock k_1 Memtile addr_1 core addr iter 1
  ## Body iteration > 50 us < 150 us
  ## k_0 -> k_2 -> k_4 Lock Chain
  LOADFM Lock k_2 Memtile addr_2 core addr iter 9
  CONV iteration 7
  WRITEFM Lock k_3 Memtile addr_3 core addr iter 9
  ## tail bottom pad > 150 us Last computation block
  LOADFM Lock k_4 Memtile addr_4 core addr iter 1
  CONV iteration 1
  WRITEFM Lock k_5 Memtile addr_5 core addr iter 1

We present in Figure 5 the execution estimate of this code. At this stage, we have all the information. Per layer, the code generation is a two pass process. First, we generate code for the all loads/stores. Second we combine them into chains with dependency, logically correct and as fast as possible.

We could estimate the time execution without a full code generation. When we annotate time information to a load, we have assurance that the load is a complete description of the DMA communication between multiple memories and as complex as the architecture. Actually, this is literally translated to a binary executable that perform the data movement.

We explain how we capture the execution time and visualize it as in Figure 5. We start by the time estimates for DDR to Memtile communications. We have two communication types: activations and weights. Per Memtile there are two dedicated channels.

• •

  If we share activations and weights in the same Memtile, we can use one channel for activations and one for weights. Thus the loads from DDR to Memtile (LOAD and LOADW) are parallel with a bandwidth of 4 GBps. Writes from Memtile to DDR (WRITE) can use both channels (8 GBps).

• •

  If activations and weights go to different Memtiles (for example weights to Memtiles ’0’ and ’3’ and activations to ’1’ and ’2’), each load is parallel and 8 GBps. Writes are identical.

The Memtile connections with AIE cores are different. We assume a few channels with 4 GBps bandwidth. One Memtile can broadcast inputs to a cores column. These communications are for activations (LOADFM). One Memtile can broadcast to rows of cores, these are for weights (LOADWM). We assume that the column and row communications are parallel.

Every communication with iteration one is synchronous and sequential. The load, convolution, and store is executed one after the other and every core is independent. For synchronization and for bookkeeping, we assume that Weights communications from Memtiles to cores are synchronous and halting (LOADWM).

Every communication with iteration larger than one, we assume that load (LOADFM), computation (COMP), and store (WRITEFM) are executed in parallel and the overall execution time is the maximum of the estimated time multiplied by the number of iterations.

We estimate the execution time of a subvolume (COMP) by the number of operations divided by the maximum number of operations per cycle which is in our scenario is $4\times 8\times 8=256$ operations per cycle and 1 GHz frequency. Sparsity reduces computation and communication time.

We do not account the wait and synchronization which are necessary to reprogram the fabric. These are very expensive running on for a few milliseconds.

We present the time estimate for a convolution with padding, dense, and with 50% sparsity, see Figure 5 and 6.

For these convolutions, the computation dominates the execution time. Sparsity cuts the execution time by half: from 200 $\mu s$ to 130 $\mu s$. On one hand, there are convolutions that realize up to $2\times$ performance; on the other, there are convolutions that are dominated by the reading or writing. In the latter case, sparsity helps in space saving and probably DDR tensors spilling. In principle, we could relax sparsity requirements for those convolutions that are communication bound (and restart training). In Figure 7, we provide the time estimates for Resnet50 using sparsity 50% and using $4\times 4$ AIE array.

Assume, we have an output tile size $u_{o}$. We repeat the projection process using the schedule. This time, we keep information of the active tensors at any time in the schedule. Projections and active tensors are sufficient to represent the maximum active space requirement (not including input and output). If we want zero DDR spill for inputs and outputs, we choose a output tile for which the maximum active space is smaller than Memtile. Say we have M tiles. With an architecture for which computation is free (see [14, 3]), we write $T_{j}$ once, we read $T_{0}$, and we reread $(M-1)(k-1)$. This is better than $\sum_{l=0}^{j}T_{l}$ especially for large $j$s.

A compiler can target easily a tiling with zero DDR communications (for layer output tensors). In practice, zero DDR communications is not our final goal.

Take the schedule $L_{0}$ …$L_{j}$, let us apply the live tensor analysis and tensor projections for a specific tile size $u_{o}$ and $M$ tiles. We add two custom layers as follows

The layer $I_{b}(M)$ is an iteration layer or input boundary that takes the input tensor $T_{-1}$ and carves out a sub-tensor shape and size $(u_{o}-1)s+k$. The layer $I_{e}(M)$ is an iteration layer or output boundary that takes the input tensor $u_{o}$ and copies it to a tensor of shape and size $T_{j}$. For the purpose of the memory allocation, we have a legal schedule that describes basically the property of the computation. We can do memory allocation.

A computation is basically defined by its layer and by the addresses in memory of its input and output tensors. The code generation unrolls the iterations as it was a loop and $I_{b}$ and $I_{e}$ generate the proper sub tensor addresses and specifying the number of iterations.

The code generation will be for the schedule:

With a little of book keeping we can write the time estimates as a time series and compare performance.

The depth-wise tiling has the tendency to increase the computation number and decrease the DDR loads and stores. As fomulated, this tiling may read multiple times the weights. Sparsity reduce the number of computations and the weights communications. Together have the opportunity to work better.

We present results for VGG16 because of simplicity and each layer tensor do not fit the three Memtile (of a $3\times 3$ system). We can apply the same approach to Resnet and inception. The generalized convolution idea is applicable.

We take VGG and we instruct the DDR to be 16 times slower (4GBs/16) highlighting the need of fewer DDR communications. We take only the part that requires DDR spills for each layer. In Figure 8, we show the performance for VGG16 using DDR spills for each layer: 0.025s total time.

We apply depth-wise tiling using two tiles and we achieve better performance at 0.022s. see Figure 9. Sparsity by itself without any tiling can achieve only 0.021s. Sparsity and tiling improves even further and we achieve 0.014s, see Figure 10. We can appreciate the reduction of activation DDR communications thanks to depth-wise tiling and the reduction of computation and weights communication by sparsity.