Breakup python pytorch deps (#2582)

This lifts the core of the jit_ir_importer and ltc out of the pt1 project, making them peers to it. As a side-effect of this layering, now the "MLIR bits" (dialects, etc) are not commingled with the various parts of the pt1 project, allowing pt1 and ltc to overlay cleanly onto a more fundamental "just MLIR" Python core. Prior to this, the Python namespace was polluted to the point that this could not happen. That "just MLIR" Python core will be introduced in a followup, which will create the space to upstream the FX and ONNX pure Python importers. This primary non-NFC change to the API is: * `torch_mlir.dialects.torch.importer.jit_ir` -> `torch_mlir.jit_ir_importer`. The rest is source code layering so that we can make the pt1 project optional without losing the other features. Progress on #2546.
llvm · Nov 19, 2023 · 5eae0ad · 5eae0ad
1 parent facbe5d
commit 5eae0ad
Show file tree

Hide file tree

Showing 143 changed files with 455 additions and 483 deletions.
diff --git a/.github/workflows/buildAndTest.yml b/.github/workflows/buildAndTest.yml
@@ -27,15 +27,12 @@ jobs:
       matrix:
         os-arch: [ubuntu-x86_64, macos-arm64, windows-x86_64]
         llvm-build: [in-tree, out-of-tree]
-        torch-binary: [ON, OFF]
+        torch-binary: [ON]
         torch-version: [nightly, stable]
         exclude:
-          # Exclude llvm in-tree and pytorch source
-          - llvm-build: in-tree
-            torch-binary: OFF
-          # Exclude llvm out-of-tree and pytorch binary
+          # Exclude llvm out-of-tree and pytorch stable (to save resources)
           - llvm-build: out-of-tree
-            torch-binary: ON
+            torch-version: stable
           # Exclude macos-arm64 and llvm out-of-tree altogether
           - os-arch: macos-arm64
             llvm-build: out-of-tree
@@ -45,9 +42,6 @@ jobs:
             llvm-build: out-of-tree
           - os-arch: windows-x86_64
             torch-version: stable
-          # For PyTorch stable builds, we don't build PyTorch from source
-          - torch-version: stable
-            torch-binary: OFF
         include:
           # Specify OS versions
           - os-arch: ubuntu-x86_64

diff --git a/.gitignore b/.gitignore
@@ -26,7 +26,7 @@ __pycache__
 bazel-*
 
 # Autogenerated files
-/projects/pt1/python/torch_mlir/csrc/base_lazy_backend/generated
+/projects/ltc/csrc/base_lazy_backend/generated
 
 #Docker builds
 build_oot/

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -149,10 +149,12 @@ endfunction()
 # Configure CMake.
 list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules)
 list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake)
+list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake)
 
 include(TableGen)
 include(AddLLVM)
 include(AddMLIR)
+include(AddMLIRPython)
 
 ################################################################################
 # Setup python.
@@ -231,6 +233,4 @@ endif()
 # Sub-projects
 #-------------------------------------------------------------------------------
 
-if(TORCH_MLIR_ENABLE_PROJECT_PT1)
-  add_subdirectory(projects/pt1)
-endif()
+add_subdirectory(projects)
diff --git a/build_tools/autogen_ltc_backend.py b/build_tools/autogen_ltc_backend.py
@@ -29,7 +29,6 @@
     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 +113,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_PT1_DIR.joinpath(
-            "python", "torch_mlir", "csrc", "base_lazy_backend"
+        self.backend_path = TORCH_MLIR_DIR.joinpath(
+            "projects", "ltc", "csrc", "base_lazy_backend"
         )
         assert self.backend_path.is_dir(), f"Backend path not found: {self.backend_path}"
         self.generated_path = self.binary_dir.joinpath(
-            "projects", "pt1", "python", "torch_mlir", "csrc", "base_lazy_backend", "generated"
+            "projects", "ltc", "csrc", "base_lazy_backend", "generated"
         )
         self.generated_path.mkdir(parents=True, exist_ok=True)
 
@@ -415,7 +414,7 @@ def extract_signatures(text):
                     // for ops that dont have a corresponding structured kernel or shape definition
 
                     #include "shape_inference.h"
-                    #include "torch_mlir/csrc/base_lazy_backend/utils/exception.h"
+                    #include "base_lazy_backend/utils/exception.h"
                     namespace torch {{
                     namespace lazy {{
                     {}
@@ -467,7 +466,7 @@ def gen_fallback_code(*args, **kwargs):
             node_base="torch::lazy::TorchMlirNode",
             node_base_hdr=str(self.backend_path.joinpath("mlir_node.h")),
             tensor_class=self.tensor_class,
-            tensor_class_hdr="torch_mlir/csrc/base_lazy_backend/tensor.h",
+            tensor_class_hdr="base_lazy_backend/tensor.h",
             create_aten_from_ltc_tensor="CreateFunctionalizedAtenFromLtcTensor",
             shape_inference_hdr=str(self.generated_path.joinpath("shape_inference.h")),
             lazy_ir_generator=GenMlirLazyIr,

diff --git a/...mlir/cmake/modules/TorchMLIRPyTorch.cmake → build_tools/cmake/TorchMLIRPyTorch.cmake b/...mlir/cmake/modules/TorchMLIRPyTorch.cmake → build_tools/cmake/TorchMLIRPyTorch.cmake
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
@@ -364,9 +364,9 @@ function setup_venv() {
 function build_out_of_tree() {
   local torch_from_bin="$1"
   local python_version="$2"
-  echo ":::: Build out-of-tree Torch from binary: $torch_from_bin with Python: $python_version"
-
   local torch_version="$3"
+  echo ":::: Build out-of-tree Torch from binary: $torch_from_bin with Python: $python_version ($torch_version)"
+
   local enable_ltc="ON"
   if [[ "${torch_version}" == "stable" ]]
   then

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
@@ -42,6 +42,6 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
 fi
 
 PYTHONPATH="${pypath}" python \
-  -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.abstract_interp_lib_gen \
+  -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
   --pytorch_op_extensions=${ext_module:-""} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
@@ -43,7 +43,7 @@ fi
 
 set +u
 PYTHONPATH="${PYTHONPATH}:${pypath}" python \
-  -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen \
+  -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \
   --debug_registry_dump="${torch_ir_include_dir}/JITOperatorRegistryDump.txt"
diff --git a/docs/Torch-ops-E2E-implementation.md b/docs/Torch-ops-E2E-implementation.md
@@ -17,7 +17,7 @@ The end-to-end test is important to check the correctness of the other steps.
 
 ### Step 2. Update ods
 
-Update [torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/main/projects/pt1/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py) with the new op and run [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/update_torch_ods.sh) to generate the ods. Running `update_torch_ods.sh` would dump all the operators with schema into `JITOperatorRegistryDump.txt`. It’s convenient to look for ops signatures and operands names in this file.
+Update [torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/main/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py) with the new op and run [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/update_torch_ods.sh) to generate the ods. Running `update_torch_ods.sh` would dump all the operators with schema into `JITOperatorRegistryDump.txt`. It’s convenient to look for ops signatures and operands names in this file.
 
 ### Step 3. Propagate types
 It’s essential to make sure the new op implements shape and dtype inference. See [abstract_interp_lib](https://github.com/llvm/torch-mlir/blob/main/docs/abstract_interp_lib.md) for information on adding shape and dtype inference.

diff --git a/docs/abstract_interp_lib.md b/docs/abstract_interp_lib.md
@@ -26,7 +26,7 @@ The two main use cases are:
 ## Architecture
 
 Functions are defined as TorchScript-able Python functions in
-`python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/abstract_interp_lib_gen.py`.
+`python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py`.
 The signatures of the functions are systematically derived from Torch JIT
 operator registry. Most shape functions are expected to reuse the upstream
 helper functions

diff --git a/docs/adding_an_e2e_test.md b/docs/adding_an_e2e_test.md
@@ -87,7 +87,7 @@ following order:
 
 1. Shape of input tensor. Use `-1` for dynamic dimensions
 2. Dtype of the input tensor
-3. Boolean representing whether the input tensor [has value semantics](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/class_annotator.h#L54-L67). This
+3. Boolean representing whether the input tensor [has value semantics](https://github.com/llvm/torch-mlir/blob/ba17a4d6c09b4bbb4ef21b1d8d4a93cb056be109/python/torch_mlir/jit_ir_importer/csrc/class_annotator.h#L54-L67). This
    will always be true for E2E tests, since the [Torch-MLIR backend contract](architecture.md#the-backend-contract) requires all tensors in the
    IR to eventually have value semantics.
 

diff --git a/docs/architecture.md b/docs/architecture.md
@@ -55,22 +55,22 @@ factored such that we can handle this with one core import path, which is
 through the PyTorch
 "[JIT IR](https://github.com/pytorch/pytorch/blob/78c8a0d75220bdd4955415b5f81509e005af4232/torch/csrc/jit/OVERVIEW.md)",
 and lives in
-[torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir](https://github.com/llvm/torch-mlir/tree/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/dialects/torch/importer/jit_ir).
+[torch-mlir/python/torch_mlir/jit_ir_importer](https://github.com/llvm/torch-mlir/tree/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/dialects/torch/importer/jit_ir).
 The JIT IR is a highly principled IR that faithfully models a Python subset (+
 tensors, the PyTorch op registry, and a few other things). All the other PyTorch
 program representations can eventually bottom-out on the JIT IR via some path
 provided by PyTorch. The `torch` dialect is almost entirely in 1:1
 correspondence with the JIT IR -- this allows the importer to be extremely small
 (the core is
-[under 500 lines of code](https://github.com/llvm/torch-mlir/blob/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp#L1)).
+[under 500 lines of code](https://github.com/llvm/torch-mlir/blob/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/jit_ir_importer/csrc/node_importer.cpp#L1)).
 
 ### Ops
 
 See [TorchOps.td](https://github.com/llvm/torch-mlir/blob/114f48e96c578ee76a6f83b3aa4aa229a8d5b76e/include/torch-mlir/Dialect/Torch/IR/TorchOps.td#L1)
 
 The ops in the `torch` dialect are almost entirely generated based on the
 PyTorch JIT IR operator registry via the script
-[torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py#L1) (invoked via [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/update_torch_ods.sh)).
+[torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/e322f6a8784009b37aa354abfa9a40a80f30877d/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py#L1) (invoked via [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/update_torch_ods.sh)).
 This script queries the registry and generates MLIR
 [ODS](https://mlir.llvm.org/docs/OpDefinitions/) in
 [GeneratedTorchOps.td](https://github.com/llvm/torch-mlir/blob/e322f6a8784009b37aa354abfa9a40a80f30877d/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td#L1). We have a guide for [adding a new op end-to-end](https://github.com/llvm/torch-mlir/wiki/Torch-ops-E2E-implementation).
@@ -195,7 +195,7 @@ values. When one `torch.jit.script`'s a `torch.nn.Module`, the result is
 actually an `IValue` that represents the module, with a hierarchy of children
 `IValue`'s. Strictly speaking, JIT IR `torch::jit::Graph`'s are only used to
 represent the bodies of methods on the modules. So in addition to importing the
-JIT IR, we also need to import the `IValue`'s. This happens inside [ivalue_importer.cpp](https://github.com/llvm/torch-mlir/blob/fde390c7669e29362b18388448ef2b188713383f/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/ivalue_importer.cpp#L1).
+JIT IR, we also need to import the `IValue`'s. This happens inside [ivalue_importer.cpp](https://github.com/llvm/torch-mlir/blob/fde390c7669e29362b18388448ef2b188713383f/python/torch_mlir/jit_ir_importer/csrc/ivalue_importer.cpp#L1).
 
 Most of the IValue modeling can reuse `torch` dialect ops that already exist
 otherwise, such as `torch.constant.int` to represent an int in the object graph.

diff --git a/docs/ltc_backend.md b/docs/ltc_backend.md
@@ -12,7 +12,7 @@
 [Lazy Tensor Core](https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/tutorial.md) is a tracing system in PyTorch which is supported as an entry point to Torch-MLIR.
 After registering an LTC backend, all operations performed on lazy tensors are recorded and handed off to the backend implementation.
 
-LTC support is provided through an abstract [`TorchMlirBackendImpl`](../python/torch_mlir/csrc/base_lazy_backend/backend_impl.h) class, which handles the conversion to MLIR.
+LTC support is provided through an abstract [`TorchMlirBackendImpl`](../projects/ltc/csrc/base_lazy_backend/backend_impl.h) class, which handles the conversion to MLIR.
 Implementations based on this abstract class will be able to specify their own compile and execution workflows.
 Additional details about how to implement a custom backend is available [below](#Implementing-a-custom-backend).
 
@@ -27,7 +27,7 @@ View examples [here](ltc_examples.md).
   - The [autogen files](#autogen-files) are generated by this script based on the list of supported ops, which includes all ops from [`GeneratedTorchOps.td`](https://github.com/llvm/torch-mlir/blob/main/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td),
     excluding those explicitly blacklisted in the YAML file
 
-### Autogen Files ([`python/torch_mlir/csrc/base_lazy_backend/generated`](../python/torch_mlir/csrc/base_lazy_backend/generated))
+### Autogen Files ([`projects/ltc/csrc/base_lazy_backend/generated`](../projects/ltc/csrc/base_lazy_backend/generated))
 Generated files are created in this directory, which is ignored by version control.
 
 - `LazyIr.h`
@@ -41,7 +41,7 @@ Generated files are created in this directory, which is ignored by version contr
 - `shape_inference.{cpp,h}`
   - Shape inference headers for supported ops and autogen'd placeholders for unimplemented functions
 
-### Base Backend ([`python/torch_mlir/csrc/base_lazy_backend`](../python/torch_mlir/csrc/base_lazy_backend))
+### Base Backend ([`projects/ltc/csrc/base_lazy_backend`](../projects/ltc/csrc/base_lazy_backend))
 
 - `backend_impl.{cpp,h}`
   - Base LTC backend to setup Torch-MLIR lowering context

diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -13,7 +13,7 @@
 // This file is automatically generated.  Please do not edit.
 // Generated via:
 // ```
-// python -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen
+// python -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen
 // ```
 //
 //===----------------------------------------------------------------------===//