diff --git a/CMakeLists.txt b/CMakeLists.txt index 181cf8b8d944..5b5f95ef71e9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -51,6 +51,9 @@ option(TORCH_MLIR_OUT_OF_TREE_BUILD "Specifies an out of tree build" OFF) # But it will be manually enabled in CI build to enable the jit_ir_importer.build_tools.torch_ods_gen # and abstract_interp_lib_gen.py. Once pure python version of build_tools finished, no need to set it in CI. option(TORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS "Enables PyTorch native extension features" OFF) +if(TORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS) + add_definitions(-DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS) +endif() # NOTE: The JIT_IR_IMPORTER paths have become unsupportable due to age and lack of maintainers. # Turning this off disables the old TorchScript path, leaving FX based import as the current supported option. # The option will be retained for a time, and if a maintainer is interested in setting up testing for it, diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 4fbd8561dcd3..6eb47b51476a 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -96,16 +96,18 @@ set(_source_components TorchMLIRPythonSources TorchMLIRPythonExtensions TorchMLIRSiteInitialize - - # Sources related to optional Torch extension dependent features. Typically - # empty unless if project features are enabled. - TorchMLIRPythonTorchExtensionsSources ) if(TORCH_MLIR_ENABLE_STABLEHLO) list(APPEND _source_components StablehloPythonExtensions) endif() +# Sources related to optional Torch extension dependent features. Typically +# empty unless if project features are enabled. +if(TORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS) + list(APPEND _source_components TorchMLIRPythonTorchExtensionsSources) +endif() + add_mlir_python_common_capi_library(TorchMLIRAggregateCAPI INSTALL_COMPONENT TorchMLIRPythonModules INSTALL_DESTINATION python_packages/torch_mlir/torch_mlir/_mlir_libs diff --git a/setup.py b/setup.py index 6f5f5d5d1c3b..71491affb988 100644 --- a/setup.py +++ b/setup.py @@ -75,7 +75,7 @@ def _check_env_flag(name: str, default=None) -> bool: # If true, enable LTC build by default TORCH_MLIR_ENABLE_LTC = _check_env_flag("TORCH_MLIR_ENABLE_LTC", True) TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = _check_env_flag( - "TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS", False + "TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS", True ) LLVM_INSTALL_DIR = os.getenv("LLVM_INSTALL_DIR", None) SRC_DIR = pathlib.Path(__file__).parent.absolute()