-
Notifications
You must be signed in to change notification settings - Fork 2.7k
/
build_rocm.sh
executable file
·67 lines (59 loc) · 2.48 KB
/
build_rocm.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#!/usr/bin/env bash
# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Environment Var Notes
# XLA_CLONE_DIR -
# Specifies filepath to where XLA repo is cloned.
# NOTE:, if this is set then XLA repo is not cloned. Must clone repo before running this script.
# Also, if this is set then setting XLA_REPO and XLA_BRANCH have no effect.
# XLA_REPO
# XLA repo to clone from. Default is https://github.com/ROCmSoftwarePlatform/tensorflow-upstream
# XLA_BRANCH
# XLA branch in the XLA repo. Default is develop-upstream-jax
#
set -eux
python -V
#If XLA_REPO is not set, then use default
if [ ! -v XLA_REPO ]; then
XLA_REPO="https://github.com/openxla/xla.git"
XLA_BRANCH="main"
elif [ -z "$XLA_REPO" ]; then
XLA_REPO="https://github.com/openxla/xla.git"
XLA_BRANCH="main"
fi
#If XLA_CLONE_PATH is not set, then use default path.
#Note, setting XLA_CLONE_PATH makes setting XLA_REPO and XLA_BRANCH a no-op
#Set this when XLA repository has been already clone. This is useful in CI
#environments and when doing local development
if [ ! -v XLA_CLONE_DIR ]; then
XLA_CLONE_DIR=/tmp/xla
rm -rf /tmp/xla || true
git clone -b ${XLA_BRANCH} ${XLA_REPO} /tmp/xla
elif [ -z "$XLA_CLONE_DIR" ]; then
XLA_CLONE_DIR=/tmp/xla
rm -rf /tmp/xla || true
git clone -b ${XLA_BRANCH} ${XLA_REPO} /tmp/xla
fi
#Export JAX_ROCM_VERSION so that it is appened in the wheel name
rocm_version=$(cat /opt/rocm/.info/version | cut -d "-" -f 1)
export JAX_ROCM_VERSION=${rocm_version//./}
#Build and install wheel
python3 ./build/build.py --enable_rocm --rocm_path=${ROCM_PATH} --bazel_options=--override_repository=xla=${XLA_CLONE_DIR}
pip3 install --force-reinstall dist/*.whl # installs jaxlib (includes XLA)
pip3 install --force-reinstall . # installs jax
#This is for CI to read without having to start the container again
if [ -v CI_RUN ]; then
pip3 list | grep jaxlib | tr -s ' ' | cut -d " " -f 2 | cut -d "+" -f 1 > jax_version_installed
cat /opt/rocm/.info/version | cut -d "-" -f 1 > jax_rocm_version
fi