Speeding up GPU copies by permuting iteration spaces

Following the recent posts on such riveting topics as integer division and addition, we turn our attention to another exciting topic: copying bytes. Most compilers end up emitting code that copies data around, and in most cases this means very large amounts of very small copies, such that any attempt at run-time cleverness is likely counterproductive. However, array languages are usually in the situation of dealing with a relatively small number of operations, that on the other hand are large, which means some interesting optimisations are possible. In this post I will explain how Futhark now makes certain copies a little faster, by exploiting information about array layouts that are only available at run-time. This has a small impact on some programs, but the technique is probably more fun than useful.

I previously wrote a detailed post on how Futhark represents array layouts, but I will repeat the salient parts in the following.

The fundamental question is how to map the n-dimensional index space of an n-dimensional array to the flat index space of computer memory. We need an index function from a tuple of n indexes to a single integer, representing a relative offset in some memory block. These index functions can in principle be essentially arbitrarily complicated, but because we want to both analyse them statically and represent them dynamically, Futhark uses a structured (and limited) representation known as linear memory access descriptors (LMADs). Simplifying slightly, an LMAD for an n-dimensional index space consists of n pairs of integers, one per dimension, with each pair denoting a size and a stride. The sizes describe the size of each dimension, and we will return to the strides in a moment. In total, an LMAD L for a d-dimensional space is written as follows:

L = [(n_1, s_1), ..., (n_d, s_d)]

A d-dimensional LMAD can be applied to d integers as follows:

L(i_1, ..., i_d) = i_1s_1 + ... + i_ds_d

That is, each index is multiplied with the corresponding stride, and the whole thing summed. If we have a two-dimensional array A represented by the LMAD L, then to index A[i,j] we compute an offset with L(i,j) and use that as an address relative to whatever memory A is stored in (multiplied with an element size, of course).

Here is an example of an LMAD describing an array of shape [4][5] in row-major order:

[(4, 5), (5, 1)]

And here it is in column-major order:

[(4,1), (5, 4)]

We can see that when we go from index [i][j] to [i+1][j], then the address increases by 5 when using row-major, and by 1 when using column-major. This is as expected. LMADs thus provide a convenient way to decouple the logical shape of an array from its physical element order. We can also do things like slicing, reshaping, and transposition in a copy-free manner, simply by fiddling with the LMADs instead. (Slicing also requires a single additional offset in addition to the dimensions, but that is immaterial for this post - read the one linked above if you are curious.)

We don’t need the LMAD sizes for the index operation, but they are useful when operating on multiple elements. In particular, they describe the size of the array. If we have two arrays A and B of the same shape, stored in memory AM and BM and with LMADs AL and BL, then we can copy A to B using a loop like this:

for i_1 < n_1: ... for i_d < n_d: BM[BL(i_1, ..., i_d)] = AM[AL(i_1, ..., i_d)]

If we assume that LMADs are injective (each index maps to different offset, which in particular bans zero strides), then the above loop has no dependencies, i.e., all iterations can be executed in parallel. The Futhark compiler only generates such LMADs, so we will assume that henceforth.

LMADs cannot represent all possible layouts, but they can represent any permutation of dimensions, which is the most important case for our purposes. They are used in the compiler itself, where the components are (usually) symbolic rather than known integers, but LMADs also exist at run-time, where they are used to dynamically optimise array copies, instead of always doing the naive loop described above.

Consider again that we have two arrays A/B with LMADs AL/BL in memory blocks AM/BM, and we have to copy A to B. We can certainly do it with the manual loop, but odds are that a platform-provided memcpy() (or GPU equivalent) is faster because it is written to take advantage of hardware-level features, like vectorisation. It would be good if instead of doing a manual loop, we could just emit memcpy(B, A, m), with m being the total size in bytes. This is possible exactly when A and B have the same layout, and the LMADs describe contiguous arrays, in the sense that there are no “gaps” in memory between elements.

Detecting such contiguous arrays is possible with run-time analysis. Specifically, two neighbouring dimensions (n_1,s_1),(n_2,s_2) in an LMAD can be collapsed into a single dimension when s_1=s_2n_2, and the result is a dimension (n_1n_2,s_2). If there exist any permutation of the LMAD dimensions such that we can collapse it into a single dimension, then the LMAD is contiguous (we can also ignore singleton dimensions). This can be done with a worst-case exponential search. Doing such a search on every copy might be a bad idea for most languages, but note that the search is exponential in the rank of the array, and usually arrays in Futhark programs have very few dimensions. Importantly, the rank is a constant no matter the size of the array. It is not a big deal to burn a few thousand instructions before copying an array with thousands (or billions) of elements.

Another important special case is to detect transpositions, which occurs when we are copying between different layouts in a way that corresponds to a transposition of a matrix. These are important cases because a naive transposition has poor locality on all hardware I know of: either the reads or writes will be accessing with a stride. However, with careful programming that exploits the memory hierarchy carefully, it is possible to implement transpositions that are almost as fast as a full-speed memory copy. (Details are out of scope for this post.)

Transpositions can be detected by collapsing contiguous dimensions until only two dimensions are left in each LMAD, and then checking whether they are transposes of each other. In fact, Futhark is a little more clever and allows batched transpositions (or map transpose as I prefer to call it), which allows an additional outer dimension that is unchanged, and additionally does not require the “matrices” to be contiguous with each other.

Transposes are quite important, because Futhark inserts them in the program to improve the locality of array traversals based on what the program actually does. This would obviously not work if those transpositions were not implemented efficiently. The capability to optimise data representation is one of my favourite consequences of using high level programming languages.

The curious can see the implementation of Futhark’s host memory copy function (macro-ized for different element types), which covers the cases discussed above.

Perhaps there are more interesting structures we will detect in the future, but at some point, we must face the possibility that the LMADs involved are just too complicated to make any sense of them, and fall back to copying elementwise with a loop. This is not the end of the world, as at least that loop is parallel, but I realised that we can (heuristically) do slightly better. Consider again the general loop for copying between arrays described by LMADs:

for i_1 < n_1: ... for i_d < n_d: BM[BL(i_1, ..., i_d)] = AM[AL(i_1, ..., i_d)]

Suppose we want to run this on a GPU. Since GPUs do not support nested parallelism (hence all this), we transform the above into a single loop:

for i < n_1 * ... n_d: (i_1, ..., i_d) = unpack_flat_index(i) BM[BL(i_1, ..., i_d)] = AM[AL(i_1, ..., i_d)]

That unpack_flat_index function is the usual mess of divisions, but it is not our concern here. Let us suppose that this flat loop is however a straightforward flattening of the original nested loop (same lexicographic order), and suppose we assign one GPU thread to each iteration of the loop. We have not fundamentally changed the structure of the loop by adding new dependencies, so we can still change the iteration order in any way we wish, in order to perform locality optimisations. (Note that this follows from the injectivity of the destination LMAD, but does not hold for arbitrary LMADs.)

A GPU executes threads by bundling consecutive threads into groups called warps or wavefronts of typically 32 or 64 threads, which execute in lockstep (fine print applies, it doesn’t matter for this explanation). Intuitively, the threads of a warp execute the same instruction at the same time. The hardware services these accesses using one or more memory transactions, each covering a contiguous block of memory. If the different threads access nearby locations, fewer transactions are needed. We call this coalesced memory access. Otherwise, multiple transactions are necessary before all threads have been serviced. We call this uncoalesced memory access. Ensuring coalesced access is one of the most fundamental performance requirements when doing GPU programming. The rule of thumb is as follows: neighbouring threads should access neighbouring elements of arrays.

In Futhark’s GPU kernel for general LMAD copying, the index space is flattened “column-major” in a sense, such that neighbouring threads process neighbouring outermost indexes. That is, when copying a two-dimensional array, two threads will be copying elements [i,j] and [i+1,j]. Note that this is still the logical iteration order, whose mapping to the physical iteration order is given by the LMADs of the source and destination arrays.

This inspires an optimisation where before copying, we permute the dimensions of the LMADs such that they are sorted in order of increasing stride by the source LMAD (and secondarily destination LMAD). The goal is to minimise the stride between neighbouring threads. Ideally we hope for perfect coalescing, but even strides greater than 1 can be useful. We prioritise locality for the source array because uncoalesced reads seem to be more costly than uncoalesced writes on the GPUs I have tested. We sort with a basic selection sort, which has quadratic complexity, which is again a small price to pay if we can improve locality.