[rocm-jaxlib-v0.9.0] Enable postsubmit benchmarks in CI

[mmakevic-amd]

Motivation

Backporting changes from #691 and #730.

Edit: Also enabled RBE in this PR, this part will need to be cherry-picked to v0.9.1

Test Plan

I will manually trigger CI check before merging

Test Result

Workflow run successful: https://github.com/ROCm/xla/actions/runs/24036556155/job/70097673616

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[i-chaochen]

use this docker image #754

[mmakevic-amd]

Not sure what would be the cleanest way to do it here. Maybe something like:

###### .github/workflows/generate_benchmark_matrix.yml
- name: Run generate_benchmark_matrices_main
  env:
    ROCM_DOCKER_IMAGE: ${{ vars.DOCKER_IMAGE }}
  run: 
   ...

and then:

##### xla/tools/benchmarks/utils/generate_benchmark_matrices.cc
{"GPU_MI250",
           []() -> std::string {
             const char* img = std::getenv("ROCM_DOCKER_IMAGE");
             return img;
           }()
}

?

[i-chaochen]

ha, I c...this is XLA code level, I guess it's best to keep what it has, and don't utilise CI config, otherwise it affects the local use. You can use our 7.2.1 docker image rocm/tensorflow-build@sha256:66eb4c1e39db76fae2eb0a1029490acbe7bfce0e00d6ab435e170f743921f4c4

[mmakevic-amd]

Same as above, I updated the image to rocm/tensorflow-build@sha256:66eb4c1e39db76fae2eb0a1029490acbe7bfce0e00d6ab435e170f743921f4c4