| 47 | |
| 48 | template<typename T> |
| 49 | void sortBatched(Param<T> pVal, int dim, bool isAscending) { |
| 50 | af::dim4 inDims; |
| 51 | for (int i = 0; i < 4; i++) inDims[i] = pVal.dims[i]; |
| 52 | |
| 53 | // Sort dimension |
| 54 | // tileDims * seqDims = inDims |
| 55 | af::dim4 tileDims(1); |
| 56 | af::dim4 seqDims = inDims; |
| 57 | tileDims[dim] = inDims[dim]; |
| 58 | seqDims[dim] = 1; |
| 59 | |
| 60 | // Create/call iota |
| 61 | Array<uint> pKey = iota<uint>(seqDims, tileDims); |
| 62 | |
| 63 | pVal = flat(pVal); |
| 64 | |
| 65 | // Sort indices |
| 66 | // sort_by_key<T, uint, isAscending>(*resVal, *resKey, val, key, 0); |
| 67 | thrustSortByKey(pVal.ptr, pKey.get(), pVal.dims[0], isAscending); |
| 68 | |
| 69 | // Needs to be ascending (true) in order to maintain the indices properly |
| 70 | thrustSortByKey(pKey.get(), pVal.ptr, pVal.dims[0], true); |
| 71 | } |
| 72 | |
| 73 | template<typename T> |
| 74 | void sort0(Param<T> val, bool isAscending) { |
nothing calls this directly
no test coverage detected