Autogen dummy _REFERENCE_LAZY_BACKEND library when LTC is disabled

llvm · Jul 28, 2022 · 18f9e35 · 18f9e35
1 parent 34b3c15
commit 18f9e35
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 34 deletions.
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -84,8 +84,9 @@ declare_mlir_python_extension(TorchMLIRPythonExtensions.Main
 
 if(TORCH_MLIR_ENABLE_LTC)
   add_subdirectory(torch_mlir/csrc/base_lazy_backend)
-  add_subdirectory(torch_mlir/csrc/reference_lazy_backend)
 endif()
+# Reference backend has a separate check for TORCH_MLIR_ENABLE_LTC.
+add_subdirectory(torch_mlir/csrc/reference_lazy_backend)
 
 ################################################################################
 # Optionally handle JIT IR importer.

diff --git a/python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt b/python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt
@@ -26,36 +26,47 @@ mlir_configure_python_dev_packages()
 # Library definition
 ###########################################################################
 
-include_directories(BEFORE
-        ${TORCH_INCLUDE_DIRS}
-        ${CMAKE_CURRENT_SOURCE_DIR}
-        ${CMAKE_CURRENT_BINARY_DIR}
-        ${Python3_INCLUDE_DIRS}
-        ${PYTHON_H_DIR}
-        ${PROJECT_SOURCE_DIR}/python
-        )
-link_directories("${TORCH_INSTALL_PREFIX}/lib")
-link_directories(${CMAKE_CURRENT_SOURCE_DIR}/lib)
-add_link_options(-Wl,-rpath,$ORIGIN/lib)
-
-add_library(reference_lazy_backend SHARED
-        backend_impl.cpp
-        reference_lazy_backend_pybind.cpp
-        )
-add_dependencies(reference_lazy_backend
-        torch_mlir_ltc_backend
-        )
-target_link_libraries(reference_lazy_backend
-        ${TORCH_LIBRARIES}
-        torch_mlir_ltc_backend
-        )
-
-message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS} -Wno-pedantic")
-set_target_properties(reference_lazy_backend PROPERTIES
-        LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/reference_lazy_backend"
-        OUTPUT_NAME _REFERENCE_LAZY_BACKEND
-        PREFIX "${PYTHON_MODULE_PREFIX}"
-        SUFFIX "${PYTHON_MODULE_EXTENSION}"
-        CXX_VISIBILITY_PRESET "hidden"
-        COMPILE_FLAGS "${TORCH_CXXFLAGS} -Wno-pedantic"
-        )
+set(LIBRARY_OUTPUT_PATH  "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/reference_lazy_backend")
+set(OUTPUT_NAME "_REFERENCE_LAZY_BACKEND")
+
+if(TORCH_MLIR_ENABLE_LTC)
+  include_directories(BEFORE
+          ${TORCH_INCLUDE_DIRS}
+          ${CMAKE_CURRENT_SOURCE_DIR}
+          ${CMAKE_CURRENT_BINARY_DIR}
+          ${Python3_INCLUDE_DIRS}
+          ${PYTHON_H_DIR}
+          ${PROJECT_SOURCE_DIR}/python
+          )
+  link_directories("${TORCH_INSTALL_PREFIX}/lib")
+  link_directories(${CMAKE_CURRENT_SOURCE_DIR}/lib)
+  add_link_options(-Wl,-rpath,$ORIGIN/lib)
+
+  add_library(reference_lazy_backend SHARED
+          backend_impl.cpp
+          reference_lazy_backend_pybind.cpp
+          )
+  add_dependencies(reference_lazy_backend
+          torch_mlir_ltc_backend
+          )
+  target_link_libraries(reference_lazy_backend
+          ${TORCH_LIBRARIES}
+          torch_mlir_ltc_backend
+          )
+
+  message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS} -Wno-pedantic")
+  set_target_properties(reference_lazy_backend PROPERTIES
+          LIBRARY_OUTPUT_DIRECTORY ${LIBRARY_OUTPUT_PATH}
+          OUTPUT_NAME ${OUTPUT_NAME}
+          PREFIX "${PYTHON_MODULE_PREFIX}"
+          SUFFIX "${PYTHON_MODULE_EXTENSION}"
+          CXX_VISIBILITY_PRESET "hidden"
+          COMPILE_FLAGS "${TORCH_CXXFLAGS} -Wno-pedantic"
+          )
+else()
+  # To avoid import errors when LTC is disabled (a bunch of checks
+  # associated with that), we will generate a dummy placeholder library.
+  execute_process(
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/gen_dummy_lib.py ${LIBRARY_OUTPUT_PATH} ${OUTPUT_NAME}
+  )
+endif()
diff --git a/python/torch_mlir/csrc/reference_lazy_backend/gen_dummy_lib.py b/python/torch_mlir/csrc/reference_lazy_backend/gen_dummy_lib.py
@@ -0,0 +1,23 @@
+# When LTC is disabled in Torch-MLIR build, we will generate a dummy module to
+# ensure that no import errors occur.
+
+import sys
+import os
+
+if __name__ == '__main__':
+    path = sys.argv[1]  # dummy script path
+    file_name = sys.argv[2]  # dummy script
+
+    contents = '''
+# This file was automatically generated due to LTC being disabled in build. 
+    
+class LazyTensorCoreTestConfig:
+    def __init__(self):
+        assert False, "LTC is not enabled. Check the value of `TORCH_MLIR_ENABLE_LTC`"
+    '''
+
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+    with open(os.path.join(path, file_name + '.py'), 'w') as file:
+        file.write(contents)