mirror of https://github.com/llvm/torch-mlir
Speed up libtorch build. (#1031)
Maksim Levental
GitHub
parent
11148e60d6
commit
1bb990afc7
|
|
@ -9,8 +9,12 @@ PYTORCH_BRANCH="${PYTORCH_BRANCH:-master}"
|
|||
LIBTORCH_VARIANT="${LIBTORCH_VARIANT:-shared}"
|
||||
PT_C_COMPILER="${PT_C_COMPILER:-clang}"
|
||||
PT_CXX_COMPILER="${PT_CXX_COMPILER:-clang++}"
|
||||
CMAKE_OSX_ARCHITECTURES="${CMAKE_OSX_ARCHITECTURES:-x86_64}"
|
||||
CMAKE_OSX_ARCHITECTURES="${CMAKE_OSX_ARCHITECTURES:-arm64;x86_64}"
|
||||
WHEELHOUSE="${WHEELHOUSE:-$SRC_ROOT/build_tools/python_deploy/wheelhouse}"
|
||||
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'
|
||||
|
|
@ -25,6 +29,8 @@ echo "LIBTORCH_SRC_BUILD=${LIBTORCH_SRC_BUILD}"
|
|||
echo "LIBTORCH_CACHE=${LIBTORCH_CACHE}"
|
||||
echo "CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES}"
|
||||
export CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES}
|
||||
export CMAKE_C_COMPILER_LAUNCHER=${CMAKE_C_COMPILER_LAUNCHER}
|
||||
export CMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER}
|
||||
|
||||
if [[ "$LIBTORCH_VARIANT" == *"cxx11-abi"* ]]; then
|
||||
echo _GLIBCXX_USE_CXX11_ABI=1
|
||||
|
|
@ -41,8 +47,8 @@ retry () {
|
|||
}
|
||||
|
||||
install_requirements() {
|
||||
pip install -qr "$PYTORCH_ROOT/requirements.txt"
|
||||
pip list
|
||||
${PIP_BIN} install -qr $PYTORCH_ROOT/requirements.txt
|
||||
${PIP_BIN} list
|
||||
}
|
||||
|
||||
# Check for an existing libtorch at $PYTORCH_ROOT
|
||||
|
|
@ -97,42 +103,93 @@ fi
|
|||
elif [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.a" ]]; then
|
||||
echo "Verifying Pytorch install -- libtorch.a found"
|
||||
return 0
|
||||
elif [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.dylib" ]]; then
|
||||
echo "Verifying Pytorch install -- libtorch.dylib found"
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
checkout_pytorch() {
|
||||
if [[ ! -d "$PYTORCH_ROOT" ]]; then
|
||||
git clone https://github.com/pytorch/pytorch "$PYTORCH_ROOT"
|
||||
git clone --depth 1 --single-branch --branch "${PYTORCH_BRANCH}" https://github.com/pytorch/pytorch "$PYTORCH_ROOT"
|
||||
fi
|
||||
cd "$PYTORCH_ROOT"
|
||||
git fetch --all
|
||||
git checkout origin/"${PYTORCH_BRANCH}"
|
||||
git submodule update --init --recursive
|
||||
git reset --hard HEAD
|
||||
git clean -df
|
||||
for dep in protobuf pocketfft cpuinfo FP16 psimd fmt sleef pybind11 onnx flatbuffers foxi; do
|
||||
git submodule update --init --depth 1 -- third_party/$dep
|
||||
done
|
||||
# setup.py will try to re-fetch
|
||||
sed -i.bak -E 's/^[[:space:]]+check_submodules()/#check_submodules()/g' setup.py
|
||||
}
|
||||
|
||||
build_pytorch() {
|
||||
cd "$PYTORCH_ROOT"
|
||||
# Uncomment the next line if you want to iterate on source builds
|
||||
python setup.py clean
|
||||
# ${PYTHON_BIN} setup.py clean
|
||||
rm -rf "${WHEELHOUSE:?}"/*
|
||||
|
||||
# Local fix for https://github.com/pytorch/pytorch/issues/81178
|
||||
sed -i.bak -E 's/namespace ao \{/namespace ao \{\n#include <ATen\/native\/ao_sparse\/quantized\/cpu\/packed_params.h>/g' aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_deserialize.cpp
|
||||
sed -i.bak -E 's/namespace ao \{/namespace ao \{\n#include <ATen\/native\/ao_sparse\/quantized\/cpu\/packed_params.h>/g' aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_serialize.cpp
|
||||
|
||||
BUILD_SHARED_LIBS=ON
|
||||
USE_LIGHTWEIGHT_DISPATCH=OFF
|
||||
STATIC_DISPATCH_BACKEND=OFF
|
||||
BUILD_LITE_INTERPRETER=OFF
|
||||
# More flags that require interop testing
|
||||
# USE_LIGHTWEIGHT_DISPATCH=OFF
|
||||
# STATIC_DISPATCH_BACKEND=CPU
|
||||
# BUILD_LITE_INTERPRETER=ON USE_LITE_INTERPRETER_PROFILER=OFF
|
||||
# INTERN_BUILD_MOBILE=ON BUILD_CAFFE2_MOBILE=ON SELECTED_OP_LIST=
|
||||
if [[ $LIBTORCH_VARIANT = *"static"* ]]; then
|
||||
BUILD_SHARED_LIBS=OFF
|
||||
# Enable after more testing.
|
||||
# USE_LIGHTWEIGHT_DISPATCH=ON
|
||||
# STATIC_DISPATCH_BACKEND=ON
|
||||
# BUILD_LITE_INTERPRETER=OFF
|
||||
else
|
||||
BUILD_SHARED_LIBS=ON
|
||||
fi
|
||||
BUILD_SHARED_LIBS=${BUILD_SHARED_LIBS} USE_LIGHTWEIGHT_DISPATCH=${USE_LIGHTWEIGHT_DISPATCH} STATIC_DISPATCH_BACKEND=${STATIC_DISPATCH_BACKEND} BUILD_LITE_INTERPRETER=${BUILD_LITE_INTERPRETER} BUILD_TEST=OFF USE_GLOO=OFF USE_MPS=OFF USE_PYTORCH_QNNPACK=OFF USE_OPENMP=OFF USE_OBSERVERS=OFF USE_KINETO=OFF USE_EIGEN_FOR_BLAS=OFF CMAKE_CXX_FLAGS="-D_GLIBCXX_USE_CXX11_ABI=${CXX_ABI}" USE_FBGEMM=OFF USE_NCCL=OFF INTERN_DISABLE_ONNX=OFF USE_CUDA=OFF USE_MKL=OFF USE_XNNPACK=OFF USE_DISTRIBUTED=OFF USE_BREAKPAD=OFF USE_MKLDNN=OFF USE_QNNPACK=OFF USE_NNPACK=OFF ONNX_ML=OFF CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES} python setup.py bdist_wheel -d "$WHEELHOUSE"
|
||||
|
||||
if [[ -z "${MAX_JOBS:-""}" ]]; then
|
||||
if [[ "$(uname)" == 'Darwin' ]]; then
|
||||
MAX_JOBS=$(sysctl -n hw.ncpu)
|
||||
else
|
||||
MAX_JOBS=$(nproc)
|
||||
fi
|
||||
fi
|
||||
|
||||
BUILD_CAFFE2_OPS=OFF \
|
||||
BUILD_SHARED_LIBS=${BUILD_SHARED_LIBS} \
|
||||
BUILD_TEST=OFF \
|
||||
CC=${PT_C_COMPILER} \
|
||||
CMAKE_CXX_FLAGS="-D_GLIBCXX_USE_CXX11_ABI=${CXX_ABI}" \
|
||||
CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES} \
|
||||
CXX=${PT_CXX_COMPILER} \
|
||||
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_DISTRIBUTED=OFF \
|
||||
USE_EIGEN_FOR_BLAS=OFF \
|
||||
USE_FBGEMM=OFF \
|
||||
USE_GLOO=OFF \
|
||||
USE_KINETO=OFF \
|
||||
USE_MKL=OFF \
|
||||
USE_MKLDNN=OFF \
|
||||
USE_MPS=OFF \
|
||||
USE_NCCL=OFF \
|
||||
USE_NNPACK=OFF \
|
||||
USE_OBSERVERS=OFF \
|
||||
USE_OPENMP=OFF \
|
||||
USE_PYTORCH_QNNPACK=OFF \
|
||||
USE_QNNPACK=OFF \
|
||||
USE_XNNPACK=OFF \
|
||||
${PYTHON_BIN} setup.py bdist_wheel -d "$WHEELHOUSE"
|
||||
}
|
||||
|
||||
package_pytorch() {
|
||||
mkdir -p libtorch/{lib,bin,include,share}
|
||||
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/
|
||||
|
|
@ -143,14 +200,19 @@ package_pytorch() {
|
|||
# Copy over all include files
|
||||
mv build/include/* libtorch/include/
|
||||
|
||||
echo "${PYTORCH_BUILD_VERSION}" > libtorch/build-version
|
||||
(pushd "$PYTORCH_ROOT" && git rev-parse HEAD) > libtorch/build-hash
|
||||
echo "Installing libtorch in ${PYTORCH_ROOT}/../../"
|
||||
echo "deleteing old ${PYTORCH_ROOT}/../../libtorch"
|
||||
echo "deleting old ${PYTORCH_ROOT}/../../libtorch"
|
||||
rm -rf "${PYTORCH_ROOT}"/../../libtorch
|
||||
mv libtorch "${PYTORCH_ROOT}"/../../
|
||||
}
|
||||
|
||||
install_pytorch() {
|
||||
echo "pip installing Pytorch.."
|
||||
${PIP_BIN} install --force-reinstall $WHEELHOUSE/*
|
||||
}
|
||||
|
||||
|
||||
#main
|
||||
if [[ $LIBTORCH_SRC_BUILD = "ON" ]]; then
|
||||
echo "Building libtorch from source"
|
||||
|
|
@ -158,6 +220,7 @@ if [[ $LIBTORCH_SRC_BUILD = "ON" ]]; then
|
|||
install_requirements
|
||||
build_pytorch
|
||||
package_pytorch
|
||||
install_pytorch
|
||||
else
|
||||
if check_existing_libtorch; then
|
||||
echo "Found existing libtorch"
|
||||
|
|
|
|||
|
|
@ -74,14 +74,15 @@ if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
|
|||
# if LIBTORCH_VARIANT=*cxx11abi* we build with the new CXX11 ABI similar to upstream pytorch/builder
|
||||
# if LIBTORCH_VARIANT=*shared* (default) we build a shared library (and pytorch .whl)
|
||||
# else if LIBTORCH_VARIANT=*static* we build a static libtorch
|
||||
set(ENV{LIBTORCH_SRC_BUILD} ${LIBTORCH_SRC_BUILD})
|
||||
set(ENV{LIBTORCH_CACHE} ${LIBTORCH_CACHE})
|
||||
set(ENV{LIBTORCH_VARIANT} ${LIBTORCH_VARIANT})
|
||||
set(ENV{CMAKE_OSX_ARCHITECTURES} ${CMAKE_OSX_ARCHITECTURES})
|
||||
set(ENV{CMAKE_C_COMPILER_LAUNCHER} ${CMAKE_C_COMPILER_LAUNCHER})
|
||||
set(ENV{CMAKE_CXX_COMPILER_LAUNCHER} ${CMAKE_CXX_COMPILER_LAUNCHER})
|
||||
execute_process(
|
||||
COMMAND ${CMAKE_COMMAND} -E env
|
||||
LIBTORCH_SRC_BUILD=${LIBTORCH_SRC_BUILD}
|
||||
LIBTORCH_CACHE=${LIBTORCH_CACHE}
|
||||
LIBTORCH_VARIANT=${LIBTORCH_VARIANT}
|
||||
CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES}
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
|
||||
RESULT_VARIABLE _result
|
||||
COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
|
||||
RESULT_VARIABLE _result
|
||||
)
|
||||
if(_result)
|
||||
message(FATAL_ERROR "Failed to run `build_libtorch.sh`")
|
||||
|
|
|
|||
Loading…
Reference in New Issue