Check value of TORCH_MLIR_USE_INSTALLED_PYTORCH during LTC build

llvm · Jul 26, 2022 · c053fd1 · c053fd1
1 parent 4106a7d
commit c053fd1
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 20 deletions.
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -13,6 +13,28 @@ set(TORCH_MLIR_PYTHON_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/torch_mlir")
 # We vendor our own MLIR instance in the `torch_mlir` namespace.
 add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=torch_mlir.")
 
+################################################################################
+# PyTorch
+################################################################################
+
+option(TORCH_MLIR_USE_INSTALLED_PYTORCH "Build from local PyTorch in environment" ON)
+
+if (NOT TORCH_MLIR_USE_INSTALLED_PYTORCH)
+  # Source builds
+  set(ENV{MACOSX_DEPLOYMENT_TARGET} ${MACOSX_DEPLOYMENT_TARGET})
+  set(ENV{CMAKE_OSX_ARCHITECTURES} ${CMAKE_OSX_ARCHITECTURES})
+  set(ENV{CMAKE_C_COMPILER_LAUNCHER} ${CMAKE_C_COMPILER_LAUNCHER})
+  set(ENV{CMAKE_CXX_COMPILER_LAUNCHER} ${CMAKE_CXX_COMPILER_LAUNCHER})
+  execute_process(
+          COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
+          RESULT_VARIABLE _result
+  )
+  if(_result)
+    message(FATAL_ERROR "Failed to run `build_libtorch.sh`")
+  endif()
+  set(TORCH_INSTALL_PREFIX "libtorch")
+endif()
+
 ################################################################################
 # Sources
 ################################################################################
@@ -67,24 +89,7 @@ add_subdirectory(torch_mlir/csrc/reference_lazy_backend)
 # Optionally handle JIT IR importer.
 ################################################################################
 
-option(TORCH_MLIR_USE_INSTALLED_PYTORCH "Build from local PyTorch in environment" ON)
-
 if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
-  if (NOT TORCH_MLIR_USE_INSTALLED_PYTORCH)
-    # Source builds
-    set(ENV{MACOSX_DEPLOYMENT_TARGET} ${MACOSX_DEPLOYMENT_TARGET})
-    set(ENV{CMAKE_OSX_ARCHITECTURES} ${CMAKE_OSX_ARCHITECTURES})
-    set(ENV{CMAKE_C_COMPILER_LAUNCHER} ${CMAKE_C_COMPILER_LAUNCHER})
-    set(ENV{CMAKE_CXX_COMPILER_LAUNCHER} ${CMAKE_CXX_COMPILER_LAUNCHER})
-    execute_process(
-      COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/../build_tools/build_libtorch.sh
-      RESULT_VARIABLE _result
-    )
-    if(_result)
-      message(FATAL_ERROR "Failed to run `build_libtorch.sh`")
-    endif()
-    set(TORCH_INSTALL_PREFIX "libtorch")
-  endif()
   add_subdirectory(torch_mlir/dialects/torch/importer/jit_ir)
   add_subdirectory(torch_mlir_e2e_test)
 endif()

diff --git a/python/torch_mlir/csrc/base_lazy_backend/CMakeLists.txt b/python/torch_mlir/csrc/base_lazy_backend/CMakeLists.txt
@@ -5,10 +5,16 @@
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/python/torch_mlir/cmake/modules")
 include(TorchMLIRPyTorch)
+
 TorchMLIRProbeForPyTorchInstall()
+if(TORCH_MLIR_USE_INSTALLED_PYTORCH)
+  TorchMLIRConfigurePyTorch()
+else()
+  set(Torch_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libtorch/share/cmake/Torch")
+endif()
+
 find_package(Torch 1.11 REQUIRED)
 
-TorchMLIRConfigurePyTorch()
 set(TORCHGEN_DIR ${Torch_ROOT}/../../../torchgen)
 
 include_directories(BEFORE

diff --git a/python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt b/python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt
@@ -4,10 +4,15 @@
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/python/torch_mlir/cmake/modules")
 include(TorchMLIRPyTorch)
+
 TorchMLIRProbeForPyTorchInstall()
-find_package(Torch 1.11 REQUIRED)
+if(TORCH_MLIR_USE_INSTALLED_PYTORCH)
+  TorchMLIRConfigurePyTorch()
+else()
+  set(Torch_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libtorch/share/cmake/Torch")
+endif()
 
-TorchMLIRConfigurePyTorch()
+find_package(Torch 1.11 REQUIRED)
 
 ###########################################################################
 # Setup Python development