PR #39744: [ROCm] Porting CUB sort FFI handler consolidation to ROCm. - #113864

Imported from GitHub PR openxla/xla#39744

📝 Summary of Changes
The CUB sort FFI handler was recently refactored and consolidated in two commits:
- openxla/xla@878c525
- openxla/xla@ce4b95a

In order to allow ROCm devices to continue using the CUB sort function, these same changes were ported to the relevant ROCm files.

Changes to files in `xla/stream_executor/rocm/`
- The build file is restructured so that it no longer compiles `cub_sort_kernel_rocm.cc` for each type.
- The various helper functions from commit openxla/xla@878c525 have been ported here.
- The missing S32_B16 and S32_B64 template instantiations have been added.

🎯 Justification
This change was breaking unit tests in XLA and JAX on ROCm:
- XLA:
  - xla/gpu/service/tests/sorting_test.cc

- JAX:
  - tests/ann_test.py::AnnTest::test_approx_max_k0
  - tests/ann_test.py::AnnTest::test_approx_max_k5
  - tests/ann_test.py::AnnTest::test_approx_max_k8
  - tests/ann_test.py::AnnTest::test_approx_min_k0
  - tests/ann_test.py::AnnTest::test_approx_min_k5
  - tests/ann_test.py::AnnTest::test_approx_min_k8
  - tests/scipy_signal_test.py::LaxBackedScipySignalTests::testWelchAgainstNumpy8

🚀 Kind of Contribution
🐛 Bug Fix

🧪 Unit Tests:
The various unit tests are now passing with the changes in the PR.
- XLA (`sorting_test.cc`):
```
================================================================================
INFO: Found 1 test target...
Target //xla/service/gpu/tests:sorting_test_amdgpu_any up-to-date:
  bazel-bin/xla/service/gpu/tests/sorting_test_amdgpu_any
INFO: Elapsed time: 192.140s, Critical Path: 153.72s
INFO: 5136 processes: 650 internal, 4486 local.
INFO: Build completed successfully, 5136 total actions
//xla/service/gpu/tests:sorting_test_amdgpu_any                          PASSED in 31.8s
  Stats over 15 runs: max = 31.8s, min = 22.3s, avg = 29.9s, dev = 2.7s

Executed 1 out of 1 test: 1 test passes.
```
- JAX (`tests/ann_test.py::AnnTest::test_approx_*`):
```
==================================================================== test session starts ====================================================================
platform linux -- Python 3.12.3, pytest-8.4.2, pluggy-1.6.0 -- /usr/bin/python3
cachedir: .pytest_cache
hypothesis profile 'default'
rootdir: /workspace/rocm-jax/jax
configfile: pyproject.toml
plugins: xdist-3.8.0, hypothesis-6.142.1
collected 42 items / 32 deselected / 10 selected

tests/ann_test.py::AnnTest::test_approx_max_k0 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k1 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k2 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k3 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k4 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k5 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k6 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k7 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k8 PASSED
tests/ann_test.py::AnnTest::test_approx_max_k9 PASSED

============================================================= 10 passed, 32 deselected in 9.25s =============================================================
```
- JAX (`tests/scipy_signal_test.py::LaxBackedScipySignalTests::testWelchAgainstNumpy8`):
```
==================================================================== test session starts ====================================================================
platform linux -- Python 3.12.3, pytest-8.4.2, pluggy-1.6.0 -- /usr/bin/python3
cachedir: .pytest_cache
hypothesis profile 'default'
rootdir: /workspace/rocm-jax/jax
configfile: pyproject.toml
plugins: xdist-3.8.0, hypothesis-6.142.1
collected 100 items / 99 deselected / 1 selected

tests/scipy_signal_test.py::LaxBackedScipySignalTests::testWelchAgainstNumpy8 PASSED

============================================================= 1 passed, 99 deselected in 8.16s ==============================================================
```
Copybara import of the project:

--
6b432c264245514680969c508d9a0c510b15aed5 by tsrw2048 <[email protected]>:

[ROCm] Porting CUB sort FFI handler consolidation to ROCm.

The CUB sort FFI handler was refactored and consolidated
in two commits:
 - 878c52525f1a72856e5c5fd42775211f19266568
 - ce4b95a3668977155ec4914b1b28f2d5bf040ef2
In order to allow ROCm devices to continue using the CUB sort
function, these same changes were ported to the relevant ROCm
files.

Changes to files in xla/stream_executor/rocm/
- The build file is restructured so that it no longer compiles
cub_sort_kernel_rocm.cc for each type.
- The various helper functions from commit 878c525 have
been ported here.
- The missing S32_B16 and S32_B64 template instantiations have
been added.

Merging this change closes #39744

PiperOrigin-RevId: 892285071