Reenable LTC in out-of-tree build (for real this time) (#1205)

* Fix OOT LTC CI build failure

* Disable LTC during macOS package gen

* Add more details about static TorchMLIRJITIRImporter library

.github/workflows/buildAndTest.yml

@@ -75,7 +75,6 @@ jobs:
           -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="${GITHUB_WORKSPACE}/externals/llvm-external-projects/torch-mlir-dialects" \
           -DLLVM_TARGETS_TO_BUILD=host \
           -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-          -DTORCH_MLIR_ENABLE_LTC=ON \
           -DTORCH_MLIR_USE_INSTALLED_PYTORCH="${{ matrix.torch-binary }}" \
           -DPython3_EXECUTABLE="$(which python)" \
           $GITHUB_WORKSPACE/externals/llvm-project/llvm
@@ -116,6 +115,7 @@ jobs:
       # cross compile, can't test arm64
       if: ${{ matrix.os-arch == 'macos-arm64' && matrix.llvm-build == 'in-tree' }}
       run: |
+        # TODO: Reenable LTC after build on macOS-arm64 is fixed (https://github.com/llvm/torch-mlir/issues/1253)
         cmake -GNinja -Bbuild_arm64 \
           -DCMAKE_BUILD_TYPE=Release \
           -DCMAKE_C_COMPILER=clang \
@@ -134,6 +134,7 @@ jobs:
           -DLLVM_ENABLE_ZSTD=OFF \
           -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
           -DTORCH_MLIR_ENABLE_MHLO=OFF \
+          -DTORCH_MLIR_ENABLE_LTC=OFF \
           -DTORCH_MLIR_USE_INSTALLED_PYTORCH="${{ matrix.torch-binary }}" \
           -DMACOSX_DEPLOYMENT_TARGET=12.0 \
           -DPython3_EXECUTABLE="$(which python)" \

CMakeLists.txt

@@ -49,8 +49,7 @@ if(TORCH_MLIR_ENABLE_MHLO)
   endif()
 endif()
 
-# TODO: Reenable LTC once OOT build is successful (https://github.com/llvm/torch-mlir/issues/1154)
+option(TORCH_MLIR_ENABLE_LTC "Enables LTC backend" ON)
-option(TORCH_MLIR_ENABLE_LTC "Enables LTC backend" OFF)
 
 if(TORCH_MLIR_ENABLE_LTC)
     set(ENV{TORCH_MLIR_ENABLE_LTC} 1)

build_tools/python_deploy/build_macos_packages.sh

@@ -35,6 +35,10 @@ export CMAKE_OSX_ARCHITECTURES="${TORCH_MLIR_OSX_ARCH:-arm64;x86_64}"
 echo "CMAKE_OSX_ARCHITECTURES: $CMAKE_OSX_ARCHITECTURES"
 echo "MACOSX_DEPLOYMENT_TARGET $MACOSX_DEPLOYMENT_TARGET"
 
+# Disable LTC build on MacOS to avoid linkage issues
+# https://github.com/llvm/torch-mlir/issues/1253
+export TORCH_MLIR_ENABLE_LTC=0
+
 function run() {
   echo "Using python versions: ${python_versions}"
 

python/CMakeLists.txt

@@ -154,6 +154,7 @@ add_mlir_python_modules(TorchMLIRPythonModules
 # Then it would "just work".
 if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
   add_dependencies(TorchMLIRPythonModules TorchMLIRJITIRImporter)
+  add_dependencies(TorchMLIRPythonModules TorchMLIRJITIRImporterPybind)
   # Build the E2E Tests (which depend on the JIT IR importer now).
   add_dependencies(TorchMLIRPythonModules TorchMLIRE2ETestPythonModules)
 endif()

python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt

@@ -42,7 +42,7 @@ if(TORCH_MLIR_ENABLE_LTC)
   link_directories(${CMAKE_CURRENT_SOURCE_DIR}/lib)
   add_link_options(-Wl,-rpath,$ORIGIN/lib)
 
-  add_library(reference_lazy_backend SHARED
+  add_library(reference_lazy_backend MODULE
           backend_impl.cpp
           reference_lazy_backend_pybind.cpp
           )
@@ -51,6 +51,7 @@ if(TORCH_MLIR_ENABLE_LTC)
           )
   target_link_libraries(reference_lazy_backend
           ${TORCH_LIBRARIES}
+          torch_python
           torch_mlir_ltc_backend
           )
 

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

@@ -10,31 +10,46 @@ include_directories(BEFORE
   )
 link_directories("${TORCH_INSTALL_PREFIX}/lib")
 
-# TODO: Currently, out-of-tree build fails when LIBRARY_TYPE is set to SHARED, so we have this toggle.
+# Static library with core functionality.
-#       See https://github.com/llvm/torch-mlir/issues/1154 for more details.
+# We can't use a shared library here, due to issues with linking on macOS-arm64 (the library itself won't build)
-if(TORCH_MLIR_ENABLE_LTC)
+# For details, see: https://github.com/llvm/torch-mlir/runs/7919012376
-  set(LIBRARY_TYPE "SHARED")
+add_library(TorchMLIRJITIRImporter STATIC
-else()
-  set(LIBRARY_TYPE "MODULE")
-endif()
-
-add_library(TorchMLIRJITIRImporter ${LIBRARY_TYPE}
   class_annotator.cpp
-  class_annotator_pybind.cpp
-  get_registered_ops.cpp
   function_importer.cpp
-  module_builder.cpp
   node_importer.cpp
-  import_options_pybind.cpp
   ivalue_importer.cpp
-  init_python_bindings.cpp
   torch_to_mlir_utils.cpp
   )
-
 target_link_libraries(TorchMLIRJITIRImporter
   TorchMLIRAggregateCAPI
+  ${TORCH_LIBRARIES}
+  )
+message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS}")
+set_target_properties(TorchMLIRJITIRImporter PROPERTIES
+  LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"
+  OUTPUT_NAME lib_jit_ir_importer
+  PREFIX ""
+  SUFFIX ".a"
+  CXX_VISIBILITY_PRESET "default"
+  COMPILE_FLAGS "${TORCH_CXXFLAGS}"
+  )
+
+# Separate Pybind MODULE due to issues with a SHARED library.
+# https://github.com/llvm/torch-mlir/issues/1154
+add_library(TorchMLIRJITIRImporterPybind MODULE
+  class_annotator_pybind.cpp
+  get_registered_ops.cpp
+  import_options_pybind.cpp
+  init_python_bindings.cpp
+  module_builder.cpp
+  )
+add_dependencies(TorchMLIRJITIRImporterPybind
+  TorchMLIRJITIRImporter
+  )
+target_link_libraries(TorchMLIRJITIRImporterPybind
   ${TORCH_LIBRARIES}
   torch_python
+  TorchMLIRJITIRImporter
   )
 
 # On static Python builds, there may not be Python libraries to link against
@@ -42,13 +57,13 @@ target_link_libraries(TorchMLIRJITIRImporter
 # this because in that case it is set to NOTFOUND and CMake will consider
 # this an error.
 if(Python3_LIBRARIES)
-  target_link_libraries(TorchMLIRJITIRImporter
+  target_link_libraries(TorchMLIRJITIRImporterPybind
     ${Python3_LIBRARIES}
   )
 endif()
 
 message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS}")
-set_target_properties(TorchMLIRJITIRImporter PROPERTIES
+set_target_properties(TorchMLIRJITIRImporterPybind PROPERTIES
   LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"
   OUTPUT_NAME _jit_ir_importer
   PREFIX "${PYTHON_MODULE_PREFIX}"
@@ -56,7 +71,7 @@ set_target_properties(TorchMLIRJITIRImporter PROPERTIES
   CXX_VISIBILITY_PRESET "hidden"
   COMPILE_FLAGS "${TORCH_CXXFLAGS}"
   )
-mlir_python_setup_extension_rpath(TorchMLIRJITIRImporter)
+mlir_python_setup_extension_rpath(TorchMLIRJITIRImporterPybind)
 
-torch_mlir_python_target_compile_options(TorchMLIRJITIRImporter)
+torch_mlir_python_target_compile_options(TorchMLIRJITIRImporterPybind)
-mlir_check_all_link_libraries(TorchMLIRJITIRImporter)
+mlir_check_all_link_libraries(TorchMLIRJITIRImporterPybind)

setup.py

@@ -20,10 +20,6 @@
 # prevent this script from attempting to build the directory, and will simply
 # use the (presumed already built) directory as-is.
 #
-# By default the lazy tensor backend is disabled and not built to avoid conflicts
-# with the out-of-tree build. To enable it, set the TORCH_MLIR_ENABLE_LTC
-# environment variable to 1.
-#
 # The package version can be set with the TORCH_MLIR_PYTHON_PACKAGE_VERSION
 # environment variable. For example, this can be "20220330.357" for a snapshot
 # release on 2022-03-30 with build number 357.
@@ -85,11 +81,8 @@ class CMakeBuild(build_py):
                 f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON",
                 f"-DCMAKE_C_VISIBILITY_PRESET=hidden",
                 f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden",
+                f"-DTORCH_MLIR_ENABLE_LTC={'ON' if int(os.environ.get('TORCH_MLIR_ENABLE_LTC', 1)) else 'OFF'}",
             ]
-            # TODO: Enable LTC by default once JIT importer linkage issue is fixed (https://github.com/llvm/torch-mlir/issues/1154)
-            enable_ltc = bool(int(os.environ.get("TORCH_MLIR_ENABLE_LTC", 0)))
-            if not enable_ltc:
-                cmake_args.append("-DTORCH_MLIR_ENABLE_LTC=OFF")
 
             os.makedirs(cmake_build_dir, exist_ok=True)
             cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")