From 73c8de6c4355b2f7cc4e8c52f0dabb7557eda7a0 Mon Sep 17 00:00:00 2001
From: Wang Kangyu <geek418@gmail.com>
Date: Thu, 9 Dec 2021 22:07:44 +0800
Subject: [PATCH] auto find pytorch install package (#3394)

---
 tools/pnnx/CMakeLists.txt          | 15 +++++++++++---
 tools/pnnx/cmake/PNNXPyTorch.cmake | 32 ++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)
 create mode 100644 tools/pnnx/cmake/PNNXPyTorch.cmake

diff --git a/tools/pnnx/CMakeLists.txt b/tools/pnnx/CMakeLists.txt
index c9dc3efe0..a8cbd1466 100644
--- a/tools/pnnx/CMakeLists.txt
+++ b/tools/pnnx/CMakeLists.txt
@@ -1,5 +1,12 @@
 project(pnnx)
-cmake_minimum_required(VERSION 3.10)
+cmake_minimum_required(VERSION 3.12)
+
+if(POLICY CMP0074)
+  cmake_policy(SET CMP0074 NEW)
+endif()
+
+list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
+include(PNNXPyTorch)
 
 # c++14 is required for using torch headers
 set(CMAKE_CXX_STANDARD 14)
@@ -12,13 +19,15 @@ option(PNNX_COVERAGE "build for coverage" OFF)
 
 #set(Torch_INSTALL_DIR "/home/nihui/.local/lib/python3.9/site-packages/torch" CACHE STRING "")
 #set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/pytorch-v1.10.0/build/install" CACHE STRING "")
-set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/libtorch" CACHE STRING "")
+#set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/libtorch" CACHE STRING "")
 set(TorchVision_INSTALL_DIR "/home/nihui/osd/vision/build/install" CACHE STRING "")
 
-set(Torch_DIR "${Torch_INSTALL_DIR}/share/cmake/Torch")
+#set(Torch_DIR "${Torch_INSTALL_DIR}/share/cmake/Torch")
 set(TorchVision_DIR "${TorchVision_INSTALL_DIR}/share/cmake/TorchVision")
 
+PNNXProbeForPyTorchInstall()
 find_package(Torch REQUIRED)
+
 find_package(TorchVision QUIET)
 
 message(STATUS "Torch_VERSION = ${Torch_VERSION}")
diff --git a/tools/pnnx/cmake/PNNXPyTorch.cmake b/tools/pnnx/cmake/PNNXPyTorch.cmake
new file mode 100644
index 000000000..bb9813015
--- /dev/null
+++ b/tools/pnnx/cmake/PNNXPyTorch.cmake
@@ -0,0 +1,32 @@
+# reference to https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/dialects/torch/importer/jit_ir/cmake/modules/TorchMLIRPyTorch.cmake
+
+# PNNXProbeForPyTorchInstall
+# Attempts to find a Torch installation and set the Torch_ROOT variable
+# based on introspecting the python environment. This allows a subsequent
+# call to find_package(Torch) to work.
+function(PNNXProbeForPyTorchInstall)
+  if(Torch_ROOT)
+    message(STATUS "Using cached Torch root = ${Torch_ROOT}")
+  elseif(Torch_INSTALL_DIR)
+    message(STATUS "Using cached Torch install dir = ${Torch_INSTALL_DIR}")
+    set(Torch_DIR "${Torch_INSTALL_DIR}/share/cmake/Torch" CACHE STRING "Torch dir" FORCE)
+  else()
+    #find_package (Python3 COMPONENTS Interpreter Development)
+    find_package (Python3)
+    message(STATUS "Checking for PyTorch using ${Python3_EXECUTABLE} ...")
+    execute_process(
+      COMMAND "${Python3_EXECUTABLE}"
+      -c "import os;import torch;print(torch.utils.cmake_prefix_path, 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 ${Python3_EXECUTABLE} (fallback to explicit config)")
+      return()
+    endif()
+    message(STATUS "Found PyTorch installation at ${PYTORCH_PACKAGE_DIR}")
+
+    set(Torch_ROOT "${PYTORCH_PACKAGE_DIR}" CACHE STRING
+        "Torch configure directory" FORCE)
+  endif()
+endfunction()