knn distance

knn distance

VEX_FUNCTION(float, l2_distance, (size_t, idx)(uint32_t, num_dim)(float*, query)(float*, candidates),
             float d = 0;
for (uint i = 0; i < num_dim; ++i) {
d += pow(query[i] - candidates[idx * num_dim + i], 2);
}
return d;
            );
            
Knn::ValueVec Knn::search_h(const ValueVec& query, size_t query_offset) const {
  prof_.tic_cl("distance");
  ValueVec distance(current_context(), num_row_);
  switch (distance_) {
  case L2:
    distance = l2_distance(element_index(), num_col_,
                           raw_pointer(query) + query_offset, raw_pointer(h_));
    break;
  case COSINE:
    distance = cosine_distance(element_index(), num_col_,
                               raw_pointer(query) + query_offset, raw_pointer(h_));
  }
  prof_.toc("distance");
  return distance;
}

pair<Knn::ValueVec, Knn::IndexVec> Knn::sort_internal(ValueVec& distance, uint32_t top_k) const {
  prof_.tic_cl("sort_by_key");
  IndexVec index(current_context(), distance.size());
  index = element_index();
  sort_by_key(distance, index);
  prof_.toc("sort_by_key");
  prof_.tic_cl("slice");
  slicer<1> slice(extents[top_k]);
  auto result = make_pair(slice[range(0, top_k)](distance), slice[range(0, top_k)](index));
  prof_.toc("slice");
  return result;
}

pair<Knn::ValueVec, Knn::IndexVec> Knn::search_similar(uint32_t obj, uint32_t top_k) const {
  auto distance = search_h(h_, obj * num_col_);
  return sort_internal(distance, top_k);
}

[    distance:     4658.446 sec.] ( 81.39%) (428566x; avg: 1.078919e+04 usec.)