PR #10291: [ROCM] fixing pmap failure on multiple gpus (forward convolution)
PR #10291: [ROCM] fixing pmap failure on multiple gpus (forward convolution)

Imported from GitHub PR https://github.com/openxla/xla/pull/10291

In this PR I fix [this issue](https://github.com/google/jax/issues/14582) for AMD GPUs.

The problem was that, in order to use MIOpen convolutions, previously one needs to call **GetMIOpenConvolveAlgorithms** to populate an intenal MIOpen cache of convolution invokers. So far this was implicitly done by conv_algorithm_picker.cc for **one stream executor** (single GPU). But when one tries to run it on multiple GPU devices, the above problem occurs.\

Note that, the fix applies to AMD GPUs only, hence no changes in behaviuor on CUDA side.

@xla-rotation: could you please take a look ?
Copybara import of the project:

--
b0aa0850a7b9e370bf23ed8b086ad63d72ec88c6 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

fixing pmap failure on multiple gpus (forward convolution)

--
f4cb4633cb293c8d99ddd650cfc91dc539869e0e by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

added missing dependency

Merging this change closes #10291

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10291 from ROCm:ci_jax_pmap_failure_fix f4cb4633cb293c8d99ddd650cfc91dc539869e0e
