117 lines
3.7 KiB
Text
117 lines
3.7 KiB
Text
#version 450
|
|
|
|
#ifdef LARDON3D_SIFT_FP64
|
|
#extension GL_ARB_gpu_shader_fp64 : require
|
|
#define Metric double
|
|
#else
|
|
#define Metric float
|
|
#endif
|
|
|
|
layout(local_size_x_id = 0) in;
|
|
|
|
layout(set = 0, binding = 0, std430) readonly buffer DescriptorsA {
|
|
float descriptors_a[];
|
|
};
|
|
|
|
layout(set = 0, binding = 1, std430) readonly buffer DescriptorsB {
|
|
float descriptors_b[];
|
|
};
|
|
|
|
struct Top2 {
|
|
uint best_index;
|
|
float best_squared_distance;
|
|
uint second_index;
|
|
float second_squared_distance;
|
|
};
|
|
|
|
layout(set = 0, binding = 2, std430) writeonly buffer Top2Output {
|
|
Top2 top2[];
|
|
};
|
|
|
|
layout(push_constant) uniform Counts {
|
|
uint count_a;
|
|
uint count_b;
|
|
} counts;
|
|
|
|
shared uint local_best_index[256];
|
|
shared Metric local_best_distance[256];
|
|
shared uint local_second_index[256];
|
|
shared Metric local_second_distance[256];
|
|
|
|
bool before(Metric left_distance, uint left_index,
|
|
Metric right_distance, uint right_index) {
|
|
return left_distance < right_distance ||
|
|
(left_distance == right_distance && left_index < right_index);
|
|
}
|
|
|
|
void insert_candidate(inout uint best_index, inout Metric best_distance,
|
|
inout uint second_index, inout Metric second_distance,
|
|
uint candidate_index, Metric candidate_distance) {
|
|
if (candidate_index == 0xffffffffu) {
|
|
return;
|
|
}
|
|
if (before(candidate_distance, candidate_index, best_distance, best_index)) {
|
|
second_index = best_index;
|
|
second_distance = best_distance;
|
|
best_index = candidate_index;
|
|
best_distance = candidate_distance;
|
|
return;
|
|
}
|
|
if (candidate_index != best_index &&
|
|
before(candidate_distance, candidate_index,
|
|
second_distance, second_index)) {
|
|
second_index = candidate_index;
|
|
second_distance = candidate_distance;
|
|
}
|
|
}
|
|
|
|
void main() {
|
|
uint query_index = gl_WorkGroupID.x;
|
|
uint lane = gl_LocalInvocationID.x;
|
|
if (query_index >= counts.count_a) {
|
|
return;
|
|
}
|
|
|
|
uint best_index = 0xffffffffu;
|
|
Metric best_distance = Metric(1.0) / Metric(0.0);
|
|
uint second_index = 0xffffffffu;
|
|
Metric second_distance = Metric(1.0) / Metric(0.0);
|
|
|
|
for (uint train_index = lane; train_index < counts.count_b;
|
|
train_index += gl_WorkGroupSize.x) {
|
|
Metric distance = Metric(0.0);
|
|
for (uint component = 0; component < 128; ++component) {
|
|
Metric difference =
|
|
Metric(descriptors_a[query_index * 128 + component]) -
|
|
Metric(descriptors_b[train_index * 128 + component]);
|
|
distance += difference * difference;
|
|
}
|
|
insert_candidate(best_index, best_distance, second_index, second_distance,
|
|
train_index, distance);
|
|
}
|
|
|
|
local_best_index[lane] = best_index;
|
|
local_best_distance[lane] = best_distance;
|
|
local_second_index[lane] = second_index;
|
|
local_second_distance[lane] = second_distance;
|
|
barrier();
|
|
|
|
// A single lane merges the bounded local results in lane order. This makes
|
|
// the reduction and its tie-break independent of subgroup scheduling.
|
|
if (lane == 0) {
|
|
best_index = 0xffffffffu;
|
|
best_distance = Metric(1.0) / Metric(0.0);
|
|
second_index = 0xffffffffu;
|
|
second_distance = Metric(1.0) / Metric(0.0);
|
|
for (uint source = 0; source < gl_WorkGroupSize.x; ++source) {
|
|
insert_candidate(best_index, best_distance, second_index, second_distance,
|
|
local_best_index[source], local_best_distance[source]);
|
|
insert_candidate(best_index, best_distance, second_index, second_distance,
|
|
local_second_index[source], local_second_distance[source]);
|
|
}
|
|
top2[query_index].best_index = best_index;
|
|
top2[query_index].best_squared_distance = float(best_distance);
|
|
top2[query_index].second_index = second_index;
|
|
top2[query_index].second_squared_distance = float(second_distance);
|
|
}
|
|
}
|