How to do hermetic ROCm PJRT plugin builds?

[joelberkeley]

Hi, how do I pin the ROCm version in builds of the PJRT plugin? In rocm_configure.bzl, I see ROCM_DISTRO_VERSION, used as

common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a"

That option is selected from this list, which is extremely restrictive (it doesn't, for example, admit any consumer cards).

For comparison, here's how I pin the CUDA versions

  bazel build \
    --config cuda_nvcc \
    --config nonccl \
    --repo_env HERMETIC_CUDA_VERSION="$cuda_version" \
    --repo_env HERMETIC_CUDNN_VERSION="$cudnn_version" \
    --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="$cuda_compute_capabilities" \
    --repo_env HERMETIC_NCCL_VERSION="$nccl_version" \
    --repo_env HERMETIC_NVSHMEM_VERSION="$nvshmem_version" \
    //xla/pjrt/c:pjrt_c_api_gpu_plugin.so

I notice the tarballs pointed to in CI are only a small subset of available tarballs https://repo.amd.com/rocm/tarball-multi-arch/, one of which therock-dist-linux-multiarch-7.14.0.tar.gz, may be what I'm after, but I can't see an entry point in the plugin build to use this. I notice ZML patch XLA for it https://github.com/zml/pjrt-artifacts/blob/master/openxla/patches/rocm/0003-Add-TheRock-multiarch-ROCm-redist.patch

Following advice from ZML, I tried installing ROCm 7.14 first, then building the plugin. The build logs mention

Downloading and extracting https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx94X-dcgpu-7.13.0.tar.gz

and

Auto-Configuration Warning: Using hermetic ROCm from: local_config_rocm

which suggests it's building with the default 7.13 listed at the top of this post. Meanwhile the plugin, run in an environment with 7.14 installled, logs at runtime

I0804 22:56:29.339227      13 service.cc:170]   StreamExecutor [0]: AMD Radeon RX 9060 XT, AMDGPU ISA version: gfx1200 (Driver: 7.14.60850; Runtime: 7.14.60850; Toolkit: 7.13.99004; DNN: 0.0.0)

@alekstheod might you know?

[sbryngelson]

ROCM_DISTRO_VERSION only accepts the versions registered in XLA’s distro table. Since you already installed ROCm 7.14, the simplest option is to point the repository rule at that installation directly:

--repo_env=ROCM_PATH=/opt/rocm-7.14.0

Use the actual installation root containing bin/hipcc, include, and lib.

XLA checks ROCM_PATH before TF_ROCM_MULTIPLE_PATHS, ROCM_DISTRO_URL, and ROCM_DISTRO_VERSION, so this should prevent it from selecting the default 7.13 archive. The “Using hermetic ROCm from local_config_rocm” message does not identify which of those sources populated local_config_rocm.

For an RX 9060 XT, gfx1200 is the correct target. XLA should detect it locally, but you can force it if needed:

--repo_env=TF_ROCM_AMDGPU_TARGETS=gfx1200

[alekstheod]

There are older distros too! You can put whatever tarball you prefer. You can also store a tarball on any of your server to bring it back for your builds.

[joelberkeley]

I've tried using ROCM_PATH, but it's still mentioning ROCm 7.15, even though I installed 10.0.0. Here's what I'm using (Ubuntu 24.04)

WORKDIR /therock

RUN wget https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz
RUN mkdir install
RUN tar -xf *.tar.gz -C install

which gives me

root@4122dfd88771:# ls /therock/install/
.info/   .kpack/  amdgcn/  bin/     etc/     include/ lib/     libexec/ llvm/    share/
root@4122dfd88771:# cat /therock/install/.info/version 
10.0.0

I then build with

  bazel build \
    --config clang_local \
    --config rocm \
    --repo_env=ROCM_PATH=/therock/install \
    --repo_env TF_ROCM_AMDGPU_TARGETS="$tf_rocm_amdgpu_targets" \
    //xla/pjrt/c:pjrt_c_api_gpu_plugin.so

When I run the executable, I see

W1031 22:17:49.180085      20 rocm_pcie_bandwidth.cc:75] rsmi_dev_gpu_metrics_info_get failed for 0000:12:00.0: RSMI_STATUS_UNEXPECTED_DATA: Data read (usually from a file) or provided to function is not what was expected
W1031 22:17:49.180097      20 rocm_executor.cc:1151] Could not determine PCIe bandwidth for device 1 via rocm_smi. Assuming PCIe Gen4 x16.
E1031 22:17:49.260812      13 platform_util.cc:279] Failed to create stream executor for device ROCM:1: StreamExecutor ROCM device (1) is of unsupported AMDGPU version 'gfx1036' . The supported AMDGPU versions are 'gfx908, gfx90a, gfx942, gfx950, gfx1030, gfx1100, gfx1101, gfx1103, gfx1150, gfx1151, gfx1200, gfx1201, gfx1250'.
I1031 22:17:49.264783      13 service.cc:154] XLA service 0x55b8872f0d80 initialized for platform ROCM (this does not guarantee that XLA will be used). Devices:
I1031 22:17:49.264797      13 service.cc:170]   StreamExecutor [0]: AMD Radeon RX 9060 XT, AMDGPU ISA version: gfx1200 (Driver: 7.15.26333; Runtime: 7.15.26333; Toolkit: 7.15.26333; DNN: 0.0.0)
I1031 22:17:49.299266      13 se_gpu_pjrt_client.cc:1342] Using BFC allocator.
I1031 22:17:49.299293      13 gpu_helpers.cc:162] XLA backend allocating 11.94GiB (12821987328 bytes) on device 0 for BFCAllocator.
I1031 22:17:50.201933      55 amdgpu_backend.cc:163] TF_XLA_HSACO_CACHE_DIR is not set: HSACO file cache is disabled!

which includes the line

I1031 22:17:49.264797      13 service.cc:170]   StreamExecutor [0]: AMD Radeon RX 9060 XT, AMDGPU ISA version: gfx1200 (Driver: 7.15.26333; Runtime: 7.15.26333; Toolkit: 7.15.26333; DNN: 0.0.0)

[sbryngelson]

@joelberkeley That’s expected: ROCm 10.0.0 ships HIP 7.15.26333, and XLA prints the HIP version in those fields.

I tested your exact tarball on an MI210. Both ROCM_PATH and ROCM_DISTRO_URL + ROCM_DISTRO_HASH built working PJRT plugins, and GPU computations passed with the same 7.15.26333 log.

Also, the checksum flag in the earlier example should be ROCM_DISTRO_HASH.

The gfx1036 error concerns your other GPU. Does the computation work with HIP_VISIBLE_DEVICES=0 to select the RX 9060 XT?

[joelberkeley]

Thanks for the quick reply!

XLA prints the HIP version in those fields.

Oh! I was wondering whether I was reading those fields right. What should I be checking to make sure the executable is using the ROCm version I tried to build it with? Or perhaps is there no straightforward way to do that?

The gfx1036 error concerns your other GPU. Does the computation work with HIP_VISIBLE_DEVICES=0 to select the RX 9060 XT?

The other device is my AMD CPU, which I didn't add support for, so it's expected. The program runs to completion, but I just want to make sure I'm shipping the right version (I'm building a shared lib artifact for downstream use).

[joelberkeley]

I've found ROCm release notes say which HIP version is used. Perhaps that's enough to know if I'm using the right ROCm version, via the HIP version.

https://rocm.docs.amd.com/en/docs-10.0.0/about/release-notes.html#rocm-core-sdk-components