torch-mlir/build_tools/build_libtorch.sh

170 lines
5.4 KiB
Bash
Raw Normal View History

#!/usr/bin/env bash
set -xeu -o pipefail
SRC_ROOT="$( cd "$(dirname "$0")" ; pwd -P)/.."
PYTORCH_ROOT=${PYTORCH_ROOT:-$SRC_ROOT/externals/pytorch}
PYTORCH_INSTALL_PATH=${PYTORCH_INSTALL_PATH:-$SRC_ROOT/libtorch}
TORCH_MLIR_SRC_PYTORCH_REPO="${TORCH_MLIR_SRC_PYTORCH_REPO:-pytorch/pytorch}"
TORCH_MLIR_SRC_PYTORCH_BRANCH="${TORCH_MLIR_SRC_PYTORCH_BRANCH:-master}"
TM_PYTORCH_INSTALL_WITHOUT_REBUILD="${TM_PYTORCH_INSTALL_WITHOUT_REBUILD:-false}"
PT_C_COMPILER="${PT_C_COMPILER:-clang}"
PT_CXX_COMPILER="${PT_CXX_COMPILER:-clang++}"
CMAKE_OSX_ARCHITECTURES="${CMAKE_OSX_ARCHITECTURES:-x86_64}"
MACOSX_DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-12.0}"
WHEELHOUSE="${WHEELHOUSE:-$SRC_ROOT/build_tools/python_deploy/wheelhouse}"
2022-07-12 09:46:49 +08:00
PYTHON_BIN="${TORCH_MLIR_PYTHON_VERSION:-python3}"
PIP_BIN="${TORCH_MLIR_PIP_VERSION:-pip3}"
CMAKE_C_COMPILER_LAUNCHER="${CMAKE_C_COMPILER_LAUNCHER:-""}"
CMAKE_CXX_COMPILER_LAUNCHER="${CMAKE_CXX_COMPILER_LAUNCHER:-""}"
Red='\033[0;31m'
Green='\033[0;32m'
Yellow='\033[1;33m'
NC='\033[0m'
echo "SRC_ROOT=${SRC_ROOT}"
echo "PYTORCH_ROOT=${PYTORCH_ROOT}"
echo "TORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO}"
echo "TORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH}"
echo "MACOSX_DEPLOYMENT_TARGET=${MACOSX_DEPLOYMENT_TARGET}"
echo "CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES}"
export CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES}
export MACOSX_DEPLOYMENT_TARGET=${MACOSX_DEPLOYMENT_TARGET}
2022-07-12 09:46:49 +08:00
export CMAKE_C_COMPILER_LAUNCHER=${CMAKE_C_COMPILER_LAUNCHER}
export CMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER}
retry () {
"$@" || (sleep 1 && "$@") || (sleep 2 && "$@") || (sleep 4 && "$@") || (sleep 8 && "$@")
}
install_requirements() {
2022-07-12 09:46:49 +08:00
${PIP_BIN} install -qr $PYTORCH_ROOT/requirements.txt
${PIP_BIN} list
}
checkout_pytorch() {
if [[ ! -d "$PYTORCH_ROOT" ]]; then
# ${TORCH_MLIR_SRC_PYTORCH_BRANCH} could be a branch name or a commit hash.
# Although `git clone` can accept a branch name, the same command does not
# accept a commit hash, so we instead use `git fetch`. The alternative is
# to clone the entire repository and then `git checkout` the requested
# branch or commit hash, but that's too expensive.
mkdir "${PYTORCH_ROOT}"
cd "${PYTORCH_ROOT}"
git init
git remote add origin "https://github.com/${TORCH_MLIR_SRC_PYTORCH_REPO}"
git fetch --depth=1 origin "${TORCH_MLIR_SRC_PYTORCH_BRANCH}"
git reset --hard FETCH_HEAD
else
cd "${PYTORCH_ROOT}"
git fetch --depth=1 origin "${TORCH_MLIR_SRC_PYTORCH_BRANCH}"
git reset --hard FETCH_HEAD
fi
2022-07-12 09:46:49 +08:00
git clean -df
git submodule update --init --depth 1 --recursive
}
build_pytorch() {
cd "$PYTORCH_ROOT"
# Uncomment the next line if you want to iterate on source builds
2022-07-12 09:46:49 +08:00
# ${PYTHON_BIN} setup.py clean
rm -rf "${WHEELHOUSE:?}"/*
if [[ -z "${MAX_JOBS:-""}" ]]; then
if [[ "$(uname)" == 'Darwin' ]]; then
MAX_JOBS=$(sysctl -n hw.ncpu)
else
MAX_JOBS=$(nproc)
fi
fi
2022-07-12 09:46:49 +08:00
BUILD_SHARED_LIBS=ON \
BUILD_CAFFE2_OPS=OFF \
INTERN_BUILD_ATEN_OPS=OFF \
ATEN_NO_TEST=OFF \
USE_LITE_INTERPRETER_PROFILER=OFF \
2022-07-12 09:46:49 +08:00
BUILD_TEST=OFF \
GLIBCXX_USE_CXX11_ABI=1 \
2022-07-12 09:46:49 +08:00
CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES} \
MACOSX_DEPLOYMENT_TARGET=${MACOSX_DEPLOYMENT_TARGET} \
2022-07-12 09:46:49 +08:00
INTERN_BUILD_ATEN_OPS=OFF \
INTERN_DISABLE_ONNX=ON \
INTERN_USE_EIGEN_BLAS=ON \
MAX_JOBS=${MAX_JOBS} \
ONNX_ML=OFF \
USE_BREAKPAD=OFF \
USE_CUDA=OFF \
USE_ITT=OFF \
2022-07-12 09:46:49 +08:00
USE_DISTRIBUTED=OFF \
USE_EIGEN_FOR_BLAS=OFF \
USE_FBGEMM=ON \
2022-07-12 09:46:49 +08:00
USE_GLOO=OFF \
USE_KINETO=ON \
2022-07-12 09:46:49 +08:00
USE_MKL=OFF \
USE_MKLDNN=OFF \
USE_MPS=OFF \
USE_NCCL=OFF \
USE_NNPACK=OFF \
USE_OBSERVERS=OFF \
USE_OPENMP=OFF \
USE_PYTORCH_QNNPACK=ON \
2022-07-12 09:46:49 +08:00
USE_QNNPACK=OFF \
USE_XNNPACK=OFF \
USE_PRECOMPILED_HEADERS=1 \
2022-07-12 09:46:49 +08:00
${PYTHON_BIN} setup.py bdist_wheel -d "$WHEELHOUSE"
}
package_pytorch() {
2022-07-12 09:46:49 +08:00
mkdir -p libtorch
if [[ -d "libtorch/lib" ]]; then
rm -rf libtorch/{lib,bin,include,share}
fi
# Copy over all of the cmake files
mv build/lib*/torch/share libtorch/
mv build/lib*/torch/include libtorch/
mv build/lib*/torch/lib libtorch/
# Copy over all lib files
mv build/lib/* libtorch/lib/
# Copy over all include files
mv build/include/* libtorch/include/
(pushd "$PYTORCH_ROOT" && git rev-parse HEAD) > libtorch/build-hash
echo "Installing libtorch in ${PYTORCH_ROOT}/../../"
2022-07-12 09:46:49 +08:00
echo "deleting old ${PYTORCH_ROOT}/../../libtorch"
rm -rf "${PYTORCH_ROOT}"/../../libtorch
mv libtorch "${PYTORCH_ROOT}"/../../
}
2022-07-12 09:46:49 +08:00
install_pytorch() {
echo "pip installing Pytorch.."
${PIP_BIN} install --force-reinstall $WHEELHOUSE/*
}
unpack_pytorch() {
PYTHON_SITE=`${PYTHON_BIN} -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`
echo "wheel unpacking Pytorch..into ${PYTHON_SITE}"
wheel unpack -d "$WHEELHOUSE"/unpack_tmp "$WHEELHOUSE"/*.whl
mv "$WHEELHOUSE"/unpack_tmp/* "$PYTHON_SITE"/
}
#main
echo "Building libtorch from source"
build: make PyTorch caching more robust (#1510) Whether or not the PyTorch build is cached should not affect the success of the Torch-MLIR build, but based on the existing code, a build may fail if the `TM_PYTORCH_INSTALL_WITHOUT_REBUILD` variable was set but the build cache doesn't exist. Although that variable is set by CI upon a cache hit, nuances of Github's caching behavior can create situations where the coupling between `TM_PYTORCH_INSTALL_WITHOUT_REBUILD` and the cache lookup fails. Specifically, a branch other than our default branch (`main`) may create the cache entry, but because Github doesn't share this cache entry with builds running on the `main` branch, the `main` branch build tries to create it's own cache entry. However, since cache identifiers are unique and because caches are immutable, the caching step running in the `main` branch appears to create an invalid cache entry (of 233 bytes, instead of the expected ~60 MB). Consequently, subsequent builds observe a cache "hit", since caches created by the `main` branch are shared with all other branches, but because this cache entry is invalid (since it doesn't actually contain the ~60 MB PyTorch WHL file), the builds fail. One workaround would be to let only the `main` branch create caches, but in doing so, we would also prevent other branches from _reading_ the cache, making the builds in those branches terribly slow. So this patch uses a different workaround, which is to check whether the PyTorch WHL file exists, even if the build observed a cache hit. If the file doesn't exist, even if it was a purported cache hit, the code builds PyTorch from source, which is probably intuitive. A longer term fix will follow, after a discussion with the wider team.
2022-10-20 21:50:18 +08:00
wheel_exists=true
compgen -G "$WHEELHOUSE/*.whl" > /dev/null || wheel_exists=false
if [[ $TM_PYTORCH_INSTALL_WITHOUT_REBUILD != "true" || ${wheel_exists} == "false" ]]; then
checkout_pytorch
install_requirements
build_pytorch
package_pytorch
fi
if [[ $CMAKE_OSX_ARCHITECTURES = "arm64" ]]; then
echo "${Yellow} Cross compiling for arm64 so unpacking PyTorch wheel for libs${NC}"
unpack_pytorch
else
echo "${Green} Installing the built PyTorch wheel ${NC}"
install_pytorch
fi