magyarsort/threepass_xbit.h

#ifndef THREE_PASS_XB_H
#define THREE_PASS_XB_H
/* Basically like threepass.h but but fewer bits - only for thiersort3! */

/* How the 28-30 bits gets separated? This is chosen because I saw that: */
/* - in float trickery I saw at around 4 billion the buckets change around 200 million */
/* - Which is smaller than 1/2 or 1/4 or 1/8 or 1/16 of it because 16*200 mil = 3.2 bill 28 bit is needed */
/* - This means those bits in numbers DO get used when we are running as internal sort of thier3 */
/* - But the numbers high bits (29..32th bits) stay the same inside each bucket so we spare it! */
/* - If you know your data does not have all the range of 32 bits you can #define these by #define CUSTOM_TPBX_BITS */
#ifndef CUSTOM_TPBX_BITS
#define TPBX1 10 // top
#define TPBX2 9 // mid
#define TPBX3 9 // bottom
#endif /* CUSTOM_TPBX_BITS */

static inline constexpr uint32_t min3u32_xb(uint32_t a, uint32_t b, uint32_t c) noexcept {
	return (a <= b) ?
		((a <= c) ? a : c) :
		((b <= c) ? b : c);
}

static inline void tpxb_process_element(uint32_t num, uint32_t* arr, uint32_t* bucket,
		uint32_t shr, uint32_t mask) {
	auto bkeyni = (num >> shr) & mask;
	auto offset = --bucket[bkeyni];
	arr[offset] = num;
}

/**
 * Simple three-pass (ok: 3 + 1) bottom-up internal radix sort writter for thiersort3
 *
 * @param a The array to sort - will be changed too!
 * @param buf Result array with the same size - result will be here
 * @param n The number of elements
 */
static inline void threepass_xb(uint32_t *a, uint32_t *buf, int n) noexcept {
	assert(buf != NULL);
	constexpr int shr1 = TPBX3 + TPBX2;
	constexpr int shr2 = TPBX3;
	constexpr int shr3 = 0;
	constexpr int mask1 = (1 << TPBX1) - 1;
	constexpr int mask2 = (1 << TPBX2) - 1;
	constexpr int mask3 = (1 << TPBX3) - 1;

	/* helper buffers. */
	int sz = n * sizeof(a[0]);

	static thread_local uint32_t bucket1[1 << TPBX1];
	memset(bucket1, 0, (1 << TPBX1) * sizeof(uint32_t));
	static thread_local uint32_t bucket2[1 << TPBX2];
	memset(bucket2, 0, (1 << TPBX2) * sizeof(uint32_t));
	static thread_local uint32_t bucket3[1 << TPBX3];
	memset(bucket3, 0, (1 << TPBX3) * sizeof(uint32_t));

	/* Count occurences (can count together with good ILP) */
	#pragma GCC unroll 64
	for(uint32_t i = 0; i < n; ++i) {
		++bucket1[(a[i] >> shr1) & mask1];
		++bucket2[(a[i] >> shr2) & mask2];
		++bucket3[(a[i] >> shr3) & mask3];
	}

	/* Count prefix sums - try as much ILP as possible because bigger arrays than usual! */
	uint32_t prev1 = 0;
	uint32_t prev2 = 0;
	uint32_t prev3 = 0;
	uint32_t common = min3u32_xb(
		(1 << TPBX1),
		(1 << TPBX2),
		(1 << TPBX3)
	);
	int i = 0;
	#pragma GCC unroll 8
	for (; i < common; ++i) {
		bucket1[i] += prev1;
		prev1 = bucket1[i];
		bucket2[i] += prev2;
		prev2 = bucket2[i];
		bucket3[i] += prev3;
		prev3 = bucket3[i];
	}
	/* Do remaining 1 */
	for (int j = i; j < (1 << TPBX1); ++j) {
		bucket1[j] += prev1;
		prev1 = bucket1[j];
	}
	/* Do remaining 2 */
	for (int j = i; j< (1 << TPBX2); ++j) {
		bucket2[j] += prev2;
		prev2 = bucket2[j];
	}
	/* Do remaining 3 */
	for (int j = i; j < (1 << TPBX3); ++j) {
		bucket3[j] += prev3;
		prev3 = bucket3[j];
	}

	// Bottom digit a->buf
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 3
	for(i = n; i >= 16; i -= 16) {
		// Prefetch the NEXT block (not current) at optimal distance
		if (i > 17) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17]);
		}
		if (i > 17*2) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17*2]);
		}
		if (i > 17*3) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17*3]);
		}

		// Process 16 elements in reverse order
		auto num15 = a[i - 1];
		auto num14 = a[i - 2];
		auto num13 = a[i - 3];
		auto num12 = a[i - 4];
		auto num11 = a[i - 5];
		auto num10 = a[i - 6];
		auto num9  = a[i - 7];
		auto num8  = a[i - 8];
		auto num7  = a[i - 9];
		auto num6  = a[i - 10];
		auto num5  = a[i - 11];
		auto num4  = a[i - 12];
		auto num3  = a[i - 13];
		auto num2  = a[i - 14];
		auto num1  = a[i - 15];
		auto num0  = a[i - 16];

		// Process all 16 elements (your bucket logic here)
		tpxb_process_element(num15, buf, bucket3, shr3, mask3);
		tpxb_process_element(num14, buf, bucket3, shr3, mask3);
		tpxb_process_element(num13, buf, bucket3, shr3, mask3);
		tpxb_process_element(num12, buf, bucket3, shr3, mask3);
		tpxb_process_element(num11, buf, bucket3, shr3, mask3);
		tpxb_process_element(num10, buf, bucket3, shr3, mask3);
		tpxb_process_element(num9, buf, bucket3, shr3, mask3);
		tpxb_process_element(num8, buf, bucket3, shr3, mask3);
		tpxb_process_element(num7, buf, bucket3, shr3, mask3);
		tpxb_process_element(num6, buf, bucket3, shr3, mask3);
		tpxb_process_element(num5, buf, bucket3, shr3, mask3);
		tpxb_process_element(num4, buf, bucket3, shr3, mask3);
		tpxb_process_element(num3, buf, bucket3, shr3, mask3);
		tpxb_process_element(num2, buf, bucket3, shr3, mask3);
		tpxb_process_element(num1, buf, bucket3, shr3, mask3);
		tpxb_process_element(num0, buf, bucket3, shr3, mask3);
	}
	// Handle remainder (less than 16 elements)
	for(uint32_t j = i; j > 0; --j) {
		auto num = a[j - 1];
		auto bkeyni = (num >> shr3) & mask3;
		auto offset = --bucket3[bkeyni];
		buf[offset] = num;
	}

	// Mid digit buf->a
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 3
	for(i = n; i >= 16; i -= 16) {
		// Prefetch the NEXT block (not current) at optimal distance
		if (i > 17) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&buf[i - 17]);
		}
		if (i > 17*2) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&buf[i - 17*2]);
		}
		if (i > 17*3) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&buf[i - 17*3]);
		}

		// Process 16 elements in reverse order
		auto num15 = buf[i - 1];
		auto num14 = buf[i - 2];
		auto num13 = buf[i - 3];
		auto num12 = buf[i - 4];
		auto num11 = buf[i - 5];
		auto num10 = buf[i - 6];
		auto num9  = buf[i - 7];
		auto num8  = buf[i - 8];
		auto num7  = buf[i - 9];
		auto num6  = buf[i - 10];
		auto num5  = buf[i - 11];
		auto num4  = buf[i - 12];
		auto num3  = buf[i - 13];
		auto num2  = buf[i - 14];
		auto num1  = buf[i - 15];
		auto num0  = buf[i - 16];

		// Process all 16 elements (your bucket logic here)
		tpxb_process_element(num15, a, bucket2, shr2, mask2);
		tpxb_process_element(num14, a, bucket2, shr2, mask2);
		tpxb_process_element(num13, a, bucket2, shr2, mask2);
		tpxb_process_element(num12, a, bucket2, shr2, mask2);
		tpxb_process_element(num11, a, bucket2, shr2, mask2);
		tpxb_process_element(num10, a, bucket2, shr2, mask2);
		tpxb_process_element(num9, a, bucket2, shr2, mask2);
		tpxb_process_element(num8, a, bucket2, shr2, mask2);
		tpxb_process_element(num7, a, bucket2, shr2, mask2);
		tpxb_process_element(num6, a, bucket2, shr2, mask2);
		tpxb_process_element(num5, a, bucket2, shr2, mask2);
		tpxb_process_element(num4, a, bucket2, shr2, mask2);
		tpxb_process_element(num3, a, bucket2, shr2, mask2);
		tpxb_process_element(num2, a, bucket2, shr2, mask2);
		tpxb_process_element(num1, a, bucket2, shr2, mask2);
		tpxb_process_element(num0, a, bucket2, shr2, mask2);
	}
	// Handle remainder (less than 16 elements)
	for(uint32_t j = i; j > 0; --j) {
		auto num = buf[j - 1];
		auto bkeyni = (num >> shr2) & mask2;
		auto offset = --bucket2[bkeyni];
		a[offset] = num;
	}
	// Top digit a->buf
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 3
	for(i = n; i >= 16; i -= 16) {
		// Prefetch the NEXT block (not current) at optimal distance
		if (i > 17) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17]);
		}
		if (i > 17*2) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17*2]);
		}
		if (i > 17*3) {  // Ensure we don't prefetch out of bounds
			__builtin_prefetch(&a[i - 17*3]);
		}

		// Process 16 elements in reverse order
		auto num15 = a[i - 1];
		auto num14 = a[i - 2];
		auto num13 = a[i - 3];
		auto num12 = a[i - 4];
		auto num11 = a[i - 5];
		auto num10 = a[i - 6];
		auto num9  = a[i - 7];
		auto num8  = a[i - 8];
		auto num7  = a[i - 9];
		auto num6  = a[i - 10];
		auto num5  = a[i - 11];
		auto num4  = a[i - 12];
		auto num3  = a[i - 13];
		auto num2  = a[i - 14];
		auto num1  = a[i - 15];
		auto num0  = a[i - 16];

		// Process all 16 elements (your bucket logic here)
		tpxb_process_element(num15, buf, bucket1, shr1, mask1);
		tpxb_process_element(num14, buf, bucket1, shr1, mask1);
		tpxb_process_element(num13, buf, bucket1, shr1, mask1);
		tpxb_process_element(num12, buf, bucket1, shr1, mask1);
		tpxb_process_element(num11, buf, bucket1, shr1, mask1);
		tpxb_process_element(num10, buf, bucket1, shr1, mask1);
		tpxb_process_element(num9, buf, bucket1, shr1, mask1);
		tpxb_process_element(num8, buf, bucket1, shr1, mask1);
		tpxb_process_element(num7, buf, bucket1, shr1, mask1);
		tpxb_process_element(num6, buf, bucket1, shr1, mask1);
		tpxb_process_element(num5, buf, bucket1, shr1, mask1);
		tpxb_process_element(num4, buf, bucket1, shr1, mask1);
		tpxb_process_element(num3, buf, bucket1, shr1, mask1);
		tpxb_process_element(num2, buf, bucket1, shr1, mask1);
		tpxb_process_element(num1, buf, bucket1, shr1, mask1);
		tpxb_process_element(num0, buf, bucket1, shr1, mask1);
	}
	// Handle remainder (less than 16 elements)
	for(uint32_t j = i; j > 0; --j) {
		auto num = a[j - 1];
		auto bkeyni = (num >> shr1) & mask1;
		auto offset = --bucket1[bkeyni];
		buf[offset] = num;
	}
}

#endif /* THREE_PASS_XB_H */