Merge pull request #80 from mlochbaum/mergeshuf
Initial multi-way split for large shuffles
This commit is contained in:
commit
135ce4620c
@ -7,23 +7,21 @@
|
|||||||
#include "../nfns.h"
|
#include "../nfns.h"
|
||||||
|
|
||||||
|
|
||||||
|
NOINLINE B intRange16(ux s, ux n) { // s+↕n with i16arr result
|
||||||
NOINLINE B intRange(ux s, ux n) { // intended for s+n≥128; assumes n≥1
|
|
||||||
assert(n>0);
|
|
||||||
ux last = n+s-1;
|
|
||||||
if (last<=I16_MAX) {
|
|
||||||
i16* rp; B r = m_i16arrv(&rp, n);
|
i16* rp; B r = m_i16arrv(&rp, n);
|
||||||
for (ux i = 0; i < n; i++) rp[i] = (i16)s + (i16)i;
|
for (ux i = 0; i < n; i++) rp[i] = (i16)s + (i16)i;
|
||||||
return r;
|
return r;
|
||||||
}
|
}
|
||||||
if (last<=I32_MAX) {
|
NOINLINE void intRange32Fill(i32* xp, ux s, ux n) { // fill xp with s+↕n
|
||||||
|
for (ux i = 0; i < n; i++) xp[i] = (i32)s + (i32)i;
|
||||||
|
}
|
||||||
|
NOINLINE B intRange32(ux s, ux n) { // s+↕n with i32arr result
|
||||||
i32* rp; B r = m_i32arrv(&rp, n);
|
i32* rp; B r = m_i32arrv(&rp, n);
|
||||||
for (ux i = 0; i < n; i++) rp[i] = (i32)s + (i32)i;
|
intRange32Fill(rp, s, n);
|
||||||
return r;
|
return r;
|
||||||
}
|
}
|
||||||
|
NOINLINE B intRangeF64(ux s, ux n) { // s+↕n with f64arr result
|
||||||
f64* rp; B r = m_f64arrv(&rp, n);
|
f64* rp; B r = m_f64arrv(&rp, n);
|
||||||
|
|
||||||
f64 c = s;
|
f64 c = s;
|
||||||
PLAINLOOP for (ux i = 0; i < n/16; i++) {
|
PLAINLOOP for (ux i = 0; i < n/16; i++) {
|
||||||
for (ux j = 0; j < 16; j++) rp[j] = c+j;
|
for (ux j = 0; j < 16; j++) rp[j] = c+j;
|
||||||
@ -33,6 +31,13 @@ NOINLINE B intRange(ux s, ux n) { // intended for s+n≥128; assumes n≥1
|
|||||||
for (ux j = 0; j < (n&15); j++) rp[j] = c+j;
|
for (ux j = 0; j < (n&15); j++) rp[j] = c+j;
|
||||||
return r;
|
return r;
|
||||||
}
|
}
|
||||||
|
B intRange(ux s, ux n) { // intended for s+n≥128; assumes n≥1
|
||||||
|
assert(n>0);
|
||||||
|
ux last = n+s-1;
|
||||||
|
if (last<=I16_MAX) return intRange16(s, n);
|
||||||
|
if (last<=I32_MAX) return intRange32(s, n);
|
||||||
|
return intRangeF64(s, n);
|
||||||
|
}
|
||||||
|
|
||||||
static B* ud_rec(B* p, usz d, usz r, i32* pos, usz* sh) {
|
static B* ud_rec(B* p, usz d, usz r, i32* pos, usz* sh) {
|
||||||
usz cl = sh[d];
|
usz cl = sh[d];
|
||||||
|
|||||||
@ -425,19 +425,87 @@ B rand_range_c2(B t, B w, B x) {
|
|||||||
return taga(r);
|
return taga(r);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
extern Arr* bitUD[3]; // from fns.c
|
||||||
|
extern B bit2x[2];
|
||||||
|
B intRange16(ux s, ux n);
|
||||||
|
B intRange32(ux s, ux n);
|
||||||
|
void intRange32Fill(i32* xp, ux s, ux n);
|
||||||
|
|
||||||
B rand_deal_c1(B t, B x) {
|
B rand_deal_c1(B t, B x) {
|
||||||
i32 xi = o2i(x);
|
i32 xi = o2i(x);
|
||||||
if (RARE(xi<0)) thrM("(rand).Deal: Argument cannot be negative");
|
if (RARE(xi<=1)) {
|
||||||
if (xi==0) return emptyIVec();
|
if (xi<0) thrM("(rand).Deal: Argument cannot be negative");
|
||||||
RAND_START;
|
return taga(ptr_inc(bitUD[xi]));
|
||||||
i32* rp; B r = m_i32arrv(&rp, xi);
|
|
||||||
for (i64 i = 0; i < xi; i++) rp[i] = i;
|
|
||||||
for (i64 i = 0; i < xi; i++) {
|
|
||||||
i32 j = wy2u0k(wyrand(&seed), xi-i) + i;
|
|
||||||
i32 c = rp[j];
|
|
||||||
rp[j] = rp[i];
|
|
||||||
rp[i] = c;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
RAND_START;
|
||||||
|
B r;
|
||||||
|
#define SHUF \
|
||||||
|
for (usz i = 0; i < xi-1; i++) { \
|
||||||
|
usz j = wy2u0k(wyrand(&seed), xi-i) + i; \
|
||||||
|
usz c=rp[j]; rp[j]=rp[i]; rp[i]=c; \
|
||||||
|
}
|
||||||
|
if (xi == 2) {
|
||||||
|
r = incG(bit2x[wyrand(&seed)&1]);
|
||||||
|
} else if (LIKELY(xi <= 128)) {
|
||||||
|
i8* rp; r = m_i8arrv(&rp, xi);
|
||||||
|
NOUNROLL for (usz i = 0; i < xi; i++) rp[i] = i;
|
||||||
|
SHUF
|
||||||
|
} else if (LIKELY(xi <= 1<<15)) {
|
||||||
|
r = intRange16(0, xi); i16* rp = i16arr_ptr(r);
|
||||||
|
SHUF
|
||||||
|
} else {
|
||||||
|
if (xi <= 1<<19) {
|
||||||
|
r = intRange32(0, xi); i32* rp = i32arr_ptr(r);
|
||||||
|
SHUF
|
||||||
|
} else {
|
||||||
|
i32* rp; r = m_i32arrv(&rp, xi);
|
||||||
|
// Initial split pass like a random radix sort
|
||||||
|
// Don't count partition size exactly; instead, assume lengths
|
||||||
|
// are within 1 and stop when a partition is full
|
||||||
|
// Shuffle leftovers in at the end
|
||||||
|
u64 n = xi;
|
||||||
|
usz log2 = 64 - CLZ(n-1);
|
||||||
|
usz thr = 16;
|
||||||
|
usz sd = log2<thr+8? log2-thr : 8;
|
||||||
|
usz m = 1<<sd; // Number of partitions
|
||||||
|
u64 mm = 0x0101010101010101ull * (m-1);
|
||||||
|
TALLOC(usz, pos, 2*m) // Current and ending positions
|
||||||
|
pos[0] = 0; pos[2*m-1] = n;
|
||||||
|
PLAINLOOP for (usz p=1; p<m; p++) pos[2*p-1] = pos[2*p] = (n*p) >> sd;
|
||||||
|
usz i = 0;
|
||||||
|
while (1) {
|
||||||
|
u64 r = wyrand(&seed) & mm;
|
||||||
|
for (usz j = 0; j < 8; j++) {
|
||||||
|
u8 k = r; r>>= 8;
|
||||||
|
usz* pp = pos+2*k;
|
||||||
|
usz p = pp[0];
|
||||||
|
if (p == pp[1]) goto split_done;
|
||||||
|
pp[0]++;
|
||||||
|
rp[p] = i++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
split_done:
|
||||||
|
for (usz p=0, b=0, s=0; p<m; p++) {
|
||||||
|
usz l = pos[2*p] - s;
|
||||||
|
i32* dp = rp+b;
|
||||||
|
i32* sp = rp+s;
|
||||||
|
for (usz i = 0; i < l; i++) {
|
||||||
|
usz j = wy2u0k(wyrand(&seed), l-i) + i;
|
||||||
|
usz c=sp[j]; sp[j]=sp[i]; dp[i]=c;
|
||||||
|
}
|
||||||
|
s = pos[2*p+1];
|
||||||
|
b+= l;
|
||||||
|
}
|
||||||
|
TFREE(pos)
|
||||||
|
intRange32Fill(rp+i, i, n-i);
|
||||||
|
for (; i < n; i++) {
|
||||||
|
usz j = wy2u0k(wyrand(&seed), 1+i);
|
||||||
|
usz c=rp[j]; rp[j]=rp[i]; rp[i]=c;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
#undef SHUF
|
||||||
RAND_END;
|
RAND_END;
|
||||||
return r;
|
return r;
|
||||||
}
|
}
|
||||||
@ -445,10 +513,11 @@ B rand_deal_c1(B t, B x) {
|
|||||||
B rand_deal_c2(B t, B w, B x) {
|
B rand_deal_c2(B t, B w, B x) {
|
||||||
i32 wi = o2i(w);
|
i32 wi = o2i(w);
|
||||||
i32 xi = o2i(x);
|
i32 xi = o2i(x);
|
||||||
if (RARE(wi<0)) thrM("(rand).Deal: 𝕨 cannot be negative");
|
|
||||||
if (RARE(xi<0)) thrM("(rand).Deal: 𝕩 cannot be negative");
|
if (RARE(xi<0)) thrM("(rand).Deal: 𝕩 cannot be negative");
|
||||||
if (RARE(wi>xi)) thrM("(rand).Deal: 𝕨 cannot exceed 𝕩");
|
if (RARE(wi<0)) thrM("(rand).Deal: 𝕨 cannot be negative");
|
||||||
if (wi==0) return emptyIVec();
|
if (wi==0) return emptyIVec();
|
||||||
|
if (RARE(wi>xi)) thrM("(rand).Deal: 𝕨 cannot exceed 𝕩");
|
||||||
|
if (wi==xi) return rand_deal_c1(t, x);
|
||||||
B r;
|
B r;
|
||||||
RAND_START;
|
RAND_START;
|
||||||
if (wi > xi/64) {
|
if (wi > xi/64) {
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user