Build jaxlib with AMD GPU ROCm support

build-jaxlib-rocm.sh

#!/bin/bash

set -ex

branch_name=rocm-jaxlib-v0.4.16

jax_repo=https://github.com/ROCmSoftwarePlatform/jax.git
jax_branch=$branch_name

xla_repo=https://github.com/ROCmSoftwarePlatform/xla.git
xla_branch=$branch_name

rocm_path=/opt/rocm

this_dir=$(dirname $(realpath $0))
echo $this_dir

if [[ ! -e ${this_dir}/jax ]]; then
    git clone $jax_repo --branch $jax_branch --filter=blob:none jax
    cd ${this_dir}/jax
else
    cd ${this_dir}/jax
    git fetch origin $jax_branch:$jax_branch && git checkout $jax_branch
fi

if [[ ! -e xla ]]; then
    git clone $xla_repo --branch $xla_branch --filter=blob:none xla
else
    pushd xla
    git fetch origin $xla_branch:$xla_branch && git checkout $xla_branch
    popd
fi

# for the the wheel name
rocm_version=$(cat /opt/rocm/.info/version | cut -d "-" -f 1)
export JAX_ROCM_VERSION=${rocm_version//./}
export JAX_RELEASE=1
export JAXLIB_RELEASE=1

# NOTE: adding version suffix to rocm_path is very important, otherwise
#
# this rule is missing dependency declarations for the following files included by ...:
#   '/opt/rocm-5.3.2/include/hip/hip_version.h'
#   '/opt/rocm-5.3.2/include/hip/hip_runtime.h'
#   '/opt/rocm-5.3.2/include/hip/hip_common.h'
#   ...
python build/build.py \
    --enable_rocm --rocm_path=$(realpath $rocm_path) \
    --bazel_options=--override_repository=xla=$(realpath xla)

python -m build --no-isolation --skip-dependency-check # --wheel