| 23 | |
| 24 | template<typename T, typename To> |
| 25 | void nearest_neighbour(Array<uint>& idx, Array<To>& dist, const Array<T>& query, |
| 26 | const Array<T>& train, const uint dist_dim, |
| 27 | const uint n_dist, const af_match_type dist_type) { |
| 28 | uint sample_dim = (dist_dim == 0) ? 1 : 0; |
| 29 | const dim4& qDims = query.dims(); |
| 30 | const dim4& tDims = train.dims(); |
| 31 | const dim4 outDims(n_dist, qDims[sample_dim]); |
| 32 | const dim4 distDims(tDims[sample_dim], qDims[sample_dim]); |
| 33 | |
| 34 | Array<To> tmp_dists = createEmptyArray<To>(distDims); |
| 35 | |
| 36 | idx = createEmptyArray<uint>(outDims); |
| 37 | dist = createEmptyArray<To>(outDims); |
| 38 | |
| 39 | switch (dist_type) { |
| 40 | case AF_SAD: |
| 41 | getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SAD>, |
| 42 | tmp_dists, query, train, dist_dim); |
| 43 | break; |
| 44 | case AF_SSD: |
| 45 | getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SSD>, |
| 46 | tmp_dists, query, train, dist_dim); |
| 47 | break; |
| 48 | case AF_SHD: |
| 49 | getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SHD>, |
| 50 | tmp_dists, query, train, dist_dim); |
| 51 | break; |
| 52 | default: AF_ERROR("Unsupported dist_type", AF_ERR_NOT_CONFIGURED); |
| 53 | } |
| 54 | |
| 55 | cpu::topk(dist, idx, tmp_dists, n_dist, 0, AF_TOPK_MIN); |
| 56 | } |
| 57 | |
| 58 | #define INSTANTIATE(T, To) \ |
| 59 | template void nearest_neighbour<T, To>( \ |