Add a script to build libtorch

After a lot of trial and error it has a good balance of working
static library. In the future we can cull it further

Linux jit_ir tracer and eager mode to libtorch.a so we can ship
to pypi with a hermtic package. Remove any python level probing

.gitignore

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

CMakeLists.txt

@@ -25,6 +25,10 @@ project(torch-mlir LANGUAGES CXX C)
 set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 14)
 
+execute_process(
+  COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/build_tools/build_libtorch.sh
+)
+
 macro(torch_mlir_add_llvm_external_project name identifier location)
   message(STATUS "Adding LLVM external project ${name} (${identifier}) -> ${location}")
   if(NOT EXISTS "${location}/CMakeLists.txt")

build_tools/build_libtorch.sh

@@ -0,0 +1,105 @@
+#!/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
+# Download and unzip into externals/pytorch/libtorch
+# macOS x86_64 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
+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
+  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}/../../"
+  rm -rf ${PYTORCH_ROOT}/../../libtorch
+  mv libtorch ${PYTORCH_ROOT}/../../
+}
+
+if check_existing_libtorch; then
+  echo "Found libtorch"
+else
+  echo "Building libtorch.."
+  checkout_pytorch
+  install_requirements
+  build_pytorch
+  package_pytorch
+fi

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

@@ -2,12 +2,8 @@
 # Setup PyTorch
 #-------------------------------------------------------------------------------
 
-list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
-include(TorchMLIRPyTorch)
-TorchMLIRProbeForPyTorchInstall()
-find_package(Torch 1.8 REQUIRED)
-
-TorchMLIRConfigurePyTorch()
+set(Torch_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../../../libtorch/share/cmake/Torch")
+find_package(Torch 1.11 REQUIRED)
 
 #-------------------------------------------------------------------------------
 # Subdirectories

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

@@ -22,6 +22,14 @@ add_library(TorchMLIRJITIRImporter MODULE
   torch_to_mlir_utils.cpp
   )
 
+message(STATUS "TORCH_INCLUDE_DIRS=${TORCH_INCLUDE_DIRS}")
+target_include_directories(TorchMLIRJITIRImporter PUBLIC
+	$<BUILD_INTERFACE:${TORCH_INCLUDE_DIRS}
+	$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}
+	$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}
+	$<BUILD_INTERFACE:${Python3_INCLUDE_DIRS}
+)
+
 target_link_libraries(TorchMLIRJITIRImporter
   TorchMLIRAggregateCAPI
   ${TORCH_LIBRARIES}
@@ -49,5 +57,26 @@ set_target_properties(TorchMLIRJITIRImporter PROPERTIES
   )
 mlir_python_setup_extension_rpath(TorchMLIRJITIRImporter)
 
+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>>:
+      # Enable RTTI and exceptions.
+      -frtti -fexceptions
+      # Noisy pybind warnings
+      -Wno-unused-value
+      -Wno-covered-switch-default
+    >
+    $<$<CXX_COMPILER_ID:MSVC>:
+    # Enable RTTI and exceptions.
+    /EHsc /GR>
+  )
+endfunction()
+
 torch_mlir_python_target_compile_options(TorchMLIRJITIRImporter)
 mlir_check_all_link_libraries(TorchMLIRJITIRImporter)
+
+add_custom_command(
+  OUTPUT libtorch.a libc10.a
+  COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../../build_tools/build_libtorch.sh
+  BYPRODUCTS libtorch.a libc10.a
+  VERBATIM)

python/torch_mlir/eager_mode/CMakeLists.txt

@@ -2,10 +2,8 @@
 # Setup PyTorch
 #-------------------------------------------------------------------------------
 
-list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
-find_package(Torch 1.8 REQUIRED)
-
-TorchMLIRConfigurePyTorch()
+#set(Torch_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../libtorch/share/cmake/Torch")
+#find_package(Torch 1.11 REQUIRED)
 
 #-------------------------------------------------------------------------------
 # Subdirectories