[BIG tensor]Fix the infinite loop issue in argsort.

[Difers]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

Pcard-73145

• 问题背景：
  在以下case argsort会导致hang住：

import numpy
import paddle
input_tensor = (numpy.random.random([2, 1140850690]) - 0.5).astype("float16")
paddle_x = paddle.to_tensor(numpy_tensor)
paddle_out = paddle.argsort(paddle_x)

• 问题原因
  在 [PHI] Fixed argsort big tensor bug #72712 中实现了拆分为多个 cub kernel；但拆分计算batch时每次参与 sort 的元素不超过 2^30 个

constexpr int64_t max_elements = 1 << 30;
const int64_t segment_size = num_cols;
const int64_t element_per_call = std::min(max_elements, total_elements);
const int64_t batch_size = (element_per_call / segment_size) * segment_size; 

在被sort的维度大于2^30时，(element_per_call / segment_size)为0，导致死循环

• 报错补充
  在被sort的维度大于2^31时，调用DeviceSegmentedRadixSort::SortPairs依旧有问题，参考torch处理，对该情况添加报错信息。
  https://github.com/pytorch/pytorch/blob/aeb5321b6360c899808d3461789b3bbd6265756e/aten/src/ATen/native/cuda/Sort.cpp#L63

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[wanghuancoder]

LGTM