Adds ability to return multiple values from nearest neighbor

* Returns multiple values from nearest neighbor using topk (removes obsolete code/kernels)
* Adds assertions and tests for maximum number of neighbors, including tests all nearest-neighbors API parameters
* Updates documentation to reflect new behavior
* Fixes out of bounds shared memory indexing
* Fixes gtest linking issues

include/af/vision.h

@@ -223,8 +223,8 @@ AFAPI void hammingMatcher(array& idx, array& dist,
    \param[in]  train is the array containing the data used as training data
    \param[in]  dist_dim indicates the dimension to analyze for distance (the dimension
                indicated here must be of equal length for both query and train arrays)
-   \param[in]  n_dist is the number of smallest distances to return (currently, only 1
-               is supported)
+   \param[in]  n_dist is the number of smallest distances to return (currently only 
+               values <= 256 are supported)
    \param[in]  dist_type is the distance computation type. Currently \ref AF_SAD (sum
                of absolute differences), \ref AF_SSD (sum of squared differences), and
                \ref AF_SHD (hamming distances) are supported.
@@ -535,8 +535,8 @@ extern "C" {
         \param[in]  train is the array containing the data used as training data
         \param[in]  dist_dim indicates the dimension to analyze for distance (the dimension
                     indicated here must be of equal length for both query and train arrays)
-        \param[in]  n_dist is the number of smallest distances to return (currently, only 1
-                    is supported)
+        \param[in]  n_dist is the number of smallest distances to return (currently, only
+                    values <= 256 are supported)
         \param[in]  dist_type is the distance computation type. Currently \ref AF_SAD (sum
                     of absolute differences), \ref AF_SSD (sum of squared differences), and
                     \ref AF_SHD (hamming distances) are supported.

src/api/c/nearest_neighbour.cpp

@@ -54,6 +54,7 @@ af_err af_nearest_neighbour(af_array* idx, af_array* dist,
         DIM_ASSERT(3, tDims[2] == 1 && tDims[3] == 1);
         DIM_ASSERT(4, (dist_dim == 0 || dist_dim == 1));
         DIM_ASSERT(5, n_dist > 0 && n_dist <= (uint)tDims[train_samples]);
+        ARG_ASSERT(5, n_dist > 0 && n_dist <= 256);
         ARG_ASSERT(6, dist_type == AF_SAD || dist_type == AF_SSD || dist_type == AF_SHD);
         TYPE_ASSERT(qType == tType);
 

src/backend/cpu/kernel/nearest_neighbour.hpp

@@ -86,7 +86,7 @@ struct dist_op<ushort, To, AF_SHD>
 };
 
 template<typename T, typename To, af_match_type dist_type>
-void nearest_neighbour(Param<uint> idx, Param<To> dist,
+void nearest_neighbour(Param<To> dists,
                        CParam<T> query, CParam<T> train,
                        const uint dist_dim, const uint n_dist)
 {
@@ -100,8 +100,7 @@ void nearest_neighbour(Param<uint> idx, Param<To> dist,
 
     const T* qPtr = query.get();
     const T* tPtr = train.get();
-    uint* iPtr = idx.get();
-    To* dPtr = dist.get();
+    To* dPtr = dists.get();
 
     dist_op<T, To, dist_type> op;
 
@@ -125,16 +124,8 @@ void nearest_neighbour(Param<uint> idx, Param<To> dist,
                 local_dist += op(qPtr[qIdx], tPtr[tIdx]);
             }
 
-            if (local_dist < best_dist) {
-                best_dist = local_dist;
-                best_idx  = j;
-            }
+            dPtr[i*nTrain + j] = local_dist;
         }
-
-        size_t oIdx;
-        oIdx = i;
-        iPtr[oIdx] = best_idx;
-        dPtr[oIdx] = best_dist;
     }
 }
 

src/backend/cpu/nearest_neighbour.cpp

@@ -12,6 +12,7 @@
 #include <math.hpp>
 #include <platform.hpp>
 #include <queue.hpp>
+#include <topk.hpp>
 #include <err_cpu.hpp>
 #include <kernel/nearest_neighbour.hpp>
 
@@ -26,35 +27,38 @@ void nearest_neighbour(Array<uint>& idx, Array<To>& dist,
                        const uint dist_dim, const uint n_dist,
                        const af_match_type dist_type)
 {
-    if (n_dist > 1) {
-        CPU_NOT_SUPPORTED("\nNumber of smallest distances can't be <1\n");
-    }
-
     idx.eval();
     dist.eval();
     query.eval();
     train.eval();
 
     uint sample_dim  = (dist_dim == 0) ? 1 : 0;
     const dim4 qDims = query.dims();
-    const dim4 outDims(n_dist, qDims[sample_dim]);
+    const dim4 tDims = train.dims();
+    const dim4 outDims (n_dist, qDims[sample_dim]);
+    const dim4 distDims(tDims[sample_dim], qDims[sample_dim]);
+
+    Array<To> tmp_dists = createEmptyArray<To>(distDims);
 
     idx  = createEmptyArray<uint>(outDims);
     dist = createEmptyArray<To  >(outDims);
 
     switch(dist_type) {
         case AF_SAD:
-            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SAD>, idx, dist, query, train, dist_dim, n_dist);
+            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SAD>, tmp_dists, query, train, dist_dim, n_dist);
             break;
         case AF_SSD:
-            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SSD>, idx, dist, query, train, dist_dim, n_dist);
+            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SSD>, tmp_dists, query, train, dist_dim, n_dist);
             break;
         case AF_SHD:
-            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SHD>, idx, dist, query, train, dist_dim, n_dist);
+            getQueue().enqueue(kernel::nearest_neighbour<T, To, AF_SHD>, tmp_dists, query, train, dist_dim, n_dist);
             break;
         default:
             AF_ERROR("Unsupported dist_type", AF_ERR_NOT_CONFIGURED);
     }
+
+    cpu::topk(dist, idx, tmp_dists, n_dist, 0, AF_TOPK_MIN);
+
 }
 
 #define INSTANTIATE(T, To)                                                              \

src/backend/cpu/topk.cpp

@@ -95,4 +95,6 @@ INSTANTIATE(float )
 INSTANTIATE(double)
 INSTANTIATE(int   )
 INSTANTIATE(uint  )
+INSTANTIATE(long long)
+INSTANTIATE(unsigned long long)
 }