Teach cmake how to find the installed PyTorch.

* In most situations, this eliminates the need to explicitly set a path to the Torch cmake files.
* Also upgrades to new Python3 find package. (should eliminate 2.x mismatches)
* Since PyTorch is located by asking Python where it is, this eliminates a lot of causes of mismatch. (one source of truth)

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.10)
+cmake_minimum_required(VERSION 3.12)
 
 if(POLICY CMP0068)
   cmake_policy(SET CMP0068 NEW)
@@ -30,6 +30,7 @@ option(NPCOMP_ENABLE_IREE "Enables the IREE backend (must configure location via
 option(NPCOMP_ENABLE_REFJIT "Enables the reference JIT backend." ON)
 option(NPCOMP_BUILD_NPCOMP_DYLIB "Enables shared build of NPCOMP dylib (depends on LLVM/MLIR dylib support)" ON)
 set(NPCOMP_IREE_SRCDIR "" CACHE STRING "If building IREE, then setting this elects to build from a source directory (versus installed package)")
+set(NPCOMP_ENABLE_PYTORCH "OPTIONAL" CACHE STRING "Enables the PyTorch frontend (OFF, OPTIONAL, REQUIRED)")
 
 # Turn on -gsplit-dwarf if requested in debug builds.
 if (NPCOMP_USE_SPLIT_DWARF AND
@@ -80,6 +81,7 @@ include(AddLLVM)
 include(AddMLIR)
 include(AddNPCOMP)
 include(HandleLLVMOptions)
+include(ConfigurePyTorch)
 include_directories(${LLVM_INCLUDE_DIRS})
 include_directories(${MLIR_INCLUDE_DIRS})
 include_directories(${PROJECT_SOURCE_DIR}/include)
@@ -133,20 +135,36 @@ endif()
 
 set(NPCOMP_PYTHON_BINDINGS_VERSION_LOCKED 1 CACHE BOOL
        "Links to specific python libraries, resolving all symbols.")
-
+find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
-# TODO(laurenzo): Rationalize with how this is done elsewhere
+message(STATUS "Found python include dirs: ${Python3_INCLUDE_DIRS}")
-find_package(PythonInterp REQUIRED)
+message(STATUS "Found python libraries: ${Python3_LIBRARIES}")
-find_package(PythonLibs)
-message(STATUS "Found python include dirs: ${PYTHON_INCLUDE_DIRS}")
-message(STATUS "Found ppython libraries: ${PYTHON_LIBRARIES}")
-find_package(pybind11 CONFIG REQUIRED)
-message(STATUS "Found pybind11 v${pybind11_VERSION}: ${pybind11_INCLUDE_DIRS}")
 
 #-------------------------------------------------------------------------------
 # Pytorch Configuration
 #-------------------------------------------------------------------------------
 
-find_package(Torch)
+if(NPCOMP_ENABLE_PYTORCH)
+  ProbeForPyTorchInstall()
+  if(NPCOMP_ENABLE_PYTORCH EQUAL "OPTIONAL")
+    find_package(Torch)
+  else()
+    find_package(Torch REQUIRED)
+  endif()
+endif()
+
+#-------------------------------------------------------------------------------
+# Pybind11 Configuration
+#-------------------------------------------------------------------------------
+
+find_package(pybind11 CONFIG REQUIRED)
+# TODO: pybind11 v2.6 switched from pybind11_INCLUDE_DIRS (plural) to
+# pybind11_INCLUDE_DIR (singular). A lot has changed in this area since this
+# was written and overall python config and pybind11 should be modernized.
+set(pybind11_INCLUDE_DIR ${pybind11_INCLUDE_DIR} ${pybind11_INCLUDE_DIRS})
+message(STATUS "Found pybind11 v${pybind11_VERSION}: ${pybind11_INCLUDE_DIR}")
+message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
+               "suffix = '${PYTHON_MODULE_SUFFIX}', "
+               "extension = '${PYTHON_MODULE_EXTENSION}")
 
 #-------------------------------------------------------------------------------
 # Directory setup

README.md

@@ -124,8 +124,7 @@ export PYTHONPATH="$(realpath python):$(realpath build/python)"
 
 ```shell
 # See note above about -D_GLIBCXX_USE_CXX11_ABI=0
-./build_tools/cmake_configure.sh \
+./build_tools/cmake_configure.sh
-  -DCMAKE_PREFIX_PATH=$CONDA_PREFIX/lib/python3.8/site-packages/torch/share/cmake/Torch
 cmake --build build --target check-npcomp check-frontends-pytorch
 ```
 

cmake/modules/ConfigurePyTorch.cmake

@@ -0,0 +1,22 @@
+function(ProbeForPyTorchInstall)
+  if(Torch_ROOT)
+    message(STATUS "Using cached Torch root = ${Torch_ROOT}")
+  else()
+    message(STATUS "Checking for PyTorch using ${PYTHON_EXECUTABLE} ...")
+    execute_process(
+      COMMAND ${PYTHON_EXECUTABLE}
+      -c "import os;import torch;print(os.path.dirname(torch.__file__), end='')"
+      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+      RESULT_VARIABLE PYTORCH_STATUS
+      OUTPUT_VARIABLE PYTORCH_PACKAGE_DIR)
+    if(NOT PYTORCH_STATUS EQUAL "0")
+      message(STATUS "Unable to 'import torch' with ${PYTHON_EXECUTABLE} (fallback to explicit config)")
+      return()
+    endif()
+    message(STATUS "Found PyTorch installation at ${PYTORCH_PACKAGE_DIR}")
+
+    # PyTorch stashes its installed .cmake files under share/cmake/Torch.
+    set(Torch_ROOT "${PYTORCH_PACKAGE_DIR}/share/cmake/Torch"
+        CACHE STRING "Torch package root")
+  endif()
+endfunction()

frontends/pytorch/csrc/CMakeLists.txt

@@ -2,13 +2,17 @@ add_subdirectory(builder)
 
 include(NpcompPython)
 
-include_directories(
+# Sharp edge: Torch extensions need to use the same pybind11 that torch
+# was compiled with, or else there will be issues in cross module exception
+# handling (which will abort instead of raise). We circumvent the possibility
+# by forcing the torch directories first.
+include_directories(BEFORE
   ${TORCH_INCLUDE_DIRS}
   ${TORCH_INSTALL_PREFIX}/include/TH
   ${TORCH_INSTALL_PREFIX}/include/THC/opt/pytorch/pytorch
   ${CMAKE_CURRENT_SOURCE_DIR}
   ${CMAKE_CURRENT_BINARY_DIR}
-  ${PYTHON_INCLUDE_DIRS}
+  ${Python3_INCLUDE_DIRS}
   )
 link_directories("${TORCH_INSTALL_PREFIX}/lib")
 
@@ -17,7 +21,7 @@ add_library(NPCOMPTorchMLIRExt SHARED
   )
 target_link_libraries(NPCOMPTorchMLIRExt
   ${TORCH_LIBRARIES}
-  ${PYTHON_LIBRARIES}
+  ${Python3_LIBRARIES}
   torch_python
   npcomp_torch_builder_bindings