WIP: No PyTorch dep (#1854)

CMakeLists.txt

@@ -41,7 +41,13 @@ if(TORCH_MLIR_ENABLE_STABLEHLO)
   add_definitions(-DTORCH_MLIR_ENABLE_STABLEHLO)
 endif()
 
+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)
@@ -109,7 +115,6 @@ 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)
-  option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 
   set(TORCH-MLIR_BUILT_STANDALONE 1)
   set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}")
@@ -119,7 +124,6 @@ else()
   # In-tree build with LLVM_EXTERNAL_PROJECTS=torch-mlir
 
   option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF)
-  option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 
   # TODO: Fix this upstream so that global include directories are not needed.
   set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir)

build-requirements.txt

@@ -1,5 +1,3 @@
--r pytorch-requirements.txt
-
 numpy
 pybind11
 wheel

python/CMakeLists.txt

@@ -45,14 +45,16 @@ endif()
 declare_mlir_python_sources(TorchMLIRPythonSources)
 declare_mlir_python_sources(TorchMLIRPythonExtensions)
 
-declare_mlir_python_sources(TorchMLIRPythonSources.TopLevel
-  ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
-  ADD_TO_PARENT TorchMLIRPythonSources
-  SOURCES
-    __init__.py
-    compiler_utils.py
-    dynamo.py
-)
+if (NOT TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS)
+  declare_mlir_python_sources(TorchMLIRPythonSources.TopLevel
+    ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
+    ADD_TO_PARENT TorchMLIRPythonSources
+    SOURCES
+      __init__.py
+      compiler_utils.py
+      dynamo.py
+  )
+endif()
 
 declare_mlir_python_sources(TorchMLIRPythonSources.Dialects
   ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
@@ -91,7 +93,9 @@ if(TORCH_MLIR_ENABLE_LTC)
 endif()
 # Reference backend has a separate check for TORCH_MLIR_ENABLE_LTC, since it
 # generates a dummy Python library when disabled.
-add_subdirectory(torch_mlir/csrc/reference_lazy_backend)
+if(NOT TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS)
+  add_subdirectory(torch_mlir/csrc/reference_lazy_backend)
+endif()
 
 ################################################################################
 # Optionally handle JIT IR importer.

requirements.txt

@@ -1,4 +1,5 @@
 -r build-requirements.txt
+-r pytorch-requirements.txt
 
 # Test Requirements
 pillow

setup.py

@@ -41,12 +41,14 @@ from setuptools import setup, Extension
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
-import torch
 
 PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION") or "0.0.1"
 
 # If true, enable LTC build by default
 TORCH_MLIR_ENABLE_LTC_DEFAULT = True
+TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = int(os.environ.get('TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS', False))
+if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS:
+    import torch
 
 # Build phase discovery is unreliable. Just tell it what phases to run.
 class CustomBuild(_build):
@@ -90,6 +92,7 @@ class CMakeBuild(build_py):
                 f"-DCMAKE_C_VISIBILITY_PRESET=hidden",
                 f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden",
                 f"-DTORCH_MLIR_ENABLE_LTC={'ON' if enable_ltc else 'OFF'}",
+                f"-DTORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS={'ON' if TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else 'OFF'}",
             ]
 
             os.makedirs(cmake_build_dir, exist_ok=True)
@@ -158,10 +161,8 @@ setup(
     },
     ext_modules=[
         CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
-    ],
-    install_requires=[
-        "numpy",
-        f"torch=={torch.__version__}".split("+", 1)[0],
-    ],
+    ] if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else [CMakeExtension("_")],
+    install_requires=["numpy", ] + (
+        [f"torch=={torch.__version__}".split("+", 1)[0], ] if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else []),
     zip_safe=False,
 )

test/python/smoketest.py

@@ -5,3 +5,8 @@ from torch_mlir.dialects import torch
 
 with torch_mlir.ir.Context() as ctx:
   torch.register_dialect(ctx)
+  with torch_mlir.ir.Location.unknown() as loc:
+    module = torch_mlir.ir.Module.create(loc)
+    with torch_mlir.ir.InsertionPoint.at_block_begin(module.body):
+      n = torch.ConstantNoneOp()
+    module.operation.print()

whl-requirements.txt

@@ -1,4 +1,5 @@
 -f build-requirements.txt
+-f pytorch-requirements.txt
 
 # Packaging requirements.
 packaging