| 54 | |
| 55 | template<typename T> |
| 56 | void sortBatched(Param<T> pVal, int dim, bool isAscending) { |
| 57 | af::dim4 inDims; |
| 58 | for (int i = 0; i < 4; i++) inDims[i] = pVal.info.dims[i]; |
| 59 | |
| 60 | // Sort dimension |
| 61 | af::dim4 tileDims(1); |
| 62 | af::dim4 seqDims = inDims; |
| 63 | tileDims[dim] = inDims[dim]; |
| 64 | seqDims[dim] = 1; |
| 65 | |
| 66 | // Create/call iota |
| 67 | Array<uint> pKey = iota<uint>(seqDims, tileDims); |
| 68 | |
| 69 | pKey.setDataDims(inDims.elements()); |
| 70 | |
| 71 | // Flat |
| 72 | pVal.info.dims[0] = inDims.elements(); |
| 73 | pVal.info.strides[0] = 1; |
| 74 | for (int i = 1; i < 4; i++) { |
| 75 | pVal.info.dims[i] = 1; |
| 76 | pVal.info.strides[i] = pVal.info.strides[i - 1] * pVal.info.dims[i - 1]; |
| 77 | } |
| 78 | |
| 79 | // Sort indices |
| 80 | auto dpl_policy = ::oneapi::dpl::execution::make_device_policy(getQueue()); |
| 81 | |
| 82 | auto key_begin = ::oneapi::dpl::begin(*pKey.get()); |
| 83 | auto key_end = key_begin + pKey.dims()[0]; |
| 84 | auto val_begin = ::oneapi::dpl::begin(*pVal.data); |
| 85 | auto val_end = val_begin + pVal.info.dims[0]; |
| 86 | auto zipped_begin = dpl::make_zip_iterator(key_begin, val_begin); |
| 87 | auto zipped_end = dpl::make_zip_iterator(key_end, val_end); |
| 88 | |
| 89 | // sort values first |
| 90 | if (isAscending) { |
| 91 | std::sort(dpl_policy, zipped_begin, zipped_end, [](auto lhs, auto rhs) { |
| 92 | return std::get<1>(lhs) < std::get<1>(rhs); |
| 93 | }); |
| 94 | } else { |
| 95 | std::sort(dpl_policy, zipped_begin, zipped_end, [](auto lhs, auto rhs) { |
| 96 | return std::get<1>(lhs) > std::get<1>(rhs); |
| 97 | }); |
| 98 | } |
| 99 | // sort according to keys second |
| 100 | std::sort(dpl_policy, zipped_begin, zipped_end, [](auto lhs, auto rhs) { |
| 101 | return std::get<0>(lhs) < std::get<0>(rhs); |
| 102 | }); |
| 103 | |
| 104 | ONEAPI_DEBUG_FINISH(getQueue()); |
| 105 | } |
| 106 | |
| 107 | template<typename T> |
| 108 | void sort0(Param<T> val, bool isAscending) { |