Branchless Lomuto Partitioning

A partition function accepts as input an array of elements, and a function returning a bool (a predicate) which indicates if an element should be in the first, or second partition. Then it returns two arrays, the two partitions:

def partition(v, pred):
    first = [x for x in v if pred(x)]
    second = [x for x in v if not pred(x)]
    return first, second

This can actually be done without needing any extra memory beyond the original array. An in-place partition algorithm reorders the array such that all the elements for which pred(x) is true come before the elements for which pred(x) is false (or vice versa - it doesn’t really matter). Finally by returning how many elements satisfied pred(x) to the caller they can then logically split the array in two slices. This scheme is used in C++’s std::partition for example.

In-place partition algorithms

There are many possible variations / implementations of in-place partition algorithms, but they usually follow one of two schools: Hoare or Lomuto, named after their original inventors Tony Hoare (the inventor of quicksort), and Nico Lomuto.

Hoare

In Hoare-style partition algorithms you have two iterators scanning the array, one left-to-right ($i$) and one right-to-left ($j$). The former tries to find elements that belong on the right, and the latter tries to find elements that belong on the left. When both iterators have found an element, you swap them, and continue. When the iterators cross each other you are done.

def hoare_partition(v, pred):
    i = 0       # Loop invariant: all(pred(x) for x in v[:i])
    j = len(v)  # Loop invariant: all(not pred(x) for x in v[j:])
    while True:
        while i < j and pred(v[i]): i += 1
        j -= 1
        while i < j and not pred(v[j]): j -= 1
        if i >= j: return i
        v[i], v[j] = v[j], v[i]
        i += 1

The loop invariant can visually be seen as such (using the symbols $<$ and $\geq$ for the predicate outcomes, as is usual in sorting):

Doing this efficiently is perhaps an article for another day, if you are curious you can check out the paper BlockQuicksort: How Branch Mispredictions don’t affect Quicksort by Stefan Edelkamp and Armin Weiß, which is the technique I implemented in pdqsort. Another take on the same idea is found in bitsetsort. Key here is that it can be done branchlessly. A branch is a point at which the CPU has to make a choice of which code to run. The most recognizable form of a branch is the if statement, but there are others (e.g. while conditions, calling a function from a lookup table, short-circuiting logical operators).

To cut a long story short, CPUs try to predict which piece of code it should run next and already starts doing it even before it knows if the code it is sending down the pipeline is the right choice. This is great in most cases as most branches are easy to predict, but it does incur a penalty when the prediction was wrong, as the CPU needs to stop, go back and restart from the right point once it realizes it was wrong (which can take a while).

Especially in sorting when the outcomes of comparisons are ideally unpredictable (the more unpredictable the outcome of a comparison, the more informative getting the answer is), it can thus be advisable to avoid branching on comparisons altogether.

Lomuto

In Lomuto-style partition algorithms the following invariant is followed:

That is, there is a single iterator scanning the array from left-to-right ($j$). If the element is found to belong in the left partition, it is swapped with the first element of the right partition (tracked by $i$), otherwise it is left where it is.

def lomuto_partition(v, pred):
    i = 0  # Loop invariant: all(pred(x) for x in v[:i])
    j = 0  # Loop invariant: all(not pred(x) for x in v[i:j])
    while j < len(v):
        if pred(v[j]):
            v[i], v[j] = v[j], v[i]
            i += 1
        j += 1

    return i

This article is focused on optimizing this style of partition.

Branchless Lomuto

A few years ago I read a post by Andrei Alexandrescu which discusses a branchless variant of the Lomuto partition. Its inner loop (in C++) looks like this:

for (; read < last; ++read) {
    auto x = *read;
    auto smaller = -int(x < pivot);
    auto delta = smaller & (read - first);
    first[delta] = *first;
    read[-delta] = x;
    first -= smaller;
}

At the time I was not overly impressed, as it does a lot of arithmetic to make the loop branchless, so I disregarded it. A while back my friend Lukas Bergdoll approached me with a new partition algorithm which was doing quite well in his benchmarks, which I recognized as being a variant of Lomuto. I then found a way the algorithm could be restructured without using conditional move instructions, which made it perform better still. I will present this algorithm now.

Simple version

First, a simplified variant which will make the more optimized variant much more readily understood:

def branchless_lomuto_partition_simplified(v, pred):
    i = 0  # Loop invariant: all(pred(x) for x in v[:i])
    j = 0  # Loop invariant: all(not pred(x) for x in v[i:j])
    while j < len(v):
        v[i], v[j] = v[j], v[i]
        i += int(pred(v[i]))
        j += 1

    return i

This is actually quite similar to the original lomuto_partition, except we now always unconditionally swap, and replace the conditional increment of $i$ with an if statement by simply converting the boolean condition to an integer and adding it to $i$.

To visualize this, the state of the array looks like this after the unconditional swap:

From this it should be pretty clear that incrementing $i$ if the predicate is true (corresponding to $v[i] < p$ for sorting) and unconditionally incrementing $j$ restores our Lomuto loop invariant. The only corner case is when ${i = j},$ but then the swap is a no-op and the algorithm remains correct.

Eliminating swaps

Swaps are equivalent to three moves, but by restructuring the algorithm we can get away with two moves per iteration. The trick is to introduce a gap in the array by temporarily moving one of the elements out of the array.

def branchless_lomuto_partition(v, pred):
    if len(v) == 0: return 0

    tmp = v[0]
    i = 0  # Loop invariant: all(pred(x) for x in v[:i])
    j = 0  # Loop invariant: all(not pred(x) for x in v[i:j])
    while j < len(v) - 1:
        v[j] = v[i]
        j += 1
        v[i] = v[j]
        i += pred(v[i])

    v[j] = v[i]
    v[i] = tmp
    i += pred(v[i])
    return i

This is our new branchless Lomuto partition algorithm. Its inner loop is incredibly tight, involving only two moves, one predicate evaluation and two additions. A full visualization of one iteration of the algorithm (the striped red area indicates the gap in the array):

We can now compare the assembly of the tight inner loops of Andrei Alexandrescu’s branchless Lomuto partition and ours:

.alexandrescu:
    mov     edi, DWORD PTR [rdx]
    mov     rsi, rdx
    mov     ebp, DWORD PTR [rax]
    cmp     edi, r8d
    setb    cl
    sub     rsi, rax
    sar     rsi, 2
    test    cl, cl
    cmove   rsi, rbx
    sal     rcx, 63
    sar     rcx, 61
    mov     DWORD PTR [rax+rsi*4], ebp
    lea     r11, [0+rsi*4]
    mov     rsi, rdx
    add     rdx, 4
    sub     rsi, r11
    sub     rax, rcx
    mov     DWORD PTR [rsi], edi
    cmp     rdx, r10
    jb      .alexandrescu

.orlp:
    lea     rdi, [r8+rax*4]
    mov     ecx, DWORD PTR [rdi]
    mov     DWORD PTR [rdx], ecx
    mov     ecx, DWORD PTR [rdx+4]
    cmp     ecx, r9d
    mov     DWORD PTR [rdi], ecx
    adc     rax, 0
    add     rdx, 4
    cmp     rdx, r10
    jne     .orlp

Half the instruction count! A neat trick the compiler did is translate the addition of the boolean result of the predicate (which is just a comparison here) to adc rax, 0. It avoids needing to create a boolean 0/1 value in a register by setting the carry flag using cmp and adding zero with carry.

Conclusion

Is the new branchless Lomuto implementation worth it? For that I’ll hand you over to my friend Lukas Bergdoll who has done an extensive write-up on the performance of an optimized implementation of this partition with a variety of real-world benchmarks and metrics.

From an algorithmic standpoint the branchless Lomuto- and Hoare-style algorithms do have a key difference: they differ in the amount of writes they must perform. Branchless Lomuto-style algorithms always do at least two moves for each element, whereas Hoare-style algorithms can get away with doing a single move for each element (crumsort), or even better, half a move per element on average for random data (BlockQuicksort, pdqsort and bitsetsort, although they spend more time figuring out what to move than crumsort does - one of many trade-offs). So a key component of choosing an algorithm is the question “How expensive are my moves?” which can vary from very cheap (small integers in cache) to very expensive (large structs not in cache).

Finally, the Lomuto-style algorithms tend to be significantly smaller in both source code and generated machine code, which can be a factor for some. They are also arguably easier to understand and prove correct, Hoare-style partition algorithms are especially prone to off-by-one errors.