PR #11880: InitializeBuffer: use maximum 2 h2d copies
PR #11880: InitializeBuffer: use maximum 2 h2d copies

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

This provides a significant speedup when autotuning, where buffers are repeatedly initialized.
Instead of repeatedly copying the same data to different locations on the device, it is copied once (using up to two host to device copies) and then replicated on the device using a custom kernel.

Using the paxml container from JAX-Toolbox and the 5B configuration on H100, the total runtime of the `gemm-algorithm-picker` pass for the main JITed function decreased from around 4.8s to 0.16s. With the same model and `--xla_gpu_triton_gemm_any=true` to make sure the Triton autotuner does a lot of work, that pass speeds up from 17.2s to 14.2s with this change.
Copybara import of the project:

--
bc9de3720516b4450d3a9791eae8a6979a3a08a9 by Olli Lupton <olupton@nvidia.com>:

InitializeBuffer: use maximum 2 h2d copies

This provides a significant speedup when autotuning, where buffers are
repeatedly re-initialized. Instead of repeatedly copying the same data
to different locations on the device, it is copied once (using up to two
host to device copies) and then replicated on the device using a custom
kernel.

--
ee92fc52d06604ab615bca9eb278faa40e22caf0 by Olli Lupton <olupton@nvidia.com>:

Do not refer to repeat_buffer_kernel::kernel() in CPU builds.

Merging this change closes #11880

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11880 from olupton:random-init-with-3-operations ee92fc52d06604ab615bca9eb278faa40e22caf0
