[PHI] Fixed argsort big tensor bug

[Enigmatisms]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

修复了大 Tensor 情况下的。argsort 算子出错。argsort 算子出错原因：
cub::DeviceSegmentedRadixSort 目前不支持超过 INT_MAX 个元素的sort，所以对于大tensor，需要一个kernel call拆分机制，将一次 sort 拆分为多次。比如下面这个例子：

paddle.argsort(Tensor([228170138, 10],"int64"), axis=1, stable=True, )

由于 Tensor numel 为 2281701380 大于 INT_MAX，直接用 DeviceSegmentedRadixSort::SortPairs 计算 temp buffer 大小将会出现溢出。目前的解决方案：每次参与 sort 的元素不超过 2^30 个（不选择 2^31 - 1 是为了避免过大的 temp buffer，显存压力大），按照 sorting 维度计算 batch，拆分为多个 cub kernel。比如上述例子将会被拆分为三次计算：

• [2^30, 10] radix sort 第一次
• [2^30, 10] radix sort 第二次
• [2281701380 - 2^31, 10] 第一次

关于 cub::DeviceSegmentedRadixSort 的 API 支持范围，见：

• Nvidia/CCCL: [EPIC]: CUB large input support
• NVidia/CCCL Issue #3133: Add support for large num_items to device_segmented_radix_sort.cuh. 实测CCCL开发者推荐的 DeviceSegmentedSort 貌似性能不如 DeviceSegmentedRadixSort（慢数倍）。

经过性能测试，此 kernel 目前平均运行时间约为 torch 的 1.25倍，部分较好的 shape 下（比如[5, 456340276] axis = -1）可达到将近50%的计算时间。

Pcard-89620

[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.