| 85 | |
| 86 | template<typename Tk, typename Tv> |
| 87 | void sortByKeyBatched(Param<Tk> pKey, Param<Tv> pVal, const int dim, |
| 88 | bool isAscending) { |
| 89 | af::dim4 inDims; |
| 90 | for (int i = 0; i < 4; i++) inDims[i] = pKey.info.dims[i]; |
| 91 | |
| 92 | const dim_t elements = inDims.elements(); |
| 93 | |
| 94 | // Sort dimension |
| 95 | // tileDims * seqDims = inDims |
| 96 | af::dim4 tileDims(1); |
| 97 | af::dim4 seqDims = inDims; |
| 98 | tileDims[dim] = inDims[dim]; |
| 99 | seqDims[dim] = 1; |
| 100 | |
| 101 | // Create/call iota |
| 102 | Array<uint> Seq = iota<uint>(seqDims, tileDims); |
| 103 | |
| 104 | auto dpl_policy = ::oneapi::dpl::execution::make_device_policy(getQueue()); |
| 105 | |
| 106 | // set up iterators for seq, key, val, and new cKey |
| 107 | auto seq_begin = ::oneapi::dpl::begin(*Seq.get()); |
| 108 | auto seq_end = seq_begin + elements; |
| 109 | auto key_begin = |
| 110 | ::oneapi::dpl::begin(pKey.data->template reinterpret<compute_t<Tk>>()); |
| 111 | auto key_end = key_begin + elements; |
| 112 | |
| 113 | auto val_begin = ::oneapi::dpl::begin(*pVal.data); |
| 114 | auto val_end = val_begin + elements; |
| 115 | |
| 116 | auto cKey = memAlloc<Tk>(elements); |
| 117 | auto cKey_get = cKey.get(); |
| 118 | getQueue().submit([&](sycl::handler &h) { |
| 119 | h.copy(pKey.data->template reinterpret<compute_t<Tk>>().get_access( |
| 120 | h, elements), |
| 121 | cKey_get->template reinterpret<compute_t<Tk>>().get_access( |
| 122 | h, elements)); |
| 123 | }); |
| 124 | auto ckey_begin = |
| 125 | ::oneapi::dpl::begin(cKey.get()->template reinterpret<compute_t<Tk>>()); |
| 126 | auto ckey_end = ckey_begin + elements; |
| 127 | |
| 128 | { |
| 129 | auto zipped_begin_KV = dpl::make_zip_iterator(key_begin, val_begin); |
| 130 | auto zipped_end_KV = dpl::make_zip_iterator(key_end, val_end); |
| 131 | auto zipped_begin_cKS = dpl::make_zip_iterator(ckey_begin, seq_begin); |
| 132 | auto zipped_end_cKS = dpl::make_zip_iterator(ckey_end, seq_end); |
| 133 | if (isAscending) { |
| 134 | std::sort(dpl_policy, zipped_begin_KV, zipped_end_KV, |
| 135 | [](auto lhs, auto rhs) { |
| 136 | return std::get<0>(lhs) < std::get<0>(rhs); |
| 137 | }); |
| 138 | std::sort(dpl_policy, zipped_begin_cKS, zipped_end_cKS, |
| 139 | [](auto lhs, auto rhs) { |
| 140 | return std::get<0>(lhs) < std::get<0>(rhs); |
| 141 | }); |
| 142 | } else { |
| 143 | std::sort(dpl_policy, zipped_begin_KV, zipped_end_KV, |
| 144 | [](auto lhs, auto rhs) { |