Re-organize project structure to separate PyTorch dependencies from core project. (#2542)

This is a first step towards the structure we discussed here:
https://gist.github.com/stellaraccident/931b068aaf7fa56f34069426740ebf20

There are two primary goals:

1. Separate the core project (C++ dialects and conversions) from the
hard PyTorch dependencies. We move all such things into projects/pt1 as
a starting point since they are presently entangled with PT1-era APIs.
Additional work can be done to disentangle components from that
(specifically LTC is identified as likely ultimately living in a
`projects/ltc`).
2. Create space for native PyTorch2 Dynamo-based infra to be upstreamed
without needing to co-exist with the original TorchScript path.

Very little changes in this path with respect to build layering or
options. These can be updated in a followup without commingling
directory structure changes.

This also takes steps toward a couple of other layering enhancements:

* Removes the llvm-external-projects/torch-mlir-dialects sub-project,
collapsing it into the main tree.
* Audits and fixes up the core C++ build to account for issues found
while moving things. This is just an opportunistic pass through but
roughly ~halves the number of build actions for the project from the
high 4000's to the low 2000's.

It deviates from the discussed plan by having a `projects/` tree instead
of `compat/`. As I was thinking about it, this will better accommodate
the follow-on code movement.

Once things are roughly in place and the CI passing, followups will
focus on more in-situ fixes and cleanups.

Author: stellaraccident

.github/workflows/buildAndTest.yml

@@ -125,9 +125,8 @@ jobs:
           -DCMAKE_OSX_ARCHITECTURES=arm64 \
           -DLLVM_ENABLE_ASSERTIONS=ON \
           -DLLVM_ENABLE_PROJECTS=mlir \
-          -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
+          -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
           -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$GITHUB_WORKSPACE" \
-          -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="${GITHUB_WORKSPACE}/externals/llvm-external-projects/torch-mlir-dialects" \
           -DLLVM_TARGETS_TO_BUILD=AArch64 \
           -DLLVM_USE_HOST_TOOLS=ON \
           -DLLVM_ENABLE_ZSTD=OFF \

.gitignore

@@ -26,7 +26,7 @@ __pycache__
 bazel-*
 
 # Autogenerated files
-/python/torch_mlir/csrc/base_lazy_backend/generated
+/projects/pt1/python/torch_mlir/csrc/base_lazy_backend/generated
 
 #Docker builds
 build_oot/

CMakeLists.txt

@@ -1,3 +1,7 @@
+#-------------------------------------------------------------------------------
+# Project setup and globals
+#-------------------------------------------------------------------------------
+
 cmake_minimum_required(VERSION 3.12)
 
 if(POLICY CMP0068)
@@ -17,52 +21,48 @@ if(POLICY CMP0116)
   cmake_policy(SET CMP0116 OLD)
 endif()
 
-#-------------------------------------------------------------------------------
-# Project setup and globals
-#-------------------------------------------------------------------------------
-
 project(torch-mlir LANGUAGES CXX C)
 set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 17)
 
-macro(torch_mlir_add_llvm_external_project name identifier location)
-  message(STATUS "Adding LLVM external project ${name} (${identifier}) -> ${location}")
-  if(NOT EXISTS "${location}/CMakeLists.txt")
-    message(FATAL_ERROR "External project location ${location} is not valid")
-  endif()
-  list(APPEND LLVM_EXTERNAL_PROJECTS ${name})
-  list(REMOVE_DUPLICATES LLVM_EXTERNAL_PROJECTS)
-  set(LLVM_EXTERNAL_${identifier}_SOURCE_DIR ${location} CACHE STRING "" FORCE)
-  set(LLVM_EXTERNAL_PROJECTS ${LLVM_EXTERNAL_PROJECTS} CACHE STRING "" FORCE)
-endmacro()
+#-------------------------------------------------------------------------------
+# Project options
+#-------------------------------------------------------------------------------
+
+option(TORCH_MLIR_ENABLE_REFBACKEND "Enable reference backend" ON)
+if(TORCH_MLIR_ENABLE_REFBACKEND)
+  add_definitions(-DTORCH_MLIR_ENABLE_REFBACKEND)
+endif()
 
 option(TORCH_MLIR_ENABLE_STABLEHLO "Add stablehlo dialect" ON)
 if(TORCH_MLIR_ENABLE_STABLEHLO)
   add_definitions(-DTORCH_MLIR_ENABLE_STABLEHLO)
 endif()
 
+option(TORCH_MLIR_OUT_OF_TREE_BUILD "Specifies an out of tree build" OFF)
+
+# PT1 options.
+option(TORCH_MLIR_ENABLE_PROJECT_PT1 "Enables the PyTorch1 project under projects/pt1" OFF)
+# TODO: Rename/scope these. They use historic names for now to ease migration
+# burden.
 option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 option(TORCH_MLIR_ENABLE_LTC "Enables LTC backend" OFF)
 option(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS "Build Torch dialect MLIR Python bindings but neither JIT IR Importer nor LTC backend" OFF)
 if(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS)
   set(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER OFF)
   set(TORCH_MLIR_ENABLE_LTC OFF)
 endif()
-
-if(TORCH_MLIR_ENABLE_LTC)
-    set(ENV{TORCH_MLIR_ENABLE_LTC} 1)
-    message(STATUS "LTC Backend build is enabled")
-else()
-    set(ENV{TORCH_MLIR_ENABLE_LTC} 0)
-    message(STATUS "LTC Backend build is disabled")
+# Force enable the PT1 project if either the JIT_IR_IMPORTER or LTC is enabled.
+if(NOT TORCH_MLIR_ENABLE_PROJECT_PT1)
+  if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER OR TORCH_MLIR_ENABLE_LTC)
+    message(STATUS "Enabling projects/pt1 because features requiring it are enabled")
+    set(TORCH_MLIR_ENABLE_PROJECT_PT1 ON)
+  endif()
 endif()
 
-torch_mlir_add_llvm_external_project(
-  torch-mlir-dialects
-  TORCH_MLIR_DIALECTS
-  ${CMAKE_CURRENT_SOURCE_DIR}/externals/llvm-external-projects/torch-mlir-dialects)
-
-option(TORCH_MLIR_OUT_OF_TREE_BUILD "Specifies an out of tree build" OFF)
+#-------------------------------------------------------------------------------
+# Configure out-of-tree vs in-tree build
+#-------------------------------------------------------------------------------
 
 if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR OR TORCH_MLIR_OUT_OF_TREE_BUILD)
   message(STATUS "Torch-MLIR out-of-tree build.")
@@ -99,11 +99,10 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR OR TORCH_MLIR_OUT_OF_TREE_
 
   # Don't try to compile the python extensions at the moment. We need
   # to import lots of dependencies from AddMLIRPython to make this work.
-  set(MLIR_ENABLE_BINDINGS_PYTHON 1)
+  set(MLIR_ENABLE_BINDINGS_PYTHON ON)
 
-  set(TORCH-MLIR_BUILT_STANDALONE 1)
+  set(TORCH-MLIR_BUILT_STANDALONE ON)
   set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}")
-  add_subdirectory(externals/llvm-external-projects/torch-mlir-dialects)
 else()
   message(STATUS "Torch-MLIR in-tree build.")
   # In-tree build with LLVM_EXTERNAL_PROJECTS=torch-mlir
@@ -169,24 +168,14 @@ add_subdirectory(lib)
 add_subdirectory(tools)
 
 add_custom_target(check-torch-mlir-all)
-add_dependencies(check-torch-mlir-all
-  check-torch-mlir
-  check-torch-mlir-dialects
-  check-torch-mlir-capi
-  )
+add_dependencies(check-torch-mlir-all check-torch-mlir)
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   # If parent projects want to configure where to place the python packages,
   # respect that.
   if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
     set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   endif()
-  add_dependencies(check-torch-mlir-all
-  check-torch-mlir-python
-  )
-  add_subdirectory(python)
-else()
-  add_custom_target(TorchMLIRPythonModules)
 endif()
 
 add_subdirectory(test)
@@ -237,3 +226,11 @@ if (TORCH_MLIR_ENABLE_STABLEHLO)
     EXCLUDE_FROM_ALL)
   include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo)
 endif()
+
+#-------------------------------------------------------------------------------
+# Sub-projects
+#-------------------------------------------------------------------------------
+
+if(TORCH_MLIR_ENABLE_PROJECT_PT1)
+  add_subdirectory(projects/pt1)
+endif()

build_tools/autogen_ltc_backend.py

@@ -29,7 +29,7 @@
     TORCH_INCLUDE_DIR = TORCH_DIR
 TORCHGEN_DIR = Path(torchgen.__path__[0]).resolve()
 TORCH_MLIR_DIR = Path(__file__).resolve().parent.parent
-
+TORCH_MLIR_PT1_DIR = TORCH_MLIR_DIR / "projects" / "pt1"
 
 def reindent(text, prefix=""):
     return indent(dedent(text), prefix)
@@ -114,12 +114,12 @@ def __init__(self, binary_dir):
         self.binary_dir = Path(binary_dir)
         assert self.binary_dir.is_dir(), f"Binary directory not found: {self.binary_dir}"
         self.source_yaml = self.binary_dir.joinpath("generated_native_functions.yaml")
-        self.backend_path = TORCH_MLIR_DIR.joinpath(
+        self.backend_path = TORCH_MLIR_PT1_DIR.joinpath(
             "python", "torch_mlir", "csrc", "base_lazy_backend"
         )
-        assert self.backend_path.is_dir()
+        assert self.backend_path.is_dir(), f"Backend path not found: {self.backend_path}"
         self.generated_path = self.binary_dir.joinpath(
-            "python", "torch_mlir", "csrc", "base_lazy_backend", "generated"
+            "projects", "pt1", "python", "torch_mlir", "csrc", "base_lazy_backend", "generated"
         )
         self.generated_path.mkdir(parents=True, exist_ok=True)
 

build_tools/build_standalone.sh

@@ -19,9 +19,8 @@ cmake -GNinja -B"$build_dir" "$llvm_project_dir/llvm" \
   -DCMAKE_BUILD_TYPE=Release \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DLLVM_ENABLE_PROJECTS=mlir \
-  -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
+  -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$project_dir" \
-  -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="${project_dir}"/externals/llvm-external-projects/torch-mlir-dialects \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DLLVM_TARGETS_TO_BUILD=host

build_tools/python_deploy/build_linux_packages.sh

@@ -234,9 +234,8 @@ function build_in_tree() {
       -DCMAKE_C_COMPILER_LAUNCHER=ccache \
       -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
       -DLLVM_ENABLE_PROJECTS=mlir \
-      -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
+      -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
       -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="/main_checkout/torch-mlir" \
-      -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="/main_checkout/torch-mlir/externals/llvm-external-projects/torch-mlir-dialects" \
       -DLLVM_TARGETS_TO_BUILD=host \
       -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
       -DTORCH_MLIR_ENABLE_LTC=${enable_ltc} \
@@ -246,7 +245,7 @@ function build_in_tree() {
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
-  cmake --build /main_checkout/torch-mlir/build
+  cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
 }
 
@@ -288,7 +287,7 @@ function test_in_tree() {
   cmake --build /main_checkout/torch-mlir/build --target check-torch-mlir-all
 
   cd /main_checkout/torch-mlir/
-  export PYTHONPATH="/main_checkout/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir"
+  export PYTHONPATH="/main_checkout/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir:/main_checkout/torch-mlir/projects/pt1"
 
   case $torch_version in
     nightly)

build_tools/python_deploy/build_windows_ci.sh

@@ -11,9 +11,8 @@ cmake -GNinja -Bbuild \
   -DLLVM_TARGETS_TO_BUILD=host \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DPython3_FIND_VIRTUALENV=ONLY \
-  -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
+  -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$PWD" \
-  -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="$PWD/externals/llvm-external-projects/torch-mlir-dialects" \
   -DPython3_EXECUTABLE="$(which python)" \
   $GITHUB_WORKSPACE/externals/llvm-project/llvm
 

docs/development.md

@@ -35,9 +35,8 @@ cmake -GNinja -Bbuild \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_FIND_VIRTUALENV=ONLY \
   -DLLVM_ENABLE_PROJECTS=mlir \
-  -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
+  -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$PWD" \
-  -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="$PWD"/externals/llvm-external-projects/torch-mlir-dialects \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_TARGETS_TO_BUILD=host \
   externals/llvm-project/llvm
@@ -360,6 +359,7 @@ Torch-MLIR has two types of tests:
 Alternatively, you can run the tests via Python directly:
 
 ```shell
+cd projects/pt1
 python -m e2e_testing.main -f 'AtenEmbeddingBag'
 ```
 
@@ -374,7 +374,7 @@ ninja check-torch-mlir-all
 This can be broken down into
 
 ```
-ninja check-torch-mlir check-torch-mlir-dialects check-torch-mlir-python
+ninja check-torch-mlir check-torch-mlir-python
 ```
 
 To run more fine-grained tests, you can do, for `check-torch-mlir`: