A radix sort on sortable 64-bit keys. It first finds the highest 16-bit digit in which the keys differ and partitions on that, with a histogram for each part of the input.
the radix sort behind it · neon-cpp/bench/lang/rust/src/main.rs:1338 · show all 143 lines
fn radix_perm_mt(amt: &[f64], pool: &Pool, threads: usize) -> Vec<u32> {
const NB: usize = 1 << 16;
let prof = PROFILE_SORT.load(Relaxed);
let t0 = Instant::now();
let n = amt.len();
let parts = threads.max(1);
let ranges = chunk_ranges(n, parts);
// the highest 16-bit digit in which the keys differ decides the partition digit
let mm: Vec<(u64, u64)> = pool.map(parts, |q| {
let (a, b) = ranges[q];
amt[a..b].iter().fold((u64::MAX, 0u64), |(lo, hi), &x| {
let k = sortable(x);
(lo.min(k), hi.max(k))
})
});
let (lo, hi) = mm.iter().fold((u64::MAX, 0u64), |(l, h), &(a, b)| (l.min(a), h.max(b)));
if lo == hi {
return (0..n as u32).collect();
}
let msd_shift = ((63 - (lo ^ hi).leading_zeros()) / 16) * 16;
let t1 = Instant::now();
// pass 1: per-part histogram, local exclusive prefix, local scatter (part-major
// layout). Afterwards hist[q][d] is the END of bucket d inside part q, so segment d
// is [hist[q][d-1], hist[q][d]).
let mut hist: Vec<u32> = Vec::with_capacity(parts * NB);
let mut pk: Vec<u64> = Vec::with_capacity(n);
let mut pi: Vec<u32> = Vec::with_capacity(n);
let (hp, pkp, pip) = (SyncPtr(hist.as_mut_ptr()), SyncPtr(pk.as_mut_ptr()), SyncPtr(pi.as_mut_ptr()));
pool.run(parts, |q| unsafe {
let (a, b) = ranges[q];
let h = hp.get().add(q * NB);
std::ptr::write_bytes(h, 0, NB);
for &x in &amt[a..b] {
*h.add((sortable(x) >> msd_shift) as usize & (NB - 1)) += 1;
}
let mut sum = 0u32;
for d in 0..NB {
let c = *h.add(d);
*h.add(d) = sum;
sum += c;
}
for i in a..b {
let k = sortable(*amt.get_unchecked(i));
let d = (k >> msd_shift) as usize & (NB - 1);
let o = a + *h.add(d) as usize;
*h.add(d) += 1;
*pkp.get().add(o) = k;
*pip.get().add(o) = i as u32;
}
});
unsafe {
hist.set_len(parts * NB);
pk.set_len(n);
pi.set_len(n);
}
let t2 = Instant::now();
// bucket totals (parallel over bucket ranges), then global bucket starts
let mut total: Vec<u32> = Vec::with_capacity(NB);
let tp = SyncPtr(total.as_mut_ptr());
let hist_ref = &hist;
let bucket_ranges = chunk_ranges(NB, threads);
pool.run(threads, |t| unsafe {
let (d0, d1) = bucket_ranges[t];
let out = std::slice::from_raw_parts_mut(tp.get().add(d0), d1 - d0);
out.fill(0);
for q in 0..parts {
let h = &hist_ref[q * NB..(q + 1) * NB];
let prev = if d0 > 0 { h[d0 - 1] } else { 0 };
let seg = &h[d0..d1];
let mut last = prev;
for (o, &e) in out.iter_mut().zip(seg) {
*o += e - last;
last = e;
}
}
});
unsafe { total.set_len(NB) };
let mut bstart = vec![0u32; NB + 1];
for d in 0..NB {
bstart[d + 1] = bstart[d] + total[d];
}
// tasks: contiguous bucket ranges holding ~n / (4 * threads) elements each
let per = (n + 4 * threads - 1) / (4 * threads);
let mut cuts = vec![0usize];
let mut acc = 0usize;
for d in 0..NB {
if acc >= per && total[d] > 0 {
cuts.push(d);
acc = 0;
}
acc += total[d] as usize;
}
cuts.push(NB);
let ntasks = cuts.len() - 1;
let t3 = Instant::now();
// pass 2: gather each bucket from the parts into scratch, sort it, emit its indices
let mut oi: Vec<u32> = Vec::with_capacity(n);
let oip = SyncPtr(oi.as_mut_ptr());
let (bstart, total, cuts, pk, pi) = (&bstart, &total, &cuts, &pk, &pi);
pool.run(ntasks, |t| unsafe {
let (d0, d1) = (cuts[t], cuts[t + 1]);
let maxb = (d0..d1).map(|d| total[d] as usize).max().unwrap_or(0);
let mut ka: Vec<u64> = Vec::with_capacity(maxb);
let mut ia: Vec<u32> = Vec::with_capacity(maxb);
let mut kb: Vec<u64> = Vec::with_capacity(maxb);
let mut ib: Vec<u32> = Vec::with_capacity(maxb);
for d in d0..d1 {
let cnt = total[d] as usize;
if cnt == 0 {
continue;
}
let mut w = 0usize;
for q in 0..parts {
let h = &hist_ref[q * NB..(q + 1) * NB];
let (s0, e0) = (if d > 0 { h[d - 1] as usize } else { 0 }, h[d] as usize);
if e0 > s0 {
let from = ranges[q].0 + s0;
std::ptr::copy_nonoverlapping(pk.as_ptr().add(from), ka.as_mut_ptr().add(w), e0 - s0);
std::ptr::copy_nonoverlapping(pi.as_ptr().add(from), ia.as_mut_ptr().add(w), e0 - s0);
w += e0 - s0;
}
}
let in_b = msd_shift > 0
&& sort_bucket(ka.as_mut_ptr(), ia.as_mut_ptr(), kb.as_mut_ptr(), ib.as_mut_ptr(), cnt, msd_shift);
let src = if in_b { ib.as_ptr() } else { ia.as_ptr() };
std::ptr::copy_nonoverlapping(src, oip.get().add(bstart[d] as usize), cnt);
}
});
unsafe { oi.set_len(n) };
if prof {
let t4 = Instant::now();
println!(
" sort phases: minmax {:.3} partition {:.3} totals+cuts {:.3} sub-sorts {:.3} ms ({} tasks, msd_shift {})",
(t1 - t0).as_secs_f64() * 1e3,
(t2 - t1).as_secs_f64() * 1e3,
(t3 - t2).as_secs_f64() * 1e3,
(t4 - t3).as_secs_f64() * 1e3,
ntasks,
msd_shift
);
}
oi
}