diff --git a/CMakeLists.txt b/CMakeLists.txt index a6d498013..790fcfeb5 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/build-requirements.txt b/build-requirements.txt index 16f7da636..db0c32051 100644 --- a/build-requirements.txt +++ b/build-requirements.txt @@ -1,5 +1,3 @@ --r pytorch-requirements.txt - numpy pybind11 wheel diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 62f9b9ec4..0ba37a8ec 100644 --- a/python/CMakeLists.txt +++ b/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. diff --git a/requirements.txt b/requirements.txt index 01478ee1f..f47a87da7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ -r build-requirements.txt +-r pytorch-requirements.txt # Test Requirements pillow diff --git a/setup.py b/setup.py index 9ba763351..049224bfd 100644 --- a/setup.py +++ b/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, ) diff --git a/test/python/smoketest.py b/test/python/smoketest.py index 88e0a10f7..bb97927e9 100644 --- a/test/python/smoketest.py +++ b/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() \ No newline at end of file diff --git a/whl-requirements.txt b/whl-requirements.txt index 554744257..f628a4180 100644 --- a/whl-requirements.txt +++ b/whl-requirements.txt @@ -1,4 +1,5 @@ -f build-requirements.txt +-f pytorch-requirements.txt # Packaging requirements. packaging