Function sort_by_key

src/backend/oneapi/sort_by_key.cpp:22–53

namespace oneapi {
template<typename Tk, typename Tv>
void sort_by_key(Array<Tk> &okey, Array<Tv> &oval, const Array<Tk> &ikey,
                 const Array<Tv> &ival, const unsigned dim, bool isAscending) {
    okey = copyArray<Tk>(ikey);
    oval = copyArray<Tv>(ival);

    switch (dim) {
        case 0: kernel::sort0ByKey<Tk, Tv>(okey, oval, isAscending); break;
        case 1:
        case 2:
        case 3:
            kernel::sortByKeyBatched<Tk, Tv>(okey, oval, dim, isAscending);
            break;
        default: AF_ERROR("Not Supported", AF_ERR_NOT_SUPPORTED);
    }

    if (dim != 0) {
        af::dim4 preorderDims = okey.dims();
        af::dim4 reorderDims(0, 1, 2, 3);
        reorderDims[dim] = 0;
        preorderDims[0]  = okey.dims()[dim];
        for (int i = 1; i <= (int)dim; i++) {
            reorderDims[i - 1] = i;
            preorderDims[i]    = okey.dims()[i - 1];
        }

        okey.setDataDims(preorderDims);
        oval.setDataDims(preorderDims);

        okey = reorder<Tk>(okey, reorderDims);
        oval = reorder<Tv>(oval, reorderDims);
    }
}

#define INSTANTIATE(Tk, Tv)                                        \
    template void sort_by_key<Tk, Tv>(                             \