How should Futhark expose irregular arrays to the programmer?
Recently Futhark has grown some powerful new features such as flattening nonuniform data parallelism and recursive functions. While the latter is obviously user-visible, and we’ll see some nice uses later in this post, flattening is mainly a program transformation. Briefly, it lets the compiler handle expressions such as
map (\n -> i64.sum (iota n)) nswhere the size of the inner parallelism (n) differs between iterations of the
outer map, and fully exploit all parallelism. Adding this did not
fundamentally change which Futhark programs could be written, or their
semantics, it merely had an effect on how good GPU code we could generate (and
in extreme cases, whether we could generate GPU code at all). In this post we
will see how the language can be extended to be fundamentally more expressive
via a somewhat small extension, and what kinds of programs we can now write in
Futhark.
Irregular arrays
The main language restriction in Futhark is that irregular arrays are not
allowed. An irregular array is a multidimensional array where subarrays differ
in size. For example, `[[1], [3,3]] is irregular, while [[1,2],[3,4]] is
regular. The reason for this restriction is ultimately rooted in concerns for
the efficiency of compilation, but the user experience is that these arrays are
simply not well-typed in Futhark’s size type
system. However, the main challenge in
implementing flattening is handling the irregular arrays that arise as
intermediate results when performing loop
distribution:
map (\n -> i64.sum (iota n)) ns
↓
let tmps = map (\n -> iota n) ns
let res = map (\tmp -> i64.sum tmp) tmpsThis means that the flattening transformation is actually perfectly able to handle irregular arrays by encoding them as regular arrays, without the rest of the compiler having to know about it. Irregular arrays are definitely useful in some cases, so the question is how we can expose some of this power in the source language without adding general support for irregular arrays. As an example of the kind of code we would like to write, here is a recursive data-parallel quicksort written in NESL:
function quicksort(a) =
if (#a < 2) then a
else
let pivot = a[#a/2];
lesser = {e in a| e < pivot};
equal = {e in a| e == pivot};
greater = {e in a| e > pivot};
result = {quicksort(v): v in [lesser,greater]};
in result[0] ++ equal ++ result[1];
Note the main trick: mapping over a two-element array that performs the
recursive quicksort calls, although in NESL this is done with an array
comprehension rather than a map function. This array is irregular in NESL,
which of course won’t fly in Futhark.
This will serve as our motivating example for a start, although we’ll see that quicksort is actually one of the simpler cases.
An initial flatmap SOAC
Futhark allows parallel programming via second-order array combinators
(SOACs): functions such as map, reduce, and scan. We extend this with a
new SOAC flatmap, that semantically is like map, except the results returned
by the mapped function may differ in size, and flatmap returns the
concatenation of the arrays:
val flatmap [n] 'a 'b : (f: (k: i64) -> (x: a) -> [k]b)
-> [n]i64
-> [n]a
-> ?[m].[m]bNote the careful use of size types: flatmap accepts two arrays both of size
n, and the first array (of type
[n]i64) contains the expected sizes
returned by the f function on the corresponding array of as. The size k is
also passed to the mapped function so we can use it in its return type. The
result of flatmap has existential size m, because we cannot know it in
advance. (Actually, we can: m is the sum of the [n]i64 array.)
This is similar to Haskell’s concatMap, although the intention is that
flatmap exploits parallelism across the n input elements and whatever
parallelism may be in f. However, semantically flatmap is nothing special,
and can easily be implemented with a sequential loop. Even the parallel
implementation of flatmap (both this version and the ones below) is largely
straightforward, as it is just a call to the “lifted” form of the provided
function, so in this post we focus on whether flatmap provides the flexibility
we need in order to express the algorithms we care about.
This flatmap is sufficient to implement a recursive quicksort in Futhark:
def quicksort [n] (xs: [n]i32) : [n]i32 =
if n <= 1
then xs
else let pivot = xs[0]
let [m][k] (lesser: [m]i32, greater: [k]i32) =
partition (<= pivot) (drop 1 xs)
let sorted =
flatmap (\k p -> quicksort (sized k (if p then lesser else greater)))
[m, k]
[true, false]
in sized n (take m sorted ++ [pivot] ++ drop m sorted)The code looks slightly odder than a normal recursive quicksort. The oddity is
constructing a two-element array just to flatmap it immediately rather than
performing two recursive calls, but this is necessary in order to run the two
recursive calls in parallel. In NESL this array is irregular, while in Futhark
we construct an array of booleans ([true, false]) and use a conditional inside
the flatmap function to pick either the lesser or greater array. I expect
this may become a common trick when writing divide-and-conquer programs in
Futhark.
Generalising flatmap
Although sufficient to express quicksort, the flatmap above has the
restriction that we must be able to determine the size of each result in
advance. This is sufficient for sorting, because sorting an array does not
change its size, but is not practical for algorithms such as
Quickhull, where the size of the
result is usually smaller than the size of the input. The obvious extension is
to modify flatmap such that the size of the array returned by the mapped
function is existentially quantified:
val flatmap [n] 'a 'b : (f: (x: a) -> ?[k].[k]b)
-> [n]a
-> ?[m].[m]bNow k is no longer a parameter to the f function, but bound in its return
type. This means no longer know anything about the size of the result, which
quickly turns out to be impractical. We therefore extend flatmap to also
return the size of each array returned by f, in the form of a shape vector:
val flatmap [n] 'a 'b : (f: (x: a) -> ?[k].[k]b)
-> [n]a
-> ?[m].([n]i64, [m]b)In fact, the real flatmap in the Futhark
prelude
returns not just the shape vector, but also various other related metadata that
flattening computes anyway, but for simplicity we’ll stick with just the shape
vector for this post.
Using this definition, we can implement Quickhull. I will not give the full program as it involves many functions for geometry calculations, but the recursive part can be written like this:
def hull [n] (a: point) (b: point) (pts: [n]point) : []point =
if n <= 1
then pts
else let p = farthest a b pts
let f i =
let (x, y, pts') =
if i == 0
then (a, p, filter (\q -> side a p q > 0) pts)
else
then (p, b, filter (\q -> side p b q > 0) pts)
in hull x y pts'
let (shape, pts') = flatmap f [0, 1]
in take shape[0] pts' ++ [p] ++ take shape[1] pts'We can actually write this in a slightly cleverer way that does not use shape
at all - see the link if you are interested.
Generalising flatmap further
As a final generalisation, although one that we don’t yet have a concrete use
case for, we also allow flatmap to return a uniform result for each input
element, like a normal map. This gives us this final type:
val flatmap [n] 'a 'b 'c : (f: (x: a) -> ?[k].([k]b, c))
-> [n]a
-> ?[m].([n]i64, [m]b, [n]c)Whenever we are not interested in a uniform result, we can simply instantiate
c with (), and the prelude contains a wrapper named flatmap' that does
this.
Although size types are by now a rather mature feature in Futhark, I do still
get a pleasant surprise whenever they allow us to express the type of a function
like flatmap in such a clear and precise manner.
With this generalisation, flatmap has now reached its final form. I like that
it is such a small extension, barely even a “language extension” at all
(although the compiler IR had to be modified), and that its semantics are
completely trivial. The question is to what extent it is sufficient to express,
say, NESL’s library of parallel
algorithms in a clean
manner, despite the language itself not supporting irregular arrays.
Performance
There’s a final wrinkle to the story: how fast is flatmap? In the flattening
post I mentioned that performance of the
nonuniform case had not yet been prioritised, and in the recursion
post I mentioned that it seemed like recursion in a
data parallel language could easily be very costly. We are now combining two
inefficient things, so it is likely the spirit of the
hedgehog may be somewhat lacking. But just how bad
is it? We are in the fortunate situation that flatmap does not fundamentally
increase the expressivity of the language, as it was always possible (if tedious
and error-prone) to flatten by hand, so we do have hand-flattened versions of
the algorithms we can now implement with recursive divide-and-conquer. For
example, let us compare the quicksort from the
https://github.com/diku-dk/sorts package with the recursive quicksort
above:
-- ==
-- entry: bench_flatmap bench_manual
-- random input { [100000]i32 }
entry bench_flatmap = quicksort -- from above
module M = import "lib/github.com/diku-dk/sorts/quick_sort"
entry bench_manual = M.qsort (i32.<=)Be aware that the qsort from the package is known to be fairly inefficient (we
usually recommend merge- or radix-sort on GPUs), so none of the sorts here are
even remotely fast in absolute terms. But let’s see how bad it is:
$ futhark bench --backend=cuda sortbench.fut
Compiling sortbench.fut...
Reporting arithmetic mean runtime of at least 10 runs for each dataset (min 0.5s).
More runs automatically performed for up to 300s to ensure accurate measurement.
sortbench.fut:bench_flatmap (no tuning file):
[100000]i32: 231986μs (95% CI: [ 230655.0, 233282.1])
sortbench.fut:bench_manual (no tuning file):
[100000]i32: 12079μs (95% CI: [ 12059.0, 12098.0])
Pretty awful stuff. I have not investigated just why bench_flatmap is so slow,
but I suspect it is because of flattening producing a huge number of small
kernels, as well as the intrinsic overhead of recursive calls, where each
activation record maintains many large-ish arrays of flattening metadata and
values. The hand-flattened quicksort (bench_manual) is very slow as sorts go
because of the excessive bookkeeping, but it is written with a tail-recursive
loop, and a minimal number of parallel operations.
It raises an interesting research question: can we automatically compile such recursive divide-and-conquer functions to something that is equivalent to manually flattened tail-recursive code? Yes, I remember writing about why I dislike tail call optimisation, but in this case it seems necessary to obtain good performance. This, however, is for the future, and if the program analysis becomes too tricky or fragile, we may need a dedicated SOAC for expressing various divide-and-conquer patterns in a way the compiler can exploit.