Speed up libtorch build. (#1031)

build_tools/build_libtorch.sh

@@ -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_BIN} install -qr $PYTORCH_ROOT/requirements.txt
-  pip list
+  ${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 reset --hard HEAD
-  git checkout origin/"${PYTORCH_BRANCH}"
+  git clean -df
-  git submodule update --init --recursive
+  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
+  # More flags that require interop testing
-  USE_LIGHTWEIGHT_DISPATCH=OFF
+  # USE_LIGHTWEIGHT_DISPATCH=OFF
-  STATIC_DISPATCH_BACKEND=OFF
+  # STATIC_DISPATCH_BACKEND=CPU
-  BUILD_LITE_INTERPRETER=OFF
+  # 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.
+  else
-    # USE_LIGHTWEIGHT_DISPATCH=ON
+    BUILD_SHARED_LIBS=ON
-    # STATIC_DISPATCH_BACKEND=ON
-    # BUILD_LITE_INTERPRETER=OFF
   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"

python/CMakeLists.txt

@@ -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
+    COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
-      LIBTORCH_SRC_BUILD=${LIBTORCH_SRC_BUILD}
+    RESULT_VARIABLE _result
-      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
   )
   if(_result)
     message(FATAL_ERROR "Failed to run `build_libtorch.sh`")