63 lines
1.5 KiB
Text
63 lines
1.5 KiB
Text
#version 450
|
|
|
|
layout(local_size_x_id = 0) in;
|
|
|
|
layout(set = 0, binding = 0, std430) readonly buffer DescriptorsA {
|
|
uint descriptors_a[];
|
|
};
|
|
|
|
layout(set = 0, binding = 1, std430) readonly buffer DescriptorsB {
|
|
uint descriptors_b[];
|
|
};
|
|
|
|
layout(set = 0, binding = 2, std430) writeonly buffer Top2Output {
|
|
uvec4 top2[];
|
|
};
|
|
|
|
layout(push_constant) uniform Counts {
|
|
uint count_a;
|
|
uint count_b;
|
|
} counts;
|
|
|
|
void main() {
|
|
uint query_index = gl_GlobalInvocationID.x;
|
|
if (query_index >= counts.count_a) {
|
|
return;
|
|
}
|
|
|
|
uint best_index = 0xffffffffu;
|
|
uint best_distance = 0xffffffffu;
|
|
uint second_index = 0xffffffffu;
|
|
uint second_distance = 0xffffffffu;
|
|
|
|
for (uint train_index = 0; train_index < counts.count_b; ++train_index) {
|
|
uint distance = 0;
|
|
for (uint word = 0; word < 8; ++word) {
|
|
uint a = descriptors_a[query_index * 8 + word];
|
|
uint b = descriptors_b[train_index * 8 + word];
|
|
distance += bitCount(a ^ b);
|
|
}
|
|
|
|
// train_index increases monotonically, so keeping the first equal distance
|
|
// implements the required smallest-index tie-break without another comparison.
|
|
if (distance < best_distance) {
|
|
second_index = best_index;
|
|
second_distance = best_distance;
|
|
best_index = train_index;
|
|
best_distance = distance;
|
|
continue;
|
|
}
|
|
|
|
if (distance < second_distance) {
|
|
second_index = train_index;
|
|
second_distance = distance;
|
|
}
|
|
}
|
|
|
|
top2[query_index] = uvec4(
|
|
best_index,
|
|
best_distance,
|
|
second_index,
|
|
second_distance
|
|
);
|
|
}
|