Bandwidth-Optimal Random Shuffling for GPUs

Shuffling is a fundamental computer science problem, the objective of which is to rearrange a set of input elements into some pseudo-random order. The classical method of Fisher-Yates (Fisher and Yates, 1943) (popularised by Knuth in (Knuth, 1997)) randomly removes elements from an input buffer, one at a time, appending them to the output buffer. This algorithm runs optimally in $O(n)$ time, outputs uniformly distributed permutations, has a simple in-place variant, and is straightforward to implement. As such, sequential shuffling has long been considered a solved problem. However, modern computing hardware such as GPUs offer massively parallel computation that cannot be effectively exploited by the standard sequential shuffling approach due to dependencies inherent in Fisher-Yates type algorithms. The work presented in this paper was motivated by a gap in GPU-oriented parallel primitive libraries, such as Thrust (Bell and Hoberock, 2012), Cub (Merrill, 2015), and Boost.Compute (Szuppe, 2016), which aim to implement the C++ standard library in equivalent form for GPUs, that prior to the method presented in this paper all lacked shuffling algorithms. The primary contributions of this paper are:

• •

  ‘Bijective shuffle’: A new shuffling algorithm, carefully optimised for GPUs.

• •

  A novel statistical test for shuffling algorithms, based on kernel space embeddings.

Some applications motivating this work are: permutation Monte Carlo tests (Good, 2006), bootstrap sampling (Mitchell and Frank, 2017), approximation of Shapley values for feature attribution (Mitchell et al., 2021), differentially private variational inference (Prediger et al., 2021), and shuffle operators in GPU accelerated dataframe libraries (Team, 2018). More generally, whenever shuffling forms part of a high-performance GPU-based algorithm, performance is lost when data is copied back to the host system, shuffled using the CPU, and copied back to the GPU. Host-to-device copies are limited by PCIE bandwidth, which is typically an order of magnitude or more smaller than available device bandwidth (Farber, 2011).

Informally, an ideal parallel algorithm would be capable of assigning elements of the input sequence randomly to unique output locations without contention, communicate minimally between processors, evenly distribute work between processors, use minimal working space, and have deterministic run-time. To achieve this goal, we utilise pseudo-random bijective functions, defining mappings between permutations. These bijective functions allow independent processors to write permuted elements to an output buffer, in parallel, without collision. Critically, we allow bijective functions defining permutations of larger sets to be applied to smaller shuffling inputs via a parallel compaction operation that preserves the pseudo-random property. The result is an $O(n)$ work algorithm for generating uniformly random permutations. We perform thorough experiments to validate the empirical performance and correctness of our algorithm. To this end, we also develop a novel test on distributions of permutations via unique kernel space embeddings.

We begin in Section 2 by introducing notation, summarising existing parallel shuffling algorithms, and describing GPUs. Section 3 introduces the idea of bijective functions, explaining their connection to shuffling and proposing two candidate bijective functions. Section 4 explains how compaction operations can be combined with bijective functions to generate bijections for arbitrary length sequences, forming the ‘bijective shuffle’ algorithm. Section 5 describes a statistical test comparing the kernel space embedding of a shuffling algorithm’s output with its expected value. This test is used to evaluate the quality of the proposed shuffling algorithms and to select parameters. Finally, in Section 6, we evaluate the runtime and throughput of the proposed GPU shuffling algorithm.

Before discussing existing work, we briefly describe some notation. The order of elements in the shuffled version of an input sequence can be defined as a permutation. We refer to the symmetric group of permutations of $n$ elements as $\mathfrak{S}_{n}$. The permutation $\sigma\in\mathfrak{S}_{n}$ assigns rank $j$ to element $i$ by $\sigma(i)=j$. For example, given the permutation written in one-line notation:

and the list of items

the items are reordered such that $x_{i}$ occupies the $\sigma(i)$ coordinate

Note that we are using non-traditional zero-based indexing of permutations for notational convenience: this simplifies description of the functions used for our shuffling approach later in this paper.

The product of permutations $\sigma\tau$ refers to the composition operator, where the element $i$ is assigned rank $\sigma(\tau(i))=j$.

Shuffling is a reordering of the elements of an array of length $n$ by a random permutation on the finite symmetric group, $\sigma\in\mathfrak{S}_{n}$, such that each possible reordering is equally likely, i.e., $p(\sigma)=\frac{1}{n!},\forall\sigma\in\mathfrak{S}_{n}$.

Our proposed approach is based on applying bijective functions. The symmetric group $\mathfrak{S}_{n}$, defined over any set of $n$ distinct items, contains all bijections of the set onto itself, or all distinct permutations of $n$ elements. A bijection is a one-to-one map between the elements of two sets, where each element from the first set is uniquely paired with an element from the second set. A bijection between two sets $X,Y$ with $|X|,|Y|=n$ is characterised by some function $f_{n}:X\rightarrow Y$. The bijective function $f_{n}$ admits a corresponding inverse $f_{n}^{-1}:Y\rightarrow X$, which is also a bijection. For any two such functions $(f_{n},g_{n})$, their composition $f_{n}\circ g_{n}$ is also a bijection. As we are dealing with bijections from the symmetric group $\mathfrak{S}_{n}$, we define functions of the form $f_{n}:X\rightarrow X$. Without loss of generality, consider the set of nonnegative integers $X=\{0,1,\cdots,n-1\}$. An example of a bijection $\sigma\in\mathfrak{S}_{4}$ is

If each of the $n$ functions can be executed on a processor $p_{i}$, independently of any other processor, without communication and in reasonable time, then this defines a simple parallel algorithm for reordering elements. If $f_{n}$ furthermore exhibit suitable pseudo-random properties such that $p(\sigma)=\frac{1}{n!},\forall\sigma\in\mathfrak{S}_{n}$ (or close enough for practical application), then we have an effective parallel shuffling algorithm. We now discuss potential candidate bijective functions.

The key observation for generating bijections of arbitrary length is this: given an input vector $X=[0,1,2,\cdots,m-1]$ of length $m$, there may not be a readily available bijective function $f_{m}$ of the same length, however, we can find the nearest applicable $f_{n}$ such that $m\leq n$ and apply the bijective function to the padded vector $\hat{X}=[0,1,2,\cdots,n-1]$. The output is a vector $W$ of length $n$, containing the randomly permuted input indices. By ‘deleting’ all $w\geq m$ from $W$, we obtain a permutation $Y$ of length $m$. Using this process, we form an algorithm generating longer permutations of length $n\geq m$, then compacting them to form permutations of length $m$. This compaction of a longer random permutation is still an unbiased random permutation as per the following proposition.

To achieve the compaction in parallel, we use the work-efficient exclusive scan of Blelloch (Blelloch, 1990). Flagging each element of the extended permutation $W$ with $0$ if $w_{i}\geq m$ or $1$ if $w_{i}<m$, the output of the exclusive scan over these flags gives the scatter indices into the shuffled output vector $Y$. This process is shown in Figure 6 for $m=5,n=8$, and pseudocode is given in Algorithm 1. Given $p=n$ independent threads, the exclusive scan operation has work complexity $O(n)$. Our bijective function, executed in parallel for each element, also has work complexity $O(n)$, and so the final algorithm has optimal work complexity $O(n)$. Note that, using the bijective functions described in Section 3, the padded array of length $n$ will have at most $2m$ elements (the nearest power of two length), so the work complexity holds regardless of input length.

Algorithm 1 can be implemented easily in several steps using a GPU parallel primitives library, for example, evaluating the bijective function, gathering elements to an output buffer, then applying a stream compaction algorithm. Achieving “bandwidth optimality”, i.e., a single global memory read and write per element, is more challenging, requiring the fusion of all operations into a single GPU kernel. The standard GPU parallel prefix sum of (Harris et al., 2007) uses two passes through global memory, the first to perform block-level scans and the second to propagate partial block-level results globally. This can be improved upon by using the specialised single-pass scan implementation of (Merrill and Garland, 2016), using a technique termed ‘decoupled look-back’ to achieve non-blocking communication between thread blocks, achieving optimal $n$ global memory reads/writes. In Section 6, we evaluate three versions of our GPU shuffling algorithm, incorporating varying levels of kernel fusion and demonstrating the effectiveness of optimising global memory read/writes for the shuffling problem.

We now consider statistical tests for pseudo-random permutations to assess the suitability of the methods discussed in Section 3 and 4 for practical applications, and to select the number of rounds for the VariablePhilox cipher. Testing the distribution of permutations is challenging as the sampling space grows super-exponentially with the length of the permutation. In particular, $21!>2^{64}$, so any computer algorithm for shuffling that uses a 64 bit integer seed cannot generate all permutations for $n\geq 21$.

Our first test considers the distribution of random permutations at $n=5$, where $5!=120$, using a standard $\chi^{2}$ test with 119 degrees of freedom, under the null hypothesis that permutations are uniformly distributed, i.e., each permutation occurs with probability $p=\frac{1}{120}$. Figure 10 shows the change in the $\chi^{2}$ statistic as the number of rounds in the VariablePhilox cipher is increased. Each data point is computed from 100,000 random permutations. We also plot the acceptance thresholds for $\alpha=0.05$, corresponding to the probability of observing a value of the $\chi^{2}$ statistic above the line indicated on the figure, if the null hypothesis is true. Considering the results shown in the figure, it is unlikely that samples from VariablePhilox with less than 20 rounds are drawn from a uniform distribution. This is somewhat surprising given that only 10 rounds are recommended for the original Philox cipher in the PRNG setting (Salmon et al., 2011). The LCG bijective function generates a $\chi^{2}$ statistic in the region of 500,000, so it clearly fails the test and is not included in this figure.

The $\chi^{2}$ test is useful for small $n$, but is intractable otherwise. We develop another test statistic suited to larger $n$, based on the maximum mean discrepancy (MMD) in a reproducing kernel Hilbert space (RKHS). Two-sample hypothesis tests using MMD are developed in (Gretton et al., 2012). In a similar spirit, we derive a one sample test comparing the uniform distribution of permutations against a finite sample generated by a shuffling algorithm. The $\mathrm{MMD}^{2}$ between two distributions $p(X)$ and $q(Y)$ in a RKHS $\mathcal{H}$, equipped with positive definite kernel $K$, is defined as

If the kernel $K$ is a characteristic kernel, the mean embedding of a distribution induced by the kernel is injective (Fukumizu et al., 2004). In other words, the mean embedding of a distribution is unique to that distribution. As a consequence, $\mathrm{MMD}(p,q)=0$ if and only if $p=q$. Thus, a strategy for statistical testing is to form the null hypothesis that $p=q$, compute the sample estimate $\hat{\mathrm{MMD}^{2}}(p,q)$ as the test statistic, and evaluate the probability of obtaining a sample estimate greater than some threshold, assuming $p=q$, using a concentration inequality. If the observed test statistic is sufficiently unlikely, this provides evidence that $p\neq q$.

To implement this idea, a characteristic kernel measuring the similarity of two permutations is needed. The Mallows kernel, for $\lambda\geq 0$, is defined for permutations as

where

We (somewhat arbitrarily) use the parameter $\lambda=5$ throughout this paper. The Mallows kernel is introduced in (Jiao and Vert, 2015) and shown to be characteristic in (Mania et al., 2018). It may be implemented in time $O(n\log n)$ using the procedure presented in (Knight, 1966). In the following, we make use of the fact (Mitchell et al., 2021) that the expected value of the Mallows kernel under a uniform distribution of permutations is

The Mallows kernel is right-invariant in the sense that $K(\sigma,\sigma^{\prime})=K(\tau\sigma,\tau\sigma^{\prime})$ for any $\tau\in\mathfrak{S}_{d}$ (Diaconis, 1988). Also note that the uniform distribution of permutations is invariant to composition. Using right invariance in conjunction with $\tau=\sigma^{-1}$, and assuming $p$ denotes the uniform distribution over permutations, then

with $I$ the identity permutation. Using the above, and assuming that $p$ is a uniform distribution of permutations, the $\mathrm{MMD}^{2}$ from (4) is simplified as

Replacing $q$ with a sample of random permutations $\Pi$, of size $|\Pi|$ and with $|\Pi|\equiv 0\pmod{2}$, we obtain an unbiased sample estimate via

where $\mathbb{E}[\hat{\mathrm{MMD}}^{2}(p,\Pi)]=0$ if and only if $p=q$. We derive two tests based on the $\hat{\mathrm{MMD}}^{2}(p,\Pi)$ statistic. The first is a distribution-free bound. Applying Hoeffding’s inequality (Hoeffding, 1963), with the fact that $0\leq K(\sigma,\sigma^{\prime})\leq 1$, we have

Let $\alpha_{H}$ be the level of significance under the Hoeffding bound. With null hypothesis $p=q$, the acceptance region of the test is

The second test uses the central limit theorem, implying that $\hat{\mathrm{MMD}}^{2}(p,\Pi)$ approaches a normal distribution as $|\Pi|\rightarrow\infty$. The variance of the Mallows kernel may be computed directly from its expectation as

So we have

The factor of two arises because the sum in (5) is of size $|\Pi|/2$. According to a normal distribution with mean 0 and variance as above,

where $\mathrm{erf}$ is the canonical error function. The acceptance region for $\alpha_{N}$ is

The threshold for $\alpha_{H}$ should be used for small $|\Pi|$ (i.e. $<100$), as it makes no assumptions on the distribution of $\hat{\mathrm{MMD}}^{2}(p,\Pi)$. For larger sample sizes, the asymptotic $\alpha_{N}$ threshold is significantly tighter and should be preferred. Comparing Figures 10 and 10, we see the MMD tests with the asymptotic acceptance threshold roughly coincide with the $\chi^{2}$ test at n = 5, where some tests fail for VariablePhilox with fewer than 20 rounds, but VariablePhilox with 24 rounds or more passes all tests. Figures 10, 10, and 10 plot the $|\hat{\mathrm{MMD}}^{2}(p,\Pi)|$ statistic for VariablePhilox, LCG, and std::shuffle using $|\Pi|=$100,000, for varying permutation lengths. These experiments lead us to recommend a 24 round VariablePhilox cipher for random permutation generation, and we use this configuration in all subsequent experiments.

We now evaluate the throughput of our bijective shuffling method, where throughput refers to (millions of keys)/time(s). Unless otherwise stated, experiments are performed on an array of 64-bit keys of length $2^{w}+1$, where $w$ ranges from 8 to 29. This represents the worst-case scenario, where our bijective shuffle algorithm must redundantly evaluate $2^{w}-1$ elements. For all throughput results, we report the average of five trials.

To consider the effect of code optimisations, we evaluate three CUDA implementations of Algorithm 1 with varying levels of GPU kernel fusion:

• •

  Bijective0: The transformation ($b=f_{n}(i)$), stream compaction, and gather are implemented in separate passes.

• •

  Bijective1: The transformation, stream compaction, and gather are fused into a single scan operation. The two-pass scan algorithm of (Harris et al., 2007) is used.

• •

  Bijective2: The transformation, stream compaction, and gather are fused into a single scan operation. The single pass scan algorithm of (Merrill and Garland, 2016) is used, for $m$ total global memory reads/writes.

Figure 11 plots the throughput of these variants, on a Tesla V100-32GB GPU, using VariablePhilox as the bijective function. For reference, the line labeled “gather” shows an upper bound on throughput for $n$ random gather operations in global memory. The optimised algorithm Bijective2 closely matches the optimal throughput of random gather, and is equivalent in performance for sizes $>2^{21}$. Bijective2(n=m) indicates the best-case performance of the algorithm, where the sequence is exactly a power of two length. The best-case and worst-case performance do not differ greatly because redundant elements do not incur global memory transactions, and are therefore relatively inexpensive to evaluate. The performance of random gather peaks around $2^{20}$, where there is sufficient L2 cache to mitigate the effects of uncoalesced reads/writes. At this size, there is a gap between Bijective2 and random gather due to arithmetic operations required in evaluating multiple rounds of the VariablePhilox function, and because the prefix-sum must be redundantly evaluated up to the nearest power of two (although no global memory transactions occur for these redundant elements). This gap disappears at larger sizes when the runtime of the kernel becomes dominated by memory operations.

We now proceed to a comparison of different shuffling algorithms. Table 1 and Figure 13 show the throughput in millions of items per second, for the Tesla V100-32GB GPU, and Table 2 and Figure 13 show data for the GeForce RTX 2080 GPU. LCG and VariablePhilox implement optimised bijective shuffling using the methods described in Section 3. The dart-throwing algorithm uses CUDA atomic exchange instructions to attempt to place items randomly in a buffer of size $2n$, a stream compaction operation is then applied to the buffer to provide the final shuffled output. SortShuffle applies the state-of-the-art radix sort algorithm of (Merrill, 2015) to randomly generated 64-bit keys. As discussed in Section 2.1, many existing work on shuffling reduce to sorting algorithms on infinite length keys. Thus, SortShuffle is representative of a wide class of divide-and-conquer algorithms.

The bijective shuffle algorithms with the VariablePhilox and LCG functions achieve near-optimal throughput at large sizes, where performance is dominated by global memory operations. Throughput for these methods is more than an order of magnitude higher than the DartThrowing or SortShuffle methods. The lesser throughput of DartThrowing can be explained by the overhead of generalised atomic instructions in the CUDA architecture, as well as contention among threads when multiple threads attempt to write to the same memory location. SortShuffle relies on radix sort, where prefix sum is applied to 4 bits in each pass, requiring 16 passes over the data to fully sort 64-bit keys. In comparison, bijective shuffle is designed to require only a single prefix sum operation. Bijective shuffle is also fully deterministic, unlike DartThrowing, and requires no global memory for working space, whereas DartThrowing and SortShuffle use memory proportional to $O(n)$.

We also evaluate the performance of our shuffling method using 2x Intel Xeon E5-2698 CPUs, with a total of 40 physical cores. Results are presented in Table 3 and Figure 14. The bijective shuffle method is not designed for CPU architectures, but the results are informative nonetheless. The Gather, VariablePhilox (bijective shuffle) and SortShuffle algorithms are implemented using the Intel Thread Building Blocks (TBB) library (Reinders, 2007). The Rao-Sandelius (RS) (Rao, 1961) and MergeShuffle (Bacher et al., 2015) algorithms use the implementation of (Bacher et al., 2015), and std::shuffle is the single-threaded C++ standard library implementation. VariablePhilox has the highest throughput for medium-sized inputs, although it is outperformed by RS when $n>2^{24}$. Curiously, RS has higher throughput at large sizes than the TBB gather implementation. This could indicate that RS has a more cache-friendly implementation, despite having a worse computational complexity of $O(n\log n)$, or simply a suboptimal implementation in TBB. Further work is needed to properly investigate these effects for CPU architectures.

While the presented algorithm for GPUs is highly effective in terms of throughput and the quality of its pseudorandom outputs, its parameterisation differs from that of typical standard library implementations of shuffle. For example, the C++ standard library implementation accepts any ‘uniform random bit generator’ (Stroustrup, 2013), such as the common Mersenne-twister generator (Matsumoto and Nishimura, 1998). Our implementation is instead parameterised by the selection of a bijective function, and a key of length $\log(n)k$ bits, where $n$ is the length of the sequence to be shuffled, and $k$ is the number of rounds in the cipher.

Future work may explore other useful constructions of bijective functions, beyond the LCG and Feistel variants discussed in this paper.

It also should be noted that the Fisher-Yates and Rao-Sandelius algorithms may operate in-place, and our proposed GPU algorithm does not, requiring the allocation of an output buffer. This is not unexpected, as there are few truly in-place GPU algorithms that reorder inputs ((Peters et al., 2009) is a notable exception). For future work, it would be interesting to consider if an in-place parallel shuffling algorithm is possible for GPU architectures.

We provide an algorithm for random shuffling specifically optimised for GPU architectures, using modified bijective functions from cryptography. Our algorithm is highly practical, achieving performance approaching the maximum possible device throughput, while being fully deterministic, using no extra working space, and providing high-quality distributions of random permutations. An outcome of this work is also a statistical test for uniform distributions of permutations based on the Mallows kernel that we expect to be useful beyond shuffling algorithms.