| 28 | namespace cpu { |
| 29 | template<typename T> |
| 30 | void topk(Array<T>& vals, Array<unsigned>& idxs, const Array<T>& in, |
| 31 | const int k, const int dim, const af::topkFunction order) { |
| 32 | // The out_dims is of size k along the dimension of the topk operation |
| 33 | // and the same as the input dimension otherwise. |
| 34 | dim4 out_dims(1); |
| 35 | int ndims = in.dims().ndims(); |
| 36 | for (int i = 0; i < ndims; i++) { |
| 37 | if (i == dim) { |
| 38 | out_dims[i] = min(k, static_cast<int>(in.dims()[i])); |
| 39 | } else { |
| 40 | out_dims[i] = in.dims()[i]; |
| 41 | } |
| 42 | } |
| 43 | |
| 44 | auto values = createEmptyArray<T>(out_dims); |
| 45 | auto indices = createEmptyArray<unsigned>(out_dims); |
| 46 | |
| 47 | auto func = [=](Param<T> values, Param<unsigned> indices, CParam<T> in) { |
| 48 | const T* ptr = in.get(); |
| 49 | unsigned* iptr = indices.get(); |
| 50 | T* vptr = values.get(); |
| 51 | |
| 52 | // Create a linear index |
| 53 | vector<uint> idx(in.dims().elements()); |
| 54 | iota(begin(idx), end(idx), 0); |
| 55 | |
| 56 | int iter = in.dims()[1] * in.dims()[2] * in.dims()[3]; |
| 57 | for (int i = 0; i < iter; i++) { |
| 58 | auto idx_itr = begin(idx) + i * in.strides()[1]; |
| 59 | auto* kiptr = iptr + k * i; |
| 60 | |
| 61 | if (order & AF_TOPK_MIN) { |
| 62 | if (order & AF_TOPK_STABLE) { |
| 63 | partial_sort_copy( |
| 64 | idx_itr, idx_itr + in.strides()[1], kiptr, kiptr + k, |
| 65 | [ptr](const uint lhs, const uint rhs) -> bool { |
| 66 | return compute_t<T>(ptr[lhs]) < |
| 67 | compute_t<T>(ptr[rhs]) |
| 68 | ? true |
| 69 | : compute_t<T>(ptr[lhs]) == |
| 70 | compute_t<T>(ptr[rhs]) |
| 71 | ? (lhs < rhs) |
| 72 | : false; |
| 73 | }); |
| 74 | } else { |
| 75 | partial_sort_copy( |
| 76 | idx_itr, idx_itr + in.strides()[1], kiptr, kiptr + k, |
| 77 | [ptr](const uint lhs, const uint rhs) -> bool { |
| 78 | return compute_t<T>(ptr[lhs]) < |
| 79 | compute_t<T>(ptr[rhs]); |
| 80 | }); |
| 81 | // Sort the top k values in each column |
| 82 | } |
| 83 | } else { |
| 84 | if (order & AF_TOPK_STABLE) { |
| 85 | partial_sort_copy( |
| 86 | idx_itr, idx_itr + in.strides()[1], kiptr, kiptr + k, |
| 87 | [ptr](const uint lhs, const uint rhs) -> bool { |