[ROCm] Use hipCUB/rocPRIM scan algorithms for large index support (#68487)
Summary:
For inclusive_scan and exclusive_scan, use hipCUB/rocPRIM scan algorithms for large index support.
Implemented for ROCm 5.0 and above.
Code reference : ROCmSoftwarePlatform/rocPRIM@5673df4#diff-47f4ef75e5af60dd5fe3906df9cf971f0635602a6b64a706dee6633d6677ef1a
Signed-off-by: Jagadish Krishnamoorthy <jagdish.krishna@gmail.com>
cc jeffdaily sunway513 jithunnair-amd ROCmSupport KyleCZH
Pull Request resolved: https://github.com/pytorch/pytorch/pull/68487
Reviewed By: ngimel
Differential Revision: D32547541
Pulled By: malfet
fbshipit-source-id: 4dd984e6906aec7634d05e2ceaa55e31cd4d7376
