Joint Program and Layout Transformations to enable Convolutional Operators on Specialized Hardware based on Constraint Programming

Conceptually, our approach is based on dataflow graphs (DFG). A DFG represents every scalar operation necessary to perform a computation as a directed graph. Nodes represent operations, and edges the dataflow. Formally:

Further, a DFG has the following properties:

• •

  Nodes with only $data$ labels generate outgoing $sequential$ edges for each operation consuming the data. They have no incoming edges. They represent the input values of the computation modelled by the DFG.

• •

  Nodes with $operation$ labels perform a scalar operation, consuming the data of the incoming $sequential$ edges and produce one or more outgoing $sequential$ edges.

• •

  Commutative reduction operations are modelled by a $sequential$ self-edge. This optimization can reduce the number of nodes and edges in a graph significantly, without loosing correctness of representation.

• •

  Nodes with $operation$ labels with only an outgoing self-edge, or no outgoing edges, represent the computation results.

• •

  Nodes labelled $operation$ can have bidirectional $spatial$ connections to other nodes, performing the same computation, but for a different output element. These connections indicate parallelism in the computation. Potentially, $spatial$ edges lead to a set of $k$ fully connected nodes. The number of edges can be reduced by pruning the connections to create a graph with one internal node and $k-1$ leaves. In this star subgraph, the transitive property maintains the parallelism information.

For the tensor-based workloads in a DNN, every output element is computed by the same sequence of operations, but with a different subset of input values. Figure 1(c) shows the partial dataflow graph of a 4x4 matrix multiplication. Each subgraph computing one output element shares a set of input nodes with its neighbours but no intermediate results. This property removes the need to cover the full operator DFG with a sequence of instructions, but instead solves the problem for a small subset and then uses a hardware-dependent inference step to determine the structure of the full program. We will demonstrate this in sections 5 and 6.

Embedding an instruction DFG $G_{i}=(N_{i},E_{i},l_{i})$ into an operator DFG $G_{o}=(N_{o},E_{o},l_{o})$ is equivalent to the subgraph isomorphism problem. For an embedding we need to find an injective function $f:G_{i}\rightarrow G_{o}$ that describes a distinct subset of nodes and edges in $G_{o}$ that exactly matches $G_{i}$, formally described as: $\forall(s,t)\in E_{i}\Rightarrow(f(s),f(t))\in E_{o}$. This function has to maintain the labeling, such that $\forall s\in N_{i}:l_{i}(s)\equiv l_{o}(f(s))$. The main advantage of this approach is that the matching problem itself is not bound to implicit implementation decisions like loop ordering or data layout of tensors or access functions. This removes the need for transformations in the search to account for these decisions. Instead, we can derive these from the result of the embedding.

However, we also identified two challenges with this approach:

1. (1)

   With a space complexity of $O(|N|^{2})$ and $O(|N|+|E|)$, adjacency matrices and lists are not suited to represent a full DFG for operators like conv2D with billions of operations. Subsection 3.2 explains how a polyhedral program representation is used to abstract the workload while maintaining the detailed information of the basic dataflow graph.

2. (2)

   Solving the embedding only with subgraph isomorphism is too imprecise of a solution formulation. Often, there are additional restrictions to the available hardware and its software interface. A possible implementation needs to be inside this restricted space in order to be valid. Section 4 explains how constraint programming is leveraged to overcome this problem.

The first challenge is addressed with a more concise program representation, inspired by the polyhedreal model, a powerful compiler methodology capable of expressing computations in quasi-affine loop nests. Several tools in the Deep Learning community leverage this representation to generate efficient DNNs kernels. TensorComprehensions (Vasilache et al., 2018) targets GPU optimizations and MLIR (Lattner et al., 2020) provides a whole polyhedral dialect for loop optimization. Polyhedral program representation is an abstraction of loop based programs, with its components abstracting the program and data dependencies on a scalar level (Verdoolaege, 2016):

• •

  The instance set $S$ is the set of all dynamic execution instances. $S$ is described by a set of integer tuples. Each tuple $s$ describes exactly one dynamic execution instance.

• •

  The data dependence relation $D$ is the union of all binary relations between pairs of instances in $S$.

Reflecting this on our DFG representation, the node set $N_{o}$ of $G_{o}$ is $S$ and the $sequential$ edges model the data dependence relations $D$. We represent $S$ as a set of integer tuples. Every tuple is a specific operation happening in the instance set. To describe the sets of integer tuples, we use the notation

where the fixed lower and upper bounds of each tuple element $e_{j}$ are defined by a condition $\tau_{j}$. The conjunction ($\wedge$) of all $\tau$ terms contains the full set. The data dependence relation $D$ is represented by a binary relation between two sets. A relation maps elements from the source set to the target set. Relations are denoted as

where every term $\Phi_{k}$ defines a condition in the relation. Elements in the target tuple are denoted $e^{\prime}$. The relation condition $\Phi_{k}$ is used to describe any source element $e$ that maps to element $e^{\prime}$ in the target set. The conjunction of all $\Phi$ terms encapsulates the whole relation domain.

To describe operators with this polyhedral representation, we move from an explicit DFG to a set-based representation. Every element from the original tensor expression $T$ (figure 1(a)), like an arithmetic operation or input tensor, is modelled by a domain set $d\subset S$, as in equation 1. For example, the domain $d_{*}$ contains all multiplication nodes in $G_{o}$. For the matrix multiply example in figure 1(a), all dynamic execution instances and input tensors are contained in the sets

where $I,J,K$ are the domain bounds in figure 1(a) and loop bounds in figure 1(b), respectively. Since the goal is to compare two different programs, we model the sequence of individual arithmetic operations explicitly. For this, $S$ contains an additional dimension $t$ that describes the order of expressions in $T$. To do this, we assign each expression (multipy, add) in $T$ an integer value. For example, we explicitly model the statements in lines 4 and 5 in figure 1(b) as individual points in the dynamic instance set. Input tensors are defined by a set of their shape. Now, every node in the original DFG is contained in a union of the above sets, for the example this means $N_{o}\equiv S\cup X\cup Y$.

Tensor access functions and dataflow are modelled by binary relations between two domains. There is a dataflow between two instances $(s1,s2)\in S$ if there exists a relation for which $s1\rightarrow s2\neq\emptyset$. The dataflow of the example in figure 1(d) is described by the following relations:

Relation $R1$ specifies that the multiplication and addition happen in the same loop iteration, but are ordered by their textual position in $T$, specifically that the node following the multiplication needs to be an add operation. $R2$ can be interpreted such as two add operations happen sequentially in iteration dimension $k$ and that the same add operation is performed. This is the self-edge in the original DFG. To accommodate commutative operations, the term controlling the reduction order can be relaxed. The $spatial$ edges of the original DFG are modelled in the same way. $A_{X}$ and $A_{Y}$ encode the access function itself and by which node the access is performed, in this case the multiplication in $T$. The union of all relations describes the edges of $G_{o}$, and specifically for the example, $E_{o}\equiv R1\cup R2\cup A_{X}\cup A_{Y}$. Bringing the sets and relation together with the original labelling function $l_{o}$ creates the expression graph shown in figure 1(d).

Not all of the relations in this representation are symmetric, being the same in both directions. The relation $*\rightarrow X$ is surjective and functional, meaning that every multiplication can exactly map to the one tensor element in $X$ it consumes. However, the inverse relation $X\rightarrow*$ is non-functional. Every input element in $X$ is used by multiple multiplications, but not all of them. The relation of $A_{X}$ and $A_{Y}$ reflects this, as they contain no term with properties for $i^{\prime}$ or $j^{\prime}$ respectively. As a result, for one specific input value in $X$, the relation describes the subset of all multiplications using this value.

This section will discuss how we represent the embedding as a CSP in $GeCode$³³3https://github.com/Gecode/gecode, the solver used for this work.

Over the space defined in 4.2 we can formulate constraints that specify solution properties. The most important constraint is matching the dataflow between operator and instruction. With other constraints we can further reduce the solution space.

This section presents the constraints placed upon the problem space defined in the previous section.

The previous sections described how the problem is modelled in CP, now we will discuss the search used in the actual solving process. Variable and value selection strategies determine how the problem space is explored by determining which variable is assigned a value next. To better fit the underlying problem the strategies can be customized. To reduce the amount of choices necessary during branching it is desirable to trigger as much propagation as possible with every assignment, such that every domain gets subsumed as early as possible.

We use a variable selection strategy based on groups of node types in $G_{i}$. From the example in figure 1(d), all multiplications are in group $g_{*}$. Changing the order of groups can result in varying degrees of propagation. When starting with a group of input nodes, less propagation is possible due to the non-functional relations to their consuming nodes. Starting with $g_{*}$, every node assigned a value automatically propagates to its inputs as well as to the following add operation. Our implementation currently begins with the output variables and propagates backwards through the DFG, which proved to be a robust heuristic for short solver runtimes. The value selection strategy implements a lexicographic search through the domain.

Padding and other data transformations necessary for an embedding can create an overhead in the number of multiply-accumulate operations (MAC) and data movement. We add an optimization term to the constraint problem, minimizing the MAC and data movement overhead $O_{MAC}$ and $O_{Data}$, compared to the theoretical minimum:

In our work, we used the number of tensor elements for the data movement computation, not the absolute size. When treating $O_{MAC}$ and $O_{Data}$ as elements of a vector $\vec{\textbf{o}}=[O_{MAC},O_{Data}]^{T}$, the CSP optimizes the term $min(||\vec{\textbf{o}}\cdot\vec{\textbf{w}}||)$, where $\vec{\textbf{w}}$ is a user-defined weight vector. It can be used shift the impact of the overheads on the candidate selection, depending on the on the used hardware and technology. This eliminates solutions with padding or other overheads, when possible.

The evaluation is performed on VTA (Moreau et al., 2019), a hardware accelerator providing a matrix-multiply (GEMM) instruction with a corresponding processing unit, as well a vector unit for activation and scaling tasks. The hardware is instantiated on a Zynq Ultrascale+ FPGA. We use the default configuration provided by the authors with a 256kbyte weight buffer, a 128kbyte data buffer and a GEMM core with $8bit$ multiplications and $32bit$ accumulation. The GEMM unit computes $C_{xy}+=A_{xz}\cdot B_{zy}^{T}$ with $(x,y,z)=(1,16,16)$ and its result can be processed by a vector-scalar unit for activation and quantization operations. Notice that matrix operand $B$ is transposed. The hardware has a load/store direct-memory access (DMA) unit for independent memory accesses of matrix operands. It can read and write full 2D operand matrices stored consecutively in memory.

Evaluation workloads are from the Baidu DeepBench Inference Benchmark Suite⁴⁴4https://github.com/baidu-research/DeepBench, accessed 01.2021, providing 108 convolution operators from a range of domains, like image and speech processing. Twelve convolutions cannot be executed on the VTA accelerator without additional zero padding. Section 6 will discuss possible solutions for this problem in detail. Furthermore, 28 layers in the convolution benchmark cannot be processed by GeCode, the solver we implemented our approach in. It uses the default C++ $int$ data type representation of the used compiler, which is 32 bits in our case. Large convolutions can yield domains substantially larger than this limit. Additional 6 layers caused various errors in the TVM compiler, like a VTA instruction buffer overflow. This leaves 62 layers to benchmark.

The conv2d reference embedding of TVM maps the three axes $x,y,z$ of the GEMM unit statically to the batch $n=x$, output channel $oc=y$ and input channel $ic=z$ dimensions of the convolution. Each of these dimensions is split and moved to be the innermost dimensions of the input, weight and out tensors. This process is specified in a static template, applying the necessary transformations during code generation.

Figure 5 shows the loop and data layout transformations for a convolution with a 2D memory tiling, where each tile is a matrix operand that can be loaded/stored by the DMA. Step 1 is the baseline conv2d. The loops and tensor are tiled as explained before (fig. 5, step 2). The resulting implementation reorders $n_{i},oc_{i},ic_{i}$ to be the innermost loops. These loops are then replaced by a call to the GEMM instruction. The DMA load and store operations are then placed around the operation to continuously load values until an output segment is complete (Figure 5 step 3). The TVM convolution implementation expects a NCHW data layout. Implementing other workloads follows a similar pattern – input and output dimensions are tiled into matrices and moved to the inside.

When TVM encounters an operator that can be accelerated by the hardware, the implementation is handled by the specified deployment strategy. A strategy implements the operator optimized for the target in TVM’s IR. We build on TVM’s code generation for VTA by integrating our approach into VTA’s deployment strategy. It generates the instruction DFG $G_{i}$ based on the hardware configuration. From there it formulates the constraint program as explained in section 4. The nodes of $G_{i}$ become the variables, the operator’s dynamic instance sets the domain. We use the following constraints to generate mappings similar to the reference :

• •

  subgraph isomorphism: Match the dataflow of the GEMM instruction. This is the central constraint of the embedding problem.

• •

  hyper-rectangle: Ensure all input and output elements are mapped into an axis aligned shape. This allows simpler memory transformations based on transpose and reshape operations.

• •

  allDiff: Prevent the same dynamic execution instance from appearing multiple times in the same instruction call.

• •

  fixed origin: The first match of all input and output tensors is fixed to the origin of the respective domain.

• •

  dense: No input or output tensor is allowed to have a stride in any dimension.

• •

  linear memory access: Only allow matches in workload dimensions with a linear memory access. This excludes, for example, strides and stencil patterns.

We can use this constraint program to attempt to embed the VTA instruction into any workload, not just convolutions. The solver produces a list of tuples, describing how each operation and data element in $G_{i}$ maps to a node in $G_{o}$. The regularity of DNN workloads like convolutions allows us to extrapolate an implementation from this information. In the solution, the variables associated with the input and output values are evaluated to compute which dimension of the instruction is matched to which dimension in the workload and what the tiling factors are. For VTA, the code generation is then straight forward. The matched dimensions are tiled by the determined factors and moved to be the innermost dimensions and loops are reordered in the same fashion. These transformations are necessary to embed the GEMM instruction. Loops and tensor dimensions not part of the embedding are free to be transformed for further performance optimization. These optimizations include loop tiling, ordering or fusion. AutoTVM is used to automatically determine the best optimization parameters for every conv2d layer. Finally, code with embedded instructions is generated by TVM’s VTA programming tool flow.

We validate our approach by comparing the performance of a micro-benchmark with the TVM reference implementation. The benchmark performs tensor packing, convolution, activation and tensor unpacking. It is implemented as a Relay (Roesch et al., 2018) program. Packing and unpacking are performed by the ARM host CPU on the Zynq board. For the convolution operator, the TVM reference uses an expert-made implementation template that specifies which memory and loop transformations are necessary, as shown in Figure 5. This template expects a NCHW tensor layout. Our method automatically generates TVM code based on the found embedding and a specified target memory layout. As intended by the constraint program design, the implementation found by our solvers maps the instruction dimensions to the same workload dimensions as the reference. From this, a Relay program for the tensor packing is generated automatically. Both the reference and our generated solutions use AutoTVM to optimize performance. To ensure comparability, both versions use the same tiling configuration found by AutoTVM.

Our tool achieved performance competitive with the TVM reference, as demonstrated in Figure 7. After a warm-up, we averaged over 200 measurements for each layer. Across the benchmark, all but two layers perform within the reference’s standard deviation ($\sigma$) envelope. The mean $\sigma$ is $26ms$ for the reference as well as our solution. The absolute differences between the two approaches range from $0.1ms$ to $90ms$. These results show that our approach can compete with existing, expert-made implementations for a hardware accelerator target.

Dynamically changing the tensor layout of a DNN for global performance optimization, such as changing from NCHW to NHWC, for instance, is a topic of interest and already supported some by existing tools (Zheng et al., 2020)(Liu et al., 2019). However, for accelerators like VTA, only some of the dimensions are free to be rearranged, as the packed, innermost dimensions are necessary for the embedding. Since the solver determines which dimensions are necessary, and thus which ones are free, the memory layout of the free dimensions can be changed during code generation.

We demonstrate this by generating code for the NHWC format, including code for tensor packing and unpacking. Figure 7 compares the results of this experiment to the TVM reference, based on NCHW. Over all layers of the tested micro benchmark, the reference outperforms NHWC in only four cases. This effect correlates with the ratio between channels and image size, as layers with larger channels show better performance in the NHWC layout. This can be explained with the data movement during the layout transformation. The packing moves data from the $ic$ and $oc$ to be the innermost tensor dimension. When dimensions are closer to their target position, the read access has less stride, yielding better CPU cache utilization.

In the previous sections, twelve convolution layers in the benchmark could not be executed on the VTA without zero padding the $ic$ dimension. We will refer to them as ”low-channel” convolutions. Depth-wise convolutions pose a similar problem to VTA, as they also lack the necessary number of input channels. To this set of workloads we added two dilated convolutions.

As explained in section 5, the $ic$ dimension is split with a factor the size of dimension $z$ in the instruction. If $ic<z$, padding $ic$ with $ic-z$ additional elements is necessary to generate code. However, padding results in lower utilization and larger tensors. Dilated convolutions are processed as by packing the original image and process it as independet dense covolution before interleaving the data back into the original shape. As a reference, we use a stencil that is blown up with zeros to match the dilated access pattern. For the evaluation, we selected the five best implementations based on the selection metric in section 4 as candidates for auto tuning.