[WIP] Barebones libtorch build

build_tools/build_libtorch.sh

@@ -54,52 +54,91 @@ checkout_pytorch() {
 }
 
 build_pytorch() {
-  cd "$PYTORCH_ROOT"
-  # Uncomment the next line if you want to iterate on source builds
-  # ${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
+  sed -i.bak 's/INTERN_USE_EIGEN_BLAS ON/INTERN_USE_EIGEN_BLAS OFF/g' "${PYTORCH_ROOT}"/CMakeLists.txt
+  sed -i.bak 's/set(BUILD_PYTHON OFF)/set(BUILD_PYTHON ON)/g' "${PYTORCH_ROOT}"/CMakeLists.txt
+  sed -i.bak 's/set(INTERN_DISABLE_ONNX ON)/set(INTERN_DISABLE_ONNX OFF)/g' "${PYTORCH_ROOT}"/CMakeLists.txt
+#  # around line 150 (under if(BUILD_LITE_INTERPRETER))
+#  sed -i.bak 's/set(all_cpu_cpp ${generated_sources} ${core_generated_sources} ${cpu_kernel_cpp})/set(all_cpu_cpp "${${CMAKE_PROJECT_NAME}_SOURCE_DIR}\/build\/aten\/src\/ATen\/RegisterSchema.cpp")/g' "${PYTORCH_ROOT}"/aten/src/ATen/CMakeLists.txt
+  sed -i.bak 's/set(all_cpu_cpp ${generated_sources} ${core_generated_sources} ${cpu_kernel_cpp})/set(all_cpu_cpp ${generated_sources} ${core_generated_sources})/g' "${PYTORCH_ROOT}"/aten/src/ATen/CMakeLists.txt
+  sed -i.bak 's/append_filelist("aten_native_source_non_codegen_list" all_cpu_cpp)/append_filelist("core_sources_full" all_cpu_cpp)/g' "${PYTORCH_ROOT}"/aten/src/ATen/CMakeLists.txt
+  sed -i.bak 's/PyTorchBackendDebugInfo/PyTorchBackendDebugInfoDummy/g' "${PYTORCH_ROOT}"/torch/csrc/jit/backends/backend_detail.cpp
+  sed -i.bak 's/backend_debug_info->setDebugInfoMap(std::move(debug_info_map));/\/\/backend_debug_info->setDebugInfoMap(std::move(debug_info_map));/g' "${PYTORCH_ROOT}"/torch/csrc/jit/backends/backend_detail.cpp
+
+  CMAKE_ARGS=()
+  if [ -x "$(command -v ninja)" ]; then
+    CMAKE_ARGS+=("-GNinja")
+  fi
+
+  CMAKE_ARGS+=("-DCMAKE_PREFIX_PATH=${PYTHON_BIN}")
+  CMAKE_ARGS+=("-DPYTHON_EXECUTABLE=$($PYTHON_BIN -c 'import sys; print(sys.executable)')")
+
+  # Necessary flags to hit minimal path
+  export BUILD_PYTORCH_MOBILE_WITH_HOST_TOOLCHAIN=1
+  CMAKE_ARGS+=("-DBUILD_LITE_INTERPRETER:BOOL=TRUE")
+  CMAKE_ARGS+=("-DBUILD_SHARED_LIBS:BOOL=FALSE")
+  # torch/csrc/jit/mobile/profiler_edge.cpp includes KinetoEdgeCPUProfiler::~KinetoEdgeCPUProfiler
+  CMAKE_ARGS+=("-DUSE_KINETO:BOOL=TRUE")
+  CMAKE_ARGS+=("-DUSE_LIGHTWEIGHT_DISPATCH:BOOL=TRUE")
+  CMAKE_ARGS+=("-DSTATIC_DISPATCH_BACKEND=CPU")
+
+#  CMAKE_ARGS+=("-DCMAKE_CXX_FLAGS=-Wl,--unresolved-symbols=ignore-all")
+#  CMAKE_ARGS+=("-DCMAKE_EXE_LINKER_FLAGS=-Wl,--unresolved-symbols=ignore-all")
+#  CMAKE_ARGS+=("-DCMAKE_SHARED_LINKER_FLAGS=-Wl,--unresolved-symbols=ignore-all")
+
+  # Disable unused dependencies
+  CMAKE_ARGS+=("-DBUILD_CUSTOM_PROTOBUF:BOOL=FALSE")
+  CMAKE_ARGS+=("-DBUILD_TEST:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_ASAN:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_BLAS:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_BREAKPAD:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_CUDA:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_CUDNN:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_DISTRIBUTED:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_FBGEMM:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_GFLAGS=OFF")
+  CMAKE_ARGS+=("-DUSE_GLOO:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_LAPACK:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_LEVELDB=OFF")
+  CMAKE_ARGS+=("-DUSE_LMDB=OFF")
+  CMAKE_ARGS+=("-DUSE_MKLDNN:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_MPI=OFF")
+  CMAKE_ARGS+=("-DUSE_NCCL:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_NNPACK:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_NUMPY:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_OBSERVERS:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_OPENCV=OFF")
+  CMAKE_ARGS+=("-DUSE_OPENMP:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_OPENMP=OFF")
+  CMAKE_ARGS+=("-DUSE_PYTORCH_QNNPACK:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_QNNPACK:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_ROCM=OFF")
+  CMAKE_ARGS+=("-DUSE_TENSORPIPE:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_VALGRIND:BOOL=FALSE")
+  CMAKE_ARGS+=("-DUSE_XNNPACK:BOOL=FALSE")
+
+  CMAKE_ARGS+=("-DCMAKE_C_COMPILER=$(which clang)")
+  CMAKE_ARGS+=("-DCMAKE_CXX_COMPILER=$(which clang++)")
+
+#  CMAKE_ARGS+=("-DCMAKE_INSTALL_RPATH=${PYTORCH_INSTALL_PATH}/lib")
+
+
+  BUILD_ROOT=${BUILD_ROOT:-"$PYTORCH_ROOT/build"}
+  mkdir -p $BUILD_ROOT
+  cd $BUILD_ROOT
+  cmake --build . --target clean || echo "No build to clean."
+
+  cmake "$PYTORCH_ROOT" \
+      -DCMAKE_INSTALL_PREFIX=$PYTORCH_INSTALL_PATH \
+      -DCMAKE_BUILD_TYPE=Release \
+      "${CMAKE_ARGS[@]}"
+
+  if [ "$(uname)" == 'Darwin' ]; then
+    MAX_JOBS=$(sysctl -n hw.ncpu)
+  else
+    MAX_JOBS=$(nproc)
   fi
 
-  BUILD_SHARED_LIBS=ON \
-  BUILD_CAFFE2_OPS=OFF \
-  INTERN_BUILD_ATEN_OPS=OFF \
-  ATEN_NO_TEST=OFF \
-  USE_LITE_INTERPRETER_PROFILER=OFF \
-  BUILD_TEST=OFF \
-  GLIBCXX_USE_CXX11_ABI=1 \
-  CMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES} \
-  MACOSX_DEPLOYMENT_TARGET=${MACOSX_DEPLOYMENT_TARGET} \
-  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 \
-  USE_DISTRIBUTED=OFF \
-  USE_EIGEN_FOR_BLAS=OFF \
-  USE_FBGEMM=ON \
-  USE_GLOO=OFF \
-  USE_KINETO=ON \
-  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 \
-  USE_QNNPACK=OFF \
-  USE_XNNPACK=OFF \
-  ${PYTHON_BIN} setup.py  bdist_wheel -d "$WHEELHOUSE"
+  cmake --build . --target install -- "-j${MAX_JOBS}"
 }
 
 package_pytorch() {
@@ -109,13 +148,11 @@ package_pytorch() {
   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/
+  mv "${PYTORCH_ROOT}"/build_minimal/install/share     libtorch/
   # Copy over all lib files
-  mv build/lib/*                libtorch/lib/
+  mv "${PYTORCH_ROOT}"/build_minimal/install/lib       libtorch/lib
   # Copy over all include files
-  mv build/include/*            libtorch/include/
+  mv "${PYTORCH_ROOT}"/build_minimal/install/include   libtorch/include
 
   (pushd "$PYTORCH_ROOT" && git rev-parse HEAD) > libtorch/build-hash
   echo "Installing libtorch in ${PYTORCH_ROOT}/../../"
@@ -129,24 +166,8 @@ install_pytorch() {
   ${PIP_BIN} install  --force-reinstall $WHEELHOUSE/*
 }
 
-unpack_pytorch() {
-  PYTHON_SITE=`${PYTHON_BIN} -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`
-  pip uninstall torch
-  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"
 checkout_pytorch
 install_requirements
 build_pytorch
-package_pytorch
-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

python/torch_mlir/__init__.py

@@ -6,7 +6,7 @@
 from typing import Sequence, Union, List
 from enum import Enum
 
-import torch
+# import torch
 
 from torch_mlir.passmanager import PassManager
 from .compiler_utils import run_pipeline_with_repro_report
@@ -85,7 +85,7 @@ class TensorPlaceholder:
     ```
     """
 
-    def __init__(self, shape: List[int], dtype: torch.dtype):
+    def __init__(self, shape: List[int], dtype):
         """Create a tensor with shape `shape` and dtype `dtype`.
 
         Args:
@@ -97,7 +97,7 @@ def __init__(self, shape: List[int], dtype: torch.dtype):
         self.dtype = dtype
 
     @staticmethod
-    def like(tensor: torch.Tensor, dynamic_axes: List[int] = None):
+    def like(tensor, dynamic_axes: List[int] = None):
         """Create a tensor placeholder that is like the given tensor.
 
         Args:
@@ -116,10 +116,10 @@ def like(tensor: torch.Tensor, dynamic_axes: List[int] = None):
         return TensorPlaceholder(shape, tensor.dtype)
 
 
-_example_arg = Union[TensorPlaceholder, torch.Tensor]
+_example_arg = Union[TensorPlaceholder]
 
 
-def compile(model: torch.nn.Module,
+def compile(model,
             example_args: Union[_example_arg, Sequence[_example_arg]],
             output_type: Union[str, "OutputType"] = OutputType.TORCH,
             use_tracing: bool = False,

python/torch_mlir/dialects/torch/importer/jit_ir/__init__.py

@@ -7,7 +7,7 @@
 # and binds names locally. It exists to allow for customization of behavior
 # prior to loading shared objects.
 
-import torch
+# import torch
 
 # Our native extension is not self-contained. It references libraries which
 # must come in via the above first.

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

@@ -22,10 +22,37 @@ add_library(TorchMLIRJITIRImporter MODULE
   torch_to_mlir_utils.cpp
   )
 
+
+function(find_static_library LIB_NAME OUT)
+  if (WIN32 OR MSVC)
+    set(CMAKE_FIND_LIBRARY_SUFFIXES ".lib")
+  elseif (UNIX)
+    set(CMAKE_FIND_LIBRARY_SUFFIXES ".a")
+  endif ()
+  find_library(
+          FOUND_${LIB_NAME}_STATIC
+          ${LIB_NAME}
+          PATHS "${TORCH_INSTALL_PREFIX}/lib"
+  )
+  if (FOUND_${LIB_NAME}_STATIC)
+    get_filename_component(ABS_FILE ${FOUND_${LIB_NAME}_STATIC} ABSOLUTE)
+  else ()
+    message(SEND_ERROR "Unable to find library ${LIB_NAME}")
+  endif ()
+  set(${OUT} ${ABS_FILE} PARENT_SCOPE)
+endfunction()
+
+find_static_library(torch TORCH)
+find_static_library(torch_cpu TORCH_CPU)
+find_static_library(c10 C10)
+find_static_library(cpuinfo CPUINFO)
+find_static_library(clog CLOG)
+find_static_library(pthreadpool PTHREADPOOL)
+find_static_library(kineto KINETO)
+
 target_link_libraries(TorchMLIRJITIRImporter
   TorchMLIRAggregateCAPI
-  ${TORCH_LIBRARIES}
-  torch_python
+  ${TORCH} ${TORCH_CPU} ${C10} ${CPUINFO} ${CLOG} ${PTHREADPOOL} ${KINETO}
 )
 
 # On static Python builds, there may not be Python libraries to link against

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/class_annotator_pybind.cpp

@@ -16,11 +16,11 @@
 using namespace torch_mlir;
 
 static c10::ScalarType convertToC10ScalarType(py::object obj) {
-  if (THPDtype_Check(obj.ptr())) {
-    // Need reinterpret_cast, since no C++-level inheritance is involved.
-    THPDtype *dtype = reinterpret_cast<THPDtype *>(obj.ptr());
-    return dtype->scalar_type;
-  }
+//  if (THPDtype_Check(obj.ptr())) {
+//    // Need reinterpret_cast, since no C++-level inheritance is involved.
+//    THPDtype *dtype = reinterpret_cast<THPDtype *>(obj.ptr());
+//    return dtype->scalar_type;
+//  }
   std::stringstream ss;
   ss << "unsupported scalar type '" << obj << "'";
   throw std::invalid_argument(ss.str());

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/ivalue_importer.cpp

@@ -11,6 +11,7 @@
 #include "class_annotator.h"
 #include "function_importer.h"
 #include "torch_to_mlir_utils.h"
+#include "packed_params.h"
 
 #include <unordered_map>
 
@@ -21,7 +22,6 @@
 #include "mlir-c/Diagnostics.h"
 #include "torch-mlir-c/TorchTypes.h"
 
-#include "ATen/native/quantized/PackedParams.h"
 #include "caffe2/core/scope_guard.h"
 
 using namespace torch_mlir;

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/packed_params.h

@@ -0,0 +1,73 @@
+#pragma once
+
+#include <ATen/core/ivalue.h>
+
+struct LinearPackedParamsBase : public torch::jit::CustomClassHolder {
+  virtual at::Tensor apply(
+      at::Tensor input,
+      double output_scale,
+      int64_t output_zero_point) = 0;
+  virtual at::Tensor apply_relu(
+      at::Tensor input,
+      double output_scale,
+      int64_t output_zero_point) = 0;
+
+  // out variant of LinearPackedParamsBase::apply
+  virtual at::Tensor& apply_out(
+      const at::Tensor& input,
+      double output_scale,
+      int64_t output_zero_point,
+      at::Tensor& output) {
+    throw std::runtime_error(
+        "apply_out is not implemented for this packed "
+        "parameter type");
+    return output;
+  }
+
+  virtual at::Tensor& apply_relu_out(
+      const at::Tensor& input,
+      double output_scale,
+      int64_t output_zero_point,
+      at::Tensor& output) {
+    throw std::runtime_error(
+        "apply_relu_out is not implemented for this packed "
+        "parameter type");
+    return output;
+  }
+
+  virtual at::Tensor apply_dynamic(
+      at::Tensor input,
+      bool reduce_range = false) = 0;
+  virtual at::Tensor apply_dynamic_relu(
+      at::Tensor input,
+      bool reduce_range = false) = 0;
+
+  virtual at::Tensor& apply_dynamic_out(
+      const at::Tensor& /* input */,
+      at::Tensor& output,
+      bool /* reduce_range */) {
+    throw std::runtime_error(
+        "apply_dynamic_out is not implemented for this packed "
+        "parameter type");
+    return output;
+  }
+  virtual at::Tensor& apply_dynamic_relu_out(
+      const at::Tensor& /* input */,
+      at::Tensor& output,
+      bool /* reduce_range */) {
+    throw std::runtime_error(
+        "apply_dynamic_relu_out is not implemented for this packed "
+        "parameter type");
+    return output;
+  }
+
+  virtual std::tuple<at::Tensor, c10::optional<at::Tensor>> unpack() = 0;
+
+  virtual c10::optional<at::Tensor> bias() = 0;
+
+  virtual void set_bias(c10::optional<at::Tensor> bias) {
+    throw std::runtime_error(
+        "set_bias is not implemented for this packed "
+        "parameter type");
+  }
+};

setup.py

@@ -41,7 +41,7 @@
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
-import torch
+# import torch
 
 PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION") or "0.0.1"
 
@@ -152,7 +152,6 @@ def build_extension(self, ext):
         # exact version we built against.
         # TODO: This includes the +cpu specifier which is overly
         # restrictive and a bit unfortunate.
-        f"torch=={torch.__version__}",
     ],
     zip_safe=False,
 )