Multiway stable partition - cuda

Question

Is there a way to perform multiway (>2) stable partition in Thrust?
Either stable partition or stable partition copy both are equally interesting. Currently I can only use two-way stable partition copy for purposes described above. It is clear how to use it to partition a sequence into a three parts using two predicates and two calls of thrust::stable_partition_copy. But I am sure it is technically possible to implement multiway stable partition.
I can imagine the following multiway stable partition copy (pseudocode):
using F = float;
thrust::device_vector< F > trianges{N * 3};
// fill triangles here
thrust::device_vector< F > A{N}, B{N}, C{N};
auto vertices_begin = thrust::make_tuple(A.begin(), B.begin(), C.begin());
using U = unsigned int;
auto selector = [] __host__ __device__ (U i) -> U { return i % 3; };
thrust::multiway_stable_partition_copy(p, triangles.cbegin(), triangles.cend(), selector, vertices_begin);
A.begin(), B.begin(), C.begin() should be incremented individually.
Also, I can imagine hypothetical dispatch iterator, which would do the same (and would be more useful I think).

Answer 1

From my knowledge of the thrust internals, there is no readily adaptable algorithm to do what you envisage.
A simple approach would be to extend your theoretical two pass three way partition to M-1 passes using a smart binary predicate, something like
template<typename T>
struct divider
{
int pass;
__host__ __device__ divider(int p) : pass(p) { };
__host__ __device__ int classify(const T &val) { .... };
__host__ __device__ bool operator()(const T &val) { return !(classify(val) > pass); };
}
which enumerates a given input into M possible subsets and returns true if the input is in the Nth or less subset, and then a loop
auto start = input.begin();
for(int i=0; i<(M-1); ++i) {
divider pred<T>(i);
result[i] = thrust::stable_partition(
thrust::device,
start,
input.end(),
pred);
start = result[i];
}
[ note all code written in a browser on a tablet while floating on a boat in the Baltic. Obviously never compiled or run. ]
This will certainly be the most space efficient, as a maximum of len(input) temporary storage is required, whereas a hypothetical single pass implementation would require M * len(input) storage, which would quickly get impractical for a large M.
Edit to add that now I'm back on land with a compiler, this seems to work as expected:
#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/copy.h>
#include <thrust/partition.h>
struct divider
{
int pass;
__host__ __device__
divider(int p) : pass(p) { };
__host__ __device__
int classify(const int &val) { return (val % 12); };
__host__ __device__
bool operator()(const int &val) { return !(classify(val) > pass); };
};
int main()
{
const int M = 12;
const int N = 120;
thrust::device_vector<int> input(N);
thrust::counting_iterator<int> iter(0);
thrust::copy(iter, iter+N, input.begin());
thrust::device_vector<int>::iterator result[M];
auto start = input.begin();
for(int i=0; i<(M-1); ++i) {
divider pred(i);
result[i] = thrust::stable_partition(
thrust::device,
start,
input.end(),
pred);
start = result[i];
}
int i = 0;
for(auto j=input.begin(); j!=input.end(); ++j) {
if (j == result[i]) {
i++;
std:: cout << std::endl;
}
std::cout << *j << " ";
}
return 0;
}
$ nvcc -std=c++11 -arch=sm_52 -o partition partition.cu
$ ./partition
0 12 24 36 48 60 72 84 96 108
1 13 25 37 49 61 73 85 97 109
2 14 26 38 50 62 74 86 98 110
3 15 27 39 51 63 75 87 99 111
4 16 28 40 52 64 76 88 100 112
5 17 29 41 53 65 77 89 101 113
6 18 30 42 54 66 78 90 102 114
7 19 31 43 55 67 79 91 103 115
8 20 32 44 56 68 80 92 104 116
9 21 33 45 57 69 81 93 105 117
10 22 34 46 58 70 82 94 106 118
11 23 35 47 59 71 83 95 107 119