| 53 | |
| 54 | template<typename T> |
| 55 | void topk(Array<T>& vals, Array<unsigned>& idxs, const Array<T>& in, |
| 56 | const int k, const int dim, const af::topkFunction order) { |
| 57 | if (getDeviceType() == CL_DEVICE_TYPE_CPU) { |
| 58 | // This branch optimizes for CPU devices by first mapping the buffer |
| 59 | // and calling partial sort on the buffer |
| 60 | |
| 61 | // TODO(umar): implement this in the kernel namespace |
| 62 | |
| 63 | // The out_dims is of size k along the dimension of the topk operation |
| 64 | // and the same as the input dimension otherwise. |
| 65 | dim4 out_dims(1); |
| 66 | int ndims = in.dims().ndims(); |
| 67 | for (int i = 0; i < ndims; i++) { |
| 68 | if (i == dim) { |
| 69 | out_dims[i] = min(k, (int)in.dims()[i]); |
| 70 | } else { |
| 71 | out_dims[i] = in.dims()[i]; |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | auto values = createEmptyArray<T>(out_dims); |
| 76 | auto indices = createEmptyArray<unsigned>(out_dims); |
| 77 | const Buffer* in_buf = in.get(); |
| 78 | Buffer* ibuf = indices.get(); |
| 79 | Buffer* vbuf = values.get(); |
| 80 | |
| 81 | cl::Event ev_in, ev_val, ev_ind; |
| 82 | |
| 83 | T* ptr = static_cast<T*>(getQueue().enqueueMapBuffer( |
| 84 | *in_buf, CL_FALSE, CL_MAP_READ, 0, in.elements() * sizeof(T), |
| 85 | nullptr, &ev_in)); |
| 86 | uint* iptr = static_cast<uint*>(getQueue().enqueueMapBuffer( |
| 87 | *ibuf, CL_FALSE, CL_MAP_READ | CL_MAP_WRITE, 0, k * sizeof(uint), |
| 88 | nullptr, &ev_ind)); |
| 89 | T* vptr = static_cast<T*>(getQueue().enqueueMapBuffer( |
| 90 | *vbuf, CL_FALSE, CL_MAP_WRITE, 0, k * sizeof(T), nullptr, &ev_val)); |
| 91 | |
| 92 | vector<uint> idx(in.elements()); |
| 93 | |
| 94 | // Create a linear index |
| 95 | iota(begin(idx), end(idx), 0); |
| 96 | cl::Event::waitForEvents({ev_in, ev_ind}); |
| 97 | |
| 98 | int iter = in.dims()[1] * in.dims()[2] * in.dims()[3]; |
| 99 | for (int i = 0; i < iter; i++) { |
| 100 | auto idx_itr = begin(idx) + i * in.strides()[1]; |
| 101 | auto kiptr = iptr + k * i; |
| 102 | |
| 103 | if (order & AF_TOPK_MIN) { |
| 104 | if (order & AF_TOPK_STABLE) { |
| 105 | partial_sort_copy( |
| 106 | idx_itr, idx_itr + in.strides()[1], kiptr, kiptr + k, |
| 107 | [ptr](const uint lhs, const uint rhs) -> bool { |
| 108 | return (compute_t<T>(ptr[lhs]) < |
| 109 | compute_t<T>(ptr[rhs])) |
| 110 | ? true |
| 111 | : compute_t<T>(ptr[lhs]) == |
| 112 | compute_t<T>(ptr[rhs]) |
no test coverage detected