Link against libtorch (#955)

This moves torch-mlir to link against libtorch on macOS and linux

TESTS: Tests pass. Tested release builds on linux and macOS

.gitignore

@@ -6,6 +6,8 @@
 .ipynb_checkpoints
 *.venv/
 mlir_venv/
+externals/pytorch/
+libtorch*
 
 /build/
 __pycache__

build_tools/build_libtorch.sh

@@ -0,0 +1,148 @@
+#!/usr/bin/env bash
+
+set -xe 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}
+PYTORCH_BRANCH="${PYTORCH_BRANCH:-master}"
+LIBTORCH_VARIANT="${LIBTORCH_VARIANT:-static-without-deps}"
+PT_C_COMPILER="${PT_C_COMPILER:-clang}"
+PT_CXX_COMPILER="${PT_CXX_COMPILER:-clang++}"
+
+echo "SRC_ROOT=${SRC_ROOT}"
+echo "PYTORCH_ROOT=${PYTORCH_ROOT}"
+echo "PYTORCH_BRANCH=${PYTORCH_BRANCH}"
+echo "LIBTORCH_VARIANT=${LIBTORCH_VARIANT}"
+
+if [[ "$LIBTORCH_VARIANT" == *"cxx11-abi"* ]]; then
+  echo _GLIBCXX_USE_CXX11_ABI=1
+  export _GLIBCXX_USE_CXX11_ABI=1
+  CXX_ABI=1
+  LIBTORCH_ABI="cxx11-abi-"
+else
+  echo _GLIBCXX_USE_CXX11_ABI=0
+  export _GLIBCXX_USE_CXX11_ABI=0
+  CXX_ABI=0
+  LIBTORCH_ABI=
+fi
+
+retry () {
+  $* || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
+}
+
+install_requirements() {
+  pip install -qr $PYTORCH_ROOT/requirements.txt
+  pip list
+}
+
+# Check for an existing libtorch at $PYTORCH_ROOT
+check_existing_libtorch() {
+  if [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.so" ]]; then
+    echo "Existing PYTORCH shared build found.. skipping build"
+    return 0
+  elif [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.a" ]]; then
+    echo "Existing PYTORCH static build found.. skipping build"
+    return 0
+  elif [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.dylib" ]]; then
+    echo "Existing PYTORCH shared dylib found.. skipping build"
+    return 0
+  fi
+  return 1
+}
+
+# Download and unzip into externals/pytorch/libtorch
+MACOS_X86_URL="https://download.pytorch.org/libtorch/nightly/cpu/libtorch-macos-latest.zip"
+# Download here (Pre-cxx11 ABI): https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
+# Download here (cxx11 ABI): https://download.pytorch.org/libtorch/nightly/cpu/libtorch-cxx11-abi-shared-with-deps-latest.zip
+# Download static here (cxx11 ABI): https://download.pytorch.org/libtorch/nightly/cpu/libtorch-cxx11-abi-static-with-deps-latest.zip
+# static builds are broken upstream and ship shared libraries anyway. Hopefully we can reland the fix upstream.
+LINUX_X86_URL="https://download.pytorch.org/libtorch/nightly/cpu/libtorch-static-without-deps-latest.zip"
+
+download_libtorch() {
+  cd $SRC_ROOT
+  if [[ $(uname -s) = 'Darwin' ]]; then
+  echo "Apple macOS detected"
+  if [[ $(uname -m) == 'arm64' ]]; then
+    echo "${Red}Apple M1 Detected...no libtorch/ binaries available"
+    return 1
+  else
+    echo "Apple x86_64 Detected"
+    DOWNLOAD_URL=${MACOS_X86_URL}
+  fi
+elif [[ $(uname -s) = 'Linux' ]]; then
+  echo "$Linux detected"
+  DOWNLOAD_URL=${LINUX_X86_URL}
+else
+  echo "OS not detected. Pray and Play"
+  return 1
+fi
+  curl -O ${DOWNLOAD_URL}
+  unzip -o libtorch-*.zip
+  if [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.so" ]]; then
+    echo "Verifying Pytorch install -- libtorch.so found"
+    return 0
+  elif [[ -f "$PYTORCH_INSTALL_PATH/lib/libtorch.a" ]]; then
+    echo "Verifying Pytorch install -- libtorch.a found"
+    return 0
+  fi
+  return 1
+}
+
+checkout_pytorch() {
+  if [[ ! -d "$PYTORCH_ROOT" ]]; then
+    git clone https://github.com/pytorch/pytorch $PYTORCH_ROOT
+  fi
+  cd $PYTORCH_ROOT
+  git fetch --all
+  git checkout ${PYTORCH_BRANCH}
+  git submodule update --init --recursive
+}
+
+build_pytorch() {
+  BUILD_SHARED_VAR="ON"
+  if [[ $LIBTORCH_VARIANT = *"static"* ]]; then
+    BUILD_SHARED_VAR="OFF"
+  fi
+  cd $PYTORCH_ROOT
+  BUILD_SHARED_LIBS=${BUILD_SHARED_VAR} BUILD_TESTS=OFF USE_GLOO=OFF USE_PYTORCH_QNNPACK=OFF USE_OPENMP=OFF  USE_OBSERVERS=OFF USE_KINETO=OFF USE_EIGEN_FOR_BLAS=OFF _GLIBCXX_USE_CXX11_ABI=${CXX_ABI} USE_NCCL=OFF INTERN_DISABLE_ONNX=OFF BUILD_PYTHONLESS=1 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 python setup.py build
+}
+
+package_pytorch() {
+  mkdir -p libtorch/{lib,bin,include,share}
+
+  # Copy over all lib files
+  cp -rv build/lib/*                libtorch/lib/
+  cp -rv build/lib*/torch/lib/*     libtorch/lib/
+
+  # Copy over all include files
+  cp -rv build/include/*            libtorch/include/
+  cp -rv build/lib*/torch/include/* libtorch/include/
+
+  # Copy over all of the cmake files
+  cp -rv build/lib*/torch/share/*   libtorch/share/
+
+  echo "${PYTORCH_BUILD_VERSION}" > libtorch/build-version
+  echo "$(pushd $PYTORCH_ROOT && git rev-parse HEAD)" > libtorch/build-hash
+  echo "Installing libtorch in ${PYTORCH_ROOT}/../../"
+  echo "deleteing old ${PYTORCH_ROOT}/../../libtorch"
+  rm -rf ${PYTORCH_ROOT}/../../libtorch
+  mv libtorch ${PYTORCH_ROOT}/../../
+}
+
+#main
+if check_existing_libtorch; then
+  echo "Found libtorch"
+  echo "Remove libtorch/ if you want to re-download or rebuild"
+else
+  if [ $SRC_BUILD ]; then
+    echo "Building libtorch from source"
+    checkout_pytorch
+    install_requirements
+    build_pytorch
+    package_pytorch
+  else
+    echo "Downloading libtorch"
+    download_libtorch
+  fi
+fi

build_tools/python_deploy/build_macos_packages.sh

@@ -105,7 +105,7 @@ function run_audit_wheel() {
     python${python_version} -m pip install -U pip
     python${python_version} -m pip install -r $repo_root/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
     python${python_version} -m pip install $generic_wheel --extra-index-url https://download.pytorch.org/whl/nightly/cpu
-    DYLD_LIBRARY_PATH=$output_dir/test_venv/lib/python${python_version}/site-packages/torch/lib delocate-wheel -v $generic_wheel
+    DYLD_LIBRARY_PATH=$repo_root/libtorch/lib delocate-wheel -v $generic_wheel
     deactivate
     rm -rf $output_dir/test_venv
   fi

python/CMakeLists.txt

@@ -61,6 +61,13 @@ declare_mlir_python_extension(TorchMLIRPythonExtensions.Main
 ################################################################################
 
 if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
+  # Build or Download libtorch
+  # if a libtorch/ exists we respect and don't override it. End user can
+  # unpack any released variant (shared/static | cxx11) of libtorch there. 
+  execute_process(
+    COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
+  )
+  set(TORCH_INSTALL_PREFIX "libtorch")
   add_subdirectory(torch_mlir/dialects/torch/importer/jit_ir)
   add_subdirectory(torch_mlir_e2e_test)
 endif()

python/torch_mlir/dialects/torch/importer/jit_ir/CMakeLists.txt

@@ -4,11 +4,12 @@
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
 include(TorchMLIRPyTorch)
-TorchMLIRProbeForPyTorchInstall()
+set(Torch_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../../../libtorch/share/cmake/Torch")
-find_package(Torch 1.8 REQUIRED)
+find_package(Torch 1.11 REQUIRED)
 
-TorchMLIRConfigurePyTorch()
+TorchMLIRConfigureLibTorch()
 
+message(STATUS "libtorch_python CXXFLAGS is ...${TORCH_CXXFLAGS}")
 #-------------------------------------------------------------------------------
 # Subdirectories
 #-------------------------------------------------------------------------------

python/torch_mlir/dialects/torch/importer/jit_ir/cmake/modules/TorchMLIRPyTorch.cmake

@@ -84,6 +84,37 @@ function(TorchMLIRConfigurePyTorch)
   endif()
 endfunction()
 
+function(TorchMLIRConfigureLibTorch)
+  message(STATUS "Checking LibTorch ABI settings...")
+  if(${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
+    message(STATUS "libtorch_python is ${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so")
+    # Check dual ABI setting first
+    execute_process(
+      COMMAND bash "-c" "cat ${TORCH_INSTALL_PREFIX}/share/cmake/Torch/TorchConfig.cmake | egrep -o '_GLIBCXX_USE_CXX11_ABI=[0-1]' | egrep -o '.$'"
+      OUTPUT_VARIABLE _use_cxx11_abi
+      OUTPUT_STRIP_TRAILING_WHITESPACE)
+    message(STATUS "LibTorch C++ Dual ABI setting: \"${_use_cxx11_abi}\"")
+
+    # Check ABI compatibility version
+    execute_process(
+      COMMAND bash "-c" "strings ${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so | egrep '^_cxxabi[0-9]{4}' | egrep -o '..$'"
+      OUTPUT_VARIABLE _cxx_abi_version
+      OUTPUT_STRIP_TRAILING_WHITESPACE)
+    message(STATUS "LibTorch C++ ABI version: \"${_cxx_abi_version}\"")
+
+    # Specialize compile flags for compiler
+    if(${CMAKE_CXX_COMPILER_ID} STREQUAL "GNU")
+      set(TORCH_CXXFLAGS "-D_GLIBCXX_USE_CXX11_ABI=${_use_cxx11_abi} -fabi-version=${_cxx_abi_version}")
+    elseif(${CMAKE_CXX_COMPILER_ID} STREQUAL "Clang")
+      set(TORCH_CXXFLAGS "-D_GLIBCXX_USE_CXX11_ABI=${_use_cxx11_abi} -U__GXX_ABI_VERSION -D__GXX_ABI_VERSION=10${_cxx_abi_version} '-DPYBIND11_COMPILER_TYPE=\"_gcc\"'")
+    else()
+      message(WARNING "Unrecognized compiler. Cannot determine ABI flags.")
+      return()
+    endif()
+    set(TORCH_CXXFLAGS "${TORCH_CXXFLAGS}" PARENT_SCOPE)
+  endif()
+endfunction()
+
 function(torch_mlir_python_target_compile_options target)
   target_compile_options(${target} PRIVATE
     $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:

python/torch_mlir/eager_mode/CMakeLists.txt

@@ -1,12 +1,3 @@
-#-------------------------------------------------------------------------------
-# Setup PyTorch
-#-------------------------------------------------------------------------------
-
-list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
-find_package(Torch 1.8 REQUIRED)
-
-TorchMLIRConfigurePyTorch()
-
 #-------------------------------------------------------------------------------
 # Subdirectories
 #-------------------------------------------------------------------------------