191 lines
3.9 KiB
Common Lisp
191 lines
3.9 KiB
Common Lisp
|
||
#ifndef DIM
|
||
#define DIM 128
|
||
#endif
|
||
|
||
#ifndef TOPK
|
||
#define TOPK 32
|
||
#endif
|
||
|
||
#ifndef QUERY_COUNT
|
||
#define QUERY_COUNT 1
|
||
#endif
|
||
|
||
/*
|
||
* Numero di work-item che collaborano per ogni query.
|
||
*/
|
||
#ifndef LANES_PER_QUERY
|
||
#define LANES_PER_QUERY 8
|
||
#endif
|
||
|
||
/*
|
||
* Ogni work-group contiene 8 work-item.
|
||
*
|
||
* Ogni work-item calcola un TOPK locale.
|
||
* Poi il work-item 0 fonde gli 8 TOPK locali.
|
||
*/
|
||
|
||
__kernel void knn(
|
||
__global const float *db,
|
||
__global const float *queries,
|
||
__global uint *indices,
|
||
__global float *values,
|
||
uint n)
|
||
{
|
||
uint lane = get_local_id(0);
|
||
uint q = get_group_id(0);
|
||
|
||
if (q >= QUERY_COUNT)
|
||
return;
|
||
|
||
/*
|
||
* TOPK locale per ogni work-item.
|
||
*/
|
||
float best_score[TOPK];
|
||
uint best_index[TOPK];
|
||
|
||
for (uint k = 0; k < TOPK; ++k) {
|
||
best_score[k] = -3.402823466e+38f;
|
||
best_index[k] = 0xffffffff;
|
||
}
|
||
|
||
__global const float *query =
|
||
&queries[q * DIM];
|
||
|
||
/*
|
||
* Ogni lane analizza una parte del database.
|
||
*
|
||
* Esempio con 8 lane:
|
||
*
|
||
* lane 0 -> 0, 8, 16, 24, ...
|
||
* lane 1 -> 1, 9, 17, 25, ...
|
||
* ...
|
||
* lane 7 -> 7, 15, 23, 31, ...
|
||
*/
|
||
for (uint i = lane; i < n; i += LANES_PER_QUERY) {
|
||
|
||
__global const float *v =
|
||
&db[i * DIM];
|
||
|
||
float s = 0.0f;
|
||
|
||
for (uint j = 0; j < DIM; ++j)
|
||
s += v[j] * query[j];
|
||
|
||
if (s <= best_score[TOPK - 1])
|
||
continue;
|
||
|
||
int pos = TOPK - 1;
|
||
|
||
while (pos > 0 &&
|
||
s > best_score[pos - 1]) {
|
||
--pos;
|
||
}
|
||
|
||
for (int k = TOPK - 1; k > pos; --k) {
|
||
best_score[k] =
|
||
best_score[k - 1];
|
||
|
||
best_index[k] =
|
||
best_index[k - 1];
|
||
}
|
||
|
||
best_score[pos] = s;
|
||
best_index[pos] = i;
|
||
}
|
||
|
||
/*
|
||
* TOPK locali nella local memory.
|
||
*
|
||
* 8 lane × 32 elementi.
|
||
*/
|
||
__local float local_scores[LANES_PER_QUERY][TOPK];
|
||
__local uint local_indices[LANES_PER_QUERY][TOPK];
|
||
|
||
for (uint k = 0; k < TOPK; ++k) {
|
||
local_scores[lane][k] = best_score[k];
|
||
local_indices[lane][k] = best_index[k];
|
||
}
|
||
|
||
/*
|
||
* Tutti i lane devono aver scritto
|
||
* la local memory prima della riduzione.
|
||
*/
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
/*
|
||
* Solo lane 0 esegue il merge finale.
|
||
*/
|
||
if (lane == 0) {
|
||
|
||
float final_score[TOPK];
|
||
uint final_index[TOPK];
|
||
|
||
for (uint k = 0; k < TOPK; ++k) {
|
||
final_score[k] =
|
||
-3.402823466e+38f;
|
||
|
||
final_index[k] =
|
||
0xffffffff;
|
||
}
|
||
|
||
/*
|
||
* Fonde gli 8 TOPK locali.
|
||
*
|
||
* Al massimo 8 × 32 = 256 candidati.
|
||
*/
|
||
for (uint l = 0;
|
||
l < LANES_PER_QUERY;
|
||
++l) {
|
||
|
||
for (uint k = 0;
|
||
k < TOPK;
|
||
++k) {
|
||
|
||
float s =
|
||
local_scores[l][k];
|
||
|
||
uint idx =
|
||
local_indices[l][k];
|
||
|
||
if (idx == 0xffffffff)
|
||
continue;
|
||
|
||
if (s <= final_score[TOPK - 1])
|
||
continue;
|
||
|
||
int pos = TOPK - 1;
|
||
|
||
while (pos > 0 &&
|
||
s > final_score[pos - 1]) {
|
||
--pos;
|
||
}
|
||
|
||
for (int p = TOPK - 1;
|
||
p > pos;
|
||
--p) {
|
||
|
||
final_score[p] =
|
||
final_score[p - 1];
|
||
|
||
final_index[p] =
|
||
final_index[p - 1];
|
||
}
|
||
|
||
final_score[pos] = s;
|
||
final_index[pos] = idx;
|
||
}
|
||
}
|
||
|
||
/*
|
||
* Risultato finale della query.
|
||
*/
|
||
for (uint k = 0; k < TOPK; ++k) {
|
||
indices[q * TOPK + k] =
|
||
final_index[k];
|
||
|
||
values[q * TOPK + k] =
|
||
final_score[k];
|
||
}
|
||
}
|
||
}
|