wehub-resource-sync
138 lines
3.9 KiB
Plaintext
138 lines
3.9 KiB
Plaintext
layout(set=0, binding=0) buffer outValueBuffer {
|
|
FLOAT data[];
|
|
} uOutValue;
|
|
|
|
layout(set=0, binding=1) buffer outIndexBuffer {
|
|
int data[];
|
|
} uOutIndex;
|
|
|
|
layout(set=0, binding=2) readonly buffer sourceBuffer {
|
|
FLOAT data[];
|
|
} uInput;
|
|
|
|
layout(set=0, binding=3) uniform constBuffer {
|
|
int rowSize;
|
|
int k;
|
|
int numRows;
|
|
int pad;
|
|
} uConst;
|
|
|
|
#define THREAD_NUMBER 128
|
|
#define LOCAL_K 8
|
|
#define CANDIDATE_NUMBER (THREAD_NUMBER * LOCAL_K)
|
|
|
|
shared FLOAT sharedValue[CANDIDATE_NUMBER];
|
|
shared int sharedIndex[CANDIDATE_NUMBER];
|
|
|
|
layout(local_size_x = THREAD_NUMBER, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
bool afterAsc(FLOAT aValue, int aIndex, FLOAT bValue, int bIndex) {
|
|
if (aValue > bValue) return true;
|
|
if (aValue < bValue) return false;
|
|
return aIndex > bIndex;
|
|
}
|
|
|
|
bool better(FLOAT aValue, int aIndex, FLOAT bValue, int bIndex) {
|
|
if (bIndex < 0) return true;
|
|
if (aIndex < 0) return false;
|
|
#ifdef SORT_DESC
|
|
if (aValue > bValue) return true;
|
|
if (aValue < bValue) return false;
|
|
#else
|
|
if (aValue < bValue) return true;
|
|
if (aValue > bValue) return false;
|
|
#endif
|
|
return aIndex < bIndex;
|
|
}
|
|
|
|
void main() {
|
|
int tid = int(gl_LocalInvocationID.x);
|
|
int row = int(gl_WorkGroupID.y);
|
|
|
|
if (row >= uConst.numRows) return;
|
|
|
|
int rowOffset = row * uConst.rowSize;
|
|
|
|
// Initialize per-thread top-k
|
|
#ifdef SORT_DESC
|
|
FLOAT initVal = FLOAT(-1e38);
|
|
#else
|
|
FLOAT initVal = FLOAT(1e38);
|
|
#endif
|
|
|
|
FLOAT localValue[LOCAL_K];
|
|
int localIndex[LOCAL_K];
|
|
for (int i = 0; i < LOCAL_K; ++i) {
|
|
localValue[i] = initVal;
|
|
localIndex[i] = -1;
|
|
}
|
|
|
|
// Read input and build per-thread top-k
|
|
for (int c = tid; c < uConst.rowSize; c += THREAD_NUMBER) {
|
|
FLOAT value = uInput.data[rowOffset + c];
|
|
int idx = c;
|
|
|
|
if (!better(value, idx, localValue[LOCAL_K - 1], localIndex[LOCAL_K - 1])) continue;
|
|
|
|
int insertPos = LOCAL_K;
|
|
for (int j = 0; j < LOCAL_K; ++j) {
|
|
if (better(value, idx, localValue[j], localIndex[j])) {
|
|
insertPos = j;
|
|
break;
|
|
}
|
|
}
|
|
if (insertPos >= LOCAL_K) continue;
|
|
|
|
for (int j = LOCAL_K - 1; j > insertPos; --j) {
|
|
localValue[j] = localValue[j - 1];
|
|
localIndex[j] = localIndex[j - 1];
|
|
}
|
|
localValue[insertPos] = value;
|
|
localIndex[insertPos] = idx;
|
|
}
|
|
|
|
// Scatter to shared memory
|
|
int base = tid * LOCAL_K;
|
|
for (int i = 0; i < LOCAL_K; ++i) {
|
|
sharedValue[base + i] = localValue[i];
|
|
sharedIndex[base + i] = localIndex[i];
|
|
}
|
|
barrier();
|
|
|
|
// Bitonic sort in shared memory
|
|
for (uint size = 2; size <= CANDIDATE_NUMBER; size <<= 1) {
|
|
for (uint stride = size >> 1; stride > 0; stride >>= 1) {
|
|
for (uint sIdx = uint(tid); sIdx < CANDIDATE_NUMBER; sIdx += THREAD_NUMBER) {
|
|
uint ixj = sIdx ^ stride;
|
|
if (ixj <= sIdx) continue;
|
|
|
|
bool up = ((sIdx & size) == 0u);
|
|
#ifdef SORT_DESC
|
|
up = !up;
|
|
#endif
|
|
bool after = afterAsc(sharedValue[sIdx], sharedIndex[sIdx],
|
|
sharedValue[ixj], sharedIndex[ixj]);
|
|
if (up == after) {
|
|
FLOAT tValue = sharedValue[sIdx];
|
|
sharedValue[sIdx] = sharedValue[ixj];
|
|
sharedValue[ixj] = tValue;
|
|
int tIndex = sharedIndex[sIdx];
|
|
sharedIndex[sIdx] = sharedIndex[ixj];
|
|
sharedIndex[ixj] = tIndex;
|
|
}
|
|
}
|
|
barrier();
|
|
}
|
|
}
|
|
|
|
// Thread 0 writes output
|
|
if (tid == 0) {
|
|
int realK = min(uConst.k, uConst.rowSize);
|
|
int outOffset = row * uConst.k;
|
|
for (int i = 0; i < realK; ++i) {
|
|
uOutValue.data[outOffset + i] = sharedValue[i];
|
|
uOutIndex.data[outOffset + i] = sharedIndex[i];
|
|
}
|
|
}
|
|
}
|