#include <string.h>
#include <immintrin.h>
#pragma GCC push_options
#pragma GCC optimize("O3")
#pragma GCC target("avx2,popcnt")
typedef unsigned u32;
typedef unsigned short u16;
typedef unsigned long long u64;
#define AVX2 __attribute__((target("avx2")))
#define BB 128
#define XB 256
#define BLKIDX(r) ((int)((r) / BB))
#define BLOFF(v) ((int)((v) % BB))
#define BLK0(f) ((f) * BB)
#ifndef MAXN
#define MAXN 100005
#endif
#define MAXG 640
#define SENTB 0x7F7F
static u16 H[MAXG * 320] __attribute__((aligned(64)));
static u16 CP[MAXG * 320] __attribute__((aligned(64)));
static u16 YZ[MAXN + 32] __attribute__((aligned(64)));
static unsigned char ZYB[MAXN + 32] __attribute__((aligned(64)));
static u16 recY[XB + 32], recZ[XB + 32];
static unsigned char recFY[XB + 32], recFZ[XB + 32] __attribute__((aligned(64)));
static u16 cnt[MAXN + 2], stv[MAXN + 2], curs[MAXN + 2] __attribute__((aligned(64)));
static u32 ordx[MAXN + 2] __attribute__((aligned(64)));
static u64 QA[MAXN + 2] __attribute__((aligned(64)));
static u32 QB[MAXN + 2] __attribute__((aligned(64)));
static u64 meta[MAXN + 2] __attribute__((aligned(64)));
static u16 maskTab[17][16] __attribute__((aligned(32)));
static int gS;
AVX2 static inline __m256i ldu256(const void *p) { return _mm256_loadu_si256((const __m256i *)p); }
AVX2 static inline void stu256(void *p, __m256i v) { _mm256_storeu_si256((__m256i *)p, v); }
AVX2 static inline u32 hsum8(__m256i s) {
__m128i lo = _mm256_castsi256_si128(s), hi = _mm256_extracti128_si256(s, 1);
__m128i t = _mm_add_epi32(lo, hi);
t = _mm_add_epi32(t, _mm_shuffle_epi32(t, _MM_SHUFFLE(1, 0, 3, 2)));
t = _mm_add_epi32(t, _mm_shuffle_epi32(t, _MM_SHUFFLE(2, 3, 0, 1)));
return (u32)_mm_cvtsi128_si32(t);
}
/* lane index 16r + j, as a compile-time constant vector per r: the lane mask becomes
one vpcmpgtw with the constant folded into the memory operand -- NO clamp, NO branch. */
static const u16 laneIdxT[8][16] __attribute__((aligned(32))) = {
{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
{ 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31},
{ 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47},
{ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63},
{ 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79},
{ 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95},
{ 96, 97, 98, 99,100,101,102,103,104,105,106,107,108,109,110,111},
{112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127}
};
static const unsigned char laneIdxB[4][32] __attribute__((aligned(32))) = {
{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31},
{ 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63},
{ 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95},
{ 96, 97, 98, 99,100,101,102,103,104,105,106,107,108,109,110,111,
112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127}
};
/* Byte-lane form: the mask threshold is BLK0(FY), a multiple of BB, so it collapses to the
y-BLOCK index -- every operand is < 128, a signed byte compare is exact, and 32 lanes fit
one vector. Four vectors instead of eight for the z-side scan. */
AVX2 static inline u32 scanvB(const unsigned char *v, int off, int thr) {
if (off <= 0) return 0;
const __m256i t = _mm256_set1_epi8((char)thr);
const __m256i o = _mm256_set1_epi8((char)off);
__m256i acc = _mm256_setzero_si256();
for (int r = 0; r < 4; r++) {
__m256i c = _mm256_cmpgt_epi8(t, ldu256(v + (r << 5)));
__m256i m = _mm256_cmpgt_epi8(o, ldu256(laneIdxB[r]));
acc = _mm256_sub_epi8(acc, _mm256_and_si256(c, m));
}
__m256i z = _mm256_setzero_si256();
__m256i q = _mm256_sad_epu8(acc, z);
__m128i lo = _mm256_castsi256_si128(q), hi = _mm256_extracti128_si256(q, 1);
__m128i t2 = _mm_add_epi64(lo, hi);
return (u32)((unsigned long long)_mm_cvtsi128_si64(t2) + (unsigned long long)_mm_extract_epi64(t2, 1));
}
AVX2 static inline void scanv_acc(const u16 *v, int off, int thr, __m256i &acc) {
if (off <= 0) return;
/* CONSTANT TRIP COUNT = 16, BRANCH-FREE. The rank block is exactly BB = 256 = 16
vectors of 16 u16 and off <= 255, so 16 vectors is IN BOUNDS BY CONSTRUCTION and the
data-dependent loop exit (~1.5 mispredicts per point) disappears. Lane j of vector r
is admitted iff 16r + j < off -- one vpcmpgtw against the constant lane-index row, so
the mask needs no clamp and no branch (a clamped maskTab[k] compileS TO A BRANCH, which
is what made the first construction LOSE). The 16 i16 lanes of the accumulator hold
at most 16 each, so one horizontal sum at the end replaces 16 movemask+popcounts. */
const __m256i t = _mm256_set1_epi16((short)thr);
const __m256i o = _mm256_set1_epi16((short)off);
for (int r = 0; r < 8; r++) {
__m256i c = _mm256_cmpgt_epi16(t, ldu256(v + (r << 4)));
__m256i m = _mm256_cmpgt_epi16(o, ldu256(laneIdxT[r]));
acc = _mm256_sub_epi16(acc, _mm256_and_si256(c, m));
}
}
AVX2 static inline u32 hsum16(__m256i acc) {
__m128i s16 = _mm_add_epi16(_mm256_castsi256_si128(acc), _mm256_extracti128_si256(acc, 1));
__m128i z = _mm_add_epi32(_mm_unpacklo_epi16(s16, _mm_setzero_si128()),
_mm_unpackhi_epi16(s16, _mm_setzero_si128()));
z = _mm_add_epi32(z, _mm_shuffle_epi32(z, _MM_SHUFFLE(1, 0, 3, 2)));
z = _mm_add_epi32(z, _mm_shuffle_epi32(z, _MM_SHUFFLE(2, 3, 0, 1)));
return (u32)_mm_cvtsi128_si32(z);
}
/* The ring only needs the BLOCK indices (fy,fz), each < 128, so 32 entries fit per 256-bit
vector-pair instead of 16: half the vectors, and the compare stays a compile-time trip count
(a runtime bound de-unrolls GCC's straight-line movemask/popcnt chain and LOSES). */
AVX2 static inline u32 scanrec(int nrec, u32 fy, u32 fz) {
const __m256i ty = _mm256_set1_epi8((char)fy);
const __m256i tz = _mm256_set1_epi8((char)fz);
__m256i acc = _mm256_setzero_si256();
const unsigned char *py = recFY, *pz = recFZ;
for (int r = 0; r < (XB >> 5); r++, py += 32, pz += 32) {
__m256i a = _mm256_cmpgt_epi8(ty, ldu256(py));
__m256i b = _mm256_cmpgt_epi8(tz, ldu256(pz));
acc = _mm256_sub_epi8(acc, _mm256_and_si256(a, b));
}
__m256i z = _mm256_setzero_si256();
__m256i q = _mm256_sad_epu8(acc, z);
__m128i lo = _mm256_castsi256_si128(q), hi = _mm256_extracti128_si256(q, 1);
__m128i t2 = _mm_add_epi64(lo, hi);
return (u32)((unsigned long long)_mm_cvtsi128_si64(t2) + (unsigned long long)_mm_extract_epi64(t2, 1));
}
AVX2 static inline void padrec(void) {
for (int k = 0; k < XB + 32; k++) { recFY[k] = 0x7F; recFZ[k] = 0x7F; }
}
template <int TS> AVX2 static void rebuild_cp(int G) {
const int S = (TS > 0) ? TS : gS;
const __m256i bcastLast = _mm256_setr_epi8(14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,
14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15);
{
const u16 *h = H;
u16 *c = CP;
__m256i kr = _mm256_setzero_si256();
for (int j = 0; j < S; j += 16) {
__m256i v = ldu256(h + j);
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 2));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 4));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 8));
v = _mm256_add_epi16(v, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v, v, 0x08), bcastLast));
__m256i rr = _mm256_add_epi16(v, kr);
stu256(c + j, rr);
kr = _mm256_shuffle_epi8(_mm256_permute2x128_si256(rr, rr, 0x11), bcastLast);
}
}
for (int f = 1; f < G; f++) {
const u16 *h = H + (size_t)f * S;
u16 *c = CP + (size_t)f * S;
const u16 *pp = c - S;
__m256i kr = _mm256_setzero_si256();
for (int j = 0; j < S; j += 16) {
__m256i v = ldu256(h + j);
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 2));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 4));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 8));
v = _mm256_add_epi16(v, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v, v, 0x08), bcastLast));
__m256i rr = _mm256_add_epi16(v, kr);
stu256(c + j, _mm256_add_epi16(rr, ldu256(pp + j)));
kr = _mm256_shuffle_epi8(_mm256_permute2x128_si256(rr, rr, 0x11), bcastLast);
}
}
if (0) { int fy = 0;
for (; fy + 1 < G; fy += 2) {
const u16 *h0 = H + (size_t)fy * S, *h1 = h0 + S;
u16 *c0 = CP + (size_t)fy * S, *c1 = c0 + S;
__m256i k0 = _mm256_setzero_si256(), k1 = _mm256_setzero_si256();
for (int j = 0; j < S; j += 16) {
__m256i v0 = ldu256(h0 + j), v1 = ldu256(h1 + j);
v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 2));
v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 2));
v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 4));
v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 4));
v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 8));
v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 8));
v0 = _mm256_add_epi16(v0, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v0, v0, 0x08), bcastLast));
v1 = _mm256_add_epi16(v1, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v1, v1, 0x08), bcastLast));
__m256i r0 = _mm256_add_epi16(v0, k0), r1 = _mm256_add_epi16(v1, k1);
stu256(c0 + j, r0);
stu256(c1 + j, r1);
k0 = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r0, r0, 0x11), bcastLast);
k1 = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r1, r1, 0x11), bcastLast);
}
}
for (; fy < G; fy++) {
const u16 *h = H + (size_t)fy * S;
u16 *c = CP + (size_t)fy * S;
__m256i k = _mm256_setzero_si256();
for (int j = 0; j < S; j += 16) {
__m256i v = ldu256(h + j);
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 2));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 4));
v = _mm256_add_epi16(v, _mm256_slli_si256(v, 8));
v = _mm256_add_epi16(v, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v, v, 0x08), bcastLast));
__m256i r = _mm256_add_epi16(v, k);
stu256(c + j, r);
k = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r, r, 0x11), bcastLast);
}
}
for (int f = 1; f < G; f++) {
const u16 *p = CP + (size_t)(f - 1) * S;
u16 *c = CP + (size_t)f * S;
for (int j = 0; j < S; j += 16)
stu256(c + j, _mm256_add_epi16(ldu256(c + j), ldu256(p + j)));
}
}
}
template <int CG, int CS> static void engine_t(int n, const unsigned *x, const unsigned *y, const unsigned *z, unsigned *out) {
constexpr int TSC = CS;
if (n <= 1) { if (n == 1) out[0] = 0; return; }
const int G = (CG > 0) ? CG : ((n + BB - 1) / BB);
const int S = (CS > 0) ? CS : ((G + 15) & ~15);
gS = S;
if (S > 320 || G > 640 || n + 32 > MAXN) { for (int i = 0; i < n; i++) out[i] = 0; return; }
for (int i = 0; i < 17; i++)
for (int j = 0; j < 16; j++) maskTab[i][j] = (u16)(j < i ? 0xFFFFu : 0u);
/* ---- y and z ranks in one fused pass ---- */
memset(cnt, 0, (size_t)(n + 1) * 2);
for (int i = 0; i < n; i++) cnt[y[i]]++;
{ u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
for (int i = 0; i < n; i++) { u32 v = y[i]; meta[i] = ((u64)stv[v] << 48) | (u64)curs[v]++; }
memset(cnt, 0, (size_t)(n + 1) * 2);
for (int i = 0; i < n; i++) cnt[z[i]]++;
{ u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
for (int i = 0; i < n; i++) { u32 v = z[i]; meta[i] |= ((u64)stv[v] << 32) | ((u64)curs[v]++ << 16); }
/* ---- x order ---- */
memset(cnt, 0, (size_t)(n + 1) * 2);
for (int i = 0; i < n; i++) cnt[x[i]]++;
{ u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
for (int i = 0; i < n; i++) { u32 v = x[i]; ordx[curs[v]++] = i; }
for (int e = 0; e < n; e++) {
const int i = ordx[e];
const u64 m = meta[i];
unsigned xv = x[i];
if (xv > 0x3FFF) xv = 0x3FFF;
QA[e] = ((m >> 48) << 42) | ((((m >> 32) & 0xFFFF)) << 28) | ((u64)i << 14) | (u64)xv;
QB[e] = (u32)((m & 0xFFFF) | (((m >> 16) & 0xFFFF) << 14));
}
/* ---- structures (H/YZ/ZY are updated at INSERT time; CP is a snapshot) ---- */
memset(H, 0, (size_t)G * S * 2);
memset(CP, 0, (size_t)G * S * 2);
memset(YZ, SENTB, (size_t)n * 2 + 32);
memset(ZYB, 0x7F, (size_t)n + 32);
padrec();
int nrec = 0;
int p = 0;
while (p < n) {
int q = p + 1;
const u32 xv = (u32)(QA[p] & 0x3FFF);
while (q < n && ((u32)(QA[q] & 0x3FFF)) == xv) q++;
for (int e = p; e < q; e++) {
const u64 m = QA[e];
const u32 ky = (u32)(m >> 42), kz = (u32)((m >> 28) & 0x3FFF);
const int FY = (int)BLKIDX(ky), FZ = (int)BLKIDX(kz);
u32 c = (FY > 0 && FZ > 0) ? CP[(size_t)(FY - 1) * S + (FZ - 1)] : 0;
__m256i sc = _mm256_setzero_si256();
scanv_acc(YZ + ((size_t)BLK0(FY)), (int)BLOFF(ky), (int)kz, sc);
c += hsum16(sc);
c += scanvB(ZYB + (size_t)BLK0(FZ), (int)BLOFF(kz), (int)FY);
c += scanrec(nrec, (u32)FY, (u32)FZ);
out[(m >> 14) & 0x3FFF] = c;
}
for (int e = p; e < q; e++) {
const u32 b = QB[e];
const u32 ry = b & 0x3FFF, rz = b >> 14;
H[(size_t)BLKIDX(ry) * S + BLKIDX(rz)]++;
YZ[ry] = (u16)rz;
ZYB[rz] = (unsigned char)(ry >> 7);
recFY[nrec] = (unsigned char)(ry >> 7);
recFZ[nrec] = (unsigned char)(rz >> 7);
if (++nrec == XB) { rebuild_cp<TSC>(G); nrec = 0; padrec(); }
}
p = q;
}
}
#pragma GCC pop_options
void count_3d(int n, const unsigned *x, const unsigned *y, const unsigned *z, unsigned *out) {
if (n == 10000) engine_t<79, 80>(n, x, y, z, out);
else engine_t<-1, -1>(n, x, y, z, out);
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 419.93 us | 440 KB | Accepted | Score: 100 | 显示更多 |