-/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

                                        memory::Allocator::kScratchpad);

    // Run sorting operation
    cub::DeviceRadixSort::SortPairsDescending<int, int>(
        d_temp_storage->ptr(), temp_storage_bytes, target_lvls_data, keys_out,
        idx_in, idx_out, roi_num);

    int* restore_idx_data =
        restore_index->mutable_data<int>({roi_num, 1}, dev_ctx.GetPlace());

    cub::DeviceRadixSort::SortPairsDescending<int, int>(
        d_temp_storage->ptr(), temp_storage_bytes, idx_out, keys_out, idx_in,
        restore_idx_data, roi_num);

                  kNumMaxinumNumBlocks);
}

static inline void TransLoD(const int* length_lod, const int lod_size,
                            int* offset_lod) {
  int offset = 0;
  for (int i = 0; i < lod_size; ++i) {
    offset_lod[i] = offset;
