| 24 | namespace opencl { |
| 25 | template<typename T> |
| 26 | void sort_index(Array<T> &okey, Array<uint> &oval, const Array<T> &in, |
| 27 | const uint dim, bool isAscending) { |
| 28 | |
| 29 | // TODO: fix half implementation of sort0bykey to support this |
| 30 | if (std::is_same_v<T, half>) { |
| 31 | OPENCL_NOT_SUPPORTED("sort_index with half"); |
| 32 | } |
| 33 | |
| 34 | try { |
| 35 | // okey contains values, oval contains indices |
| 36 | okey = copyArray<T>(in); |
| 37 | oval = range<uint>(in.dims(), dim); |
| 38 | oval.eval(); |
| 39 | |
| 40 | switch (dim) { |
| 41 | case 0: kernel::sort0ByKey<T, uint>(okey, oval, isAscending); break; |
| 42 | case 1: |
| 43 | case 2: |
| 44 | case 3: |
| 45 | kernel::sortByKeyBatched<T, uint>(okey, oval, dim, isAscending); |
| 46 | break; |
| 47 | default: AF_ERROR("Not Supported", AF_ERR_NOT_SUPPORTED); |
| 48 | } |
| 49 | |
| 50 | if (dim != 0) { |
| 51 | af::dim4 preorderDims = okey.dims(); |
| 52 | af::dim4 reorderDims(0, 1, 2, 3); |
| 53 | reorderDims[dim] = 0; |
| 54 | preorderDims[0] = okey.dims()[dim]; |
| 55 | for (uint i = 1; i <= dim; i++) { |
| 56 | reorderDims[i - 1] = i; |
| 57 | preorderDims[i] = okey.dims()[i - 1]; |
| 58 | } |
| 59 | |
| 60 | okey.setDataDims(preorderDims); |
| 61 | oval.setDataDims(preorderDims); |
| 62 | |
| 63 | okey = reorder<T>(okey, reorderDims); |
| 64 | oval = reorder<uint>(oval, reorderDims); |
| 65 | } |
| 66 | } catch (const std::exception &ex) { AF_ERROR(ex.what(), AF_ERR_INTERNAL); } |
| 67 | } |
| 68 | |
| 69 | #define INSTANTIATE(T) \ |
| 70 | template void sort_index<T>(Array<T> & val, Array<uint> & idx, \ |
no test coverage detected