Make torch-mlir-dialects optional Added IREE_BUILD_TORCH_MLIR_DIALECTS cmake option to opt-out building torch-mlir-dialects. The IREE_HAVE_TORCH_MLIR_DIALECTS preprocessor variable would be defined when the cmake option is turned on. Bazel build is ignored for now.
diff --git a/CMakeLists.txt b/CMakeLists.txt index 78109eb..1c03fa1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -63,6 +63,7 @@ option(IREE_BUILD_EXPERIMENTAL_REMOTING "Builds experimental remoting support." OFF) option(IREE_BUILD_EXPERIMENTAL_WEB_SAMPLES "Builds experimental web samples." OFF) option(IREE_HAL_DRIVER_EXPERIMENTAL_ROCM "Builds the experimental ROCm Backend." OFF) +option(IREE_BUILD_TORCH_MLIR_SUPPORT "Builds support for compiling torch-mlir programs." ON) #------------------------------------------------------------------------------- # Derived flags based on primary options @@ -450,7 +451,9 @@ # Add default external projects. iree_add_llvm_external_project(mlir-iree-dialects MLIR_IREE_DIALECTS ${CMAKE_CURRENT_SOURCE_DIR}/llvm-external-projects/iree-dialects) iree_add_llvm_external_project(mlir-hlo MLIR_HLO ${CMAKE_CURRENT_SOURCE_DIR}/third_party/mlir-hlo) - iree_add_llvm_external_project(torch-mlir-dialects TORCH_MLIR_DIALECTS ${CMAKE_CURRENT_SOURCE_DIR}/third_party/torch-mlir-dialects) + if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + iree_add_llvm_external_project(torch-mlir-dialects TORCH_MLIR_DIALECTS ${CMAKE_CURRENT_SOURCE_DIR}/third_party/torch-mlir-dialects) + endif() add_subdirectory("third_party/llvm-project/llvm" EXCLUDE_FROM_ALL)
diff --git a/docs/developers/get_started/cmake_options_and_variables.md b/docs/developers/get_started/cmake_options_and_variables.md index 6a340bb..5f1b78c 100644 --- a/docs/developers/get_started/cmake_options_and_variables.md +++ b/docs/developers/get_started/cmake_options_and_variables.md
@@ -141,6 +141,10 @@ "${CMAKE_SOURCE_DIR}/integrations/tensorflow/bazel-bin/iree_tf_compiler", which is where they would be placed by a `bazel build` invocation. +#### `IREE_BUILD_TORCH_MLIR_SUPPORT`:BOOL + +Enables building of the torch-mlir-dialects to IREE compiler. Defaults to `ON`. + ## MLIR-specific CMake Options and Variables #### `MLIR_DIR`:STRING
diff --git a/iree/compiler/InputConversion/TMTensor/BUILD b/iree/compiler/InputConversion/TMTensor/BUILD index a8ee733..7760a2e 100644 --- a/iree/compiler/InputConversion/TMTensor/BUILD +++ b/iree/compiler/InputConversion/TMTensor/BUILD
@@ -5,6 +5,7 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") +load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content") package( default_visibility = ["//visibility:public"], @@ -12,6 +13,17 @@ licenses = ["notice"], # Apache 2.0 ) +iree_cmake_extra_content( + content = """ +if(NOT "${IREE_BUILD_TORCH_MLIR_SUPPORT}") + iree_cc_library( + NAME TMTensor + ) + return() +endif() +""", +) + gentbl_cc_library( name = "PassesIncGen", tbl_outs = [ @@ -50,6 +62,9 @@ hdrs = [ "Passes.h", ], + defines = [ + "IREE_HAVE_TORCH_MLIR_DIALECTS", + ], deps = [ ":PassHeaders", ":PassesIncGen",
diff --git a/iree/compiler/InputConversion/TMTensor/CMakeLists.txt b/iree/compiler/InputConversion/TMTensor/CMakeLists.txt index 51de03a..472853c 100644 --- a/iree/compiler/InputConversion/TMTensor/CMakeLists.txt +++ b/iree/compiler/InputConversion/TMTensor/CMakeLists.txt
@@ -8,6 +8,13 @@ # To disable autogeneration for this file entirely, delete this header. # ################################################################################ +if(NOT "${IREE_BUILD_TORCH_MLIR_SUPPORT}") + iree_cc_library( + NAME TMTensor + ) + return() +endif() + iree_add_all_subdirs() iree_tablegen_library( @@ -49,6 +56,8 @@ MLIRPass MLIRTransforms TorchMLIRTMTensorDialect + DEFINES + "IREE_HAVE_TORCH_MLIR_DIALECTS" PUBLIC )
diff --git a/iree/compiler/InputConversion/TMTensor/test/BUILD b/iree/compiler/InputConversion/TMTensor/test/BUILD index 6b4c4b9c..1c4a38a 100644 --- a/iree/compiler/InputConversion/TMTensor/test/BUILD +++ b/iree/compiler/InputConversion/TMTensor/test/BUILD
@@ -8,6 +8,7 @@ load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite") load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob") +load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content") package( default_visibility = ["//visibility:public"], @@ -15,6 +16,14 @@ licenses = ["notice"], # Apache 2.0 ) +iree_cmake_extra_content( + content = """ +if(NOT "${IREE_BUILD_TORCH_MLIR_SUPPORT}") + return() +endif() +""", +) + iree_lit_test_suite( name = "lit", srcs = enforce_glob(
diff --git a/iree/compiler/InputConversion/TMTensor/test/CMakeLists.txt b/iree/compiler/InputConversion/TMTensor/test/CMakeLists.txt index 4d20573..1586cdf 100644 --- a/iree/compiler/InputConversion/TMTensor/test/CMakeLists.txt +++ b/iree/compiler/InputConversion/TMTensor/test/CMakeLists.txt
@@ -8,6 +8,10 @@ # To disable autogeneration for this file entirely, delete this header. # ################################################################################ +if(NOT "${IREE_BUILD_TORCH_MLIR_SUPPORT}") + return() +endif() + iree_add_all_subdirs() iree_lit_test_suite(
diff --git a/iree/compiler/Pipelines/BUILD b/iree/compiler/Pipelines/BUILD index 4eaf898..fb03a81 100644 --- a/iree/compiler/Pipelines/BUILD +++ b/iree/compiler/Pipelines/BUILD
@@ -15,6 +15,7 @@ srcs = ["Options.cpp"], hdrs = ["Options.h"], deps = [ + "//iree/compiler/InputConversion/TMTensor", "//iree/compiler/Utils", ], )
diff --git a/iree/compiler/Pipelines/CMakeLists.txt b/iree/compiler/Pipelines/CMakeLists.txt index 5b1e662..8398e68 100644 --- a/iree/compiler/Pipelines/CMakeLists.txt +++ b/iree/compiler/Pipelines/CMakeLists.txt
@@ -18,6 +18,7 @@ SRCS "Options.cpp" DEPS + iree::compiler::InputConversion::TMTensor iree::compiler::Utils PUBLIC )
diff --git a/iree/compiler/Pipelines/Options.cpp b/iree/compiler/Pipelines/Options.cpp index 5623ee5..b5b5a2b 100644 --- a/iree/compiler/Pipelines/Options.cpp +++ b/iree/compiler/Pipelines/Options.cpp
@@ -38,8 +38,10 @@ "Legalize from TOSA ops."), clEnumValN(InputDialectOptions::Type::mhlo, "mhlo", "Legalize from MHLO ops."), +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS clEnumValN(InputDialectOptions::Type::tm_tensor, "tm_tensor", "Legalize from TMTensor ops."), +#endif clEnumValN( InputDialectOptions::Type::xla, "xla", "Legalize from MHLO ops (with XLA cleanup preprocessing).")),
diff --git a/iree/compiler/Pipelines/Options.h b/iree/compiler/Pipelines/Options.h index 495228e..6f784e7 100644 --- a/iree/compiler/Pipelines/Options.h +++ b/iree/compiler/Pipelines/Options.h
@@ -38,8 +38,10 @@ tosa, // Legalizes input defined over MHLO ops. mhlo, +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS // Legalizes input defined over TMTensor ops. tm_tensor, +#endif // Special case of 'mhlo' legalization which also performs some XLA // cleanup activities. xla,
diff --git a/iree/compiler/Pipelines/Pipelines.cpp b/iree/compiler/Pipelines/Pipelines.cpp index 446e57b..001ea3b 100644 --- a/iree/compiler/Pipelines/Pipelines.cpp +++ b/iree/compiler/Pipelines/Pipelines.cpp
@@ -15,7 +15,9 @@ #include "iree/compiler/Dialect/VM/Transforms/Passes.h" #include "iree/compiler/InputConversion/Common/Passes.h" #include "iree/compiler/InputConversion/MHLO/Passes.h" +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS #include "iree/compiler/InputConversion/TMTensor/Passes.h" +#endif #include "iree/compiler/InputConversion/TOSA/Passes.h" namespace mlir { @@ -41,10 +43,12 @@ case InputDialectOptions::Type::mhlo: MHLO::buildMHLOInputConversionPassPipeline(passManager); break; +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS case InputDialectOptions::Type::tm_tensor: passManager.addNestedPass<func::FuncOp>( TMTensor::createConvertTMTensorToLinalgExtPass()); break; +#endif case InputDialectOptions::Type::xla: MHLO::buildXLACleanupPassPipeline(passManager); MHLO::buildMHLOInputConversionPassPipeline(passManager);
diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt index ce127de..6c3acca 100644 --- a/iree/tools/CMakeLists.txt +++ b/iree/tools/CMakeLists.txt
@@ -302,6 +302,10 @@ PUBLIC ) + if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + set(_OPTIONAL_INIT_TORCH_MLIR_DIALECT_LIB ::init_torch_mlir_dialects) + endif() + iree_cc_library( NAME init_passes_and_dialects @@ -312,7 +316,7 @@ ::init_compiler_modules ::init_iree_passes_and_dialects ::init_mlir_passes_and_dialects - ::init_torch_mlir_dialects + ${_OPTIONAL_INIT_TORCH_MLIR_DIALECT_LIB} ::init_xla_dialects iree::compiler::Codegen::Codegen iree::compiler::Dialect::HAL::Conversion::Passes @@ -365,15 +369,19 @@ PUBLIC ) - iree_cc_library( - NAME - init_torch_mlir_dialects - HDRS - "init_torch_mlir_dialects.h" - DEPS - TorchMLIRTMTensorDialect - PUBLIC - ) + if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + iree_cc_library( + NAME + init_torch_mlir_dialects + HDRS + "init_torch_mlir_dialects.h" + DEPS + TorchMLIRTMTensorDialect + DEFINES + "IREE_HAVE_TORCH_MLIR_DIALECTS" + PUBLIC + ) + endif() iree_cc_library( NAME
diff --git a/iree/tools/init_dialects.h b/iree/tools/init_dialects.h index a8dbee8..5be36b5 100644 --- a/iree/tools/init_dialects.h +++ b/iree/tools/init_dialects.h
@@ -15,7 +15,9 @@ #include "iree/tools/init_compiler_modules.h" #include "iree/tools/init_iree_dialects.h" #include "iree/tools/init_mlir_dialects.h" +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS #include "iree/tools/init_torch_mlir_dialects.h" +#endif #include "iree/tools/init_xla_dialects.h" namespace mlir { @@ -24,7 +26,9 @@ inline void registerAllDialects(DialectRegistry ®istry) { registerMlirDialects(registry); registerIreeDialects(registry); +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS registerTorchMLIRDialects(registry); +#endif mlir::registerXLADialects(registry); mlir::iree_compiler::registerIreeCompilerModuleDialects(registry); }
diff --git a/iree/tools/init_iree_passes.h b/iree/tools/init_iree_passes.h index 3cb1078..68c4075 100644 --- a/iree/tools/init_iree_passes.h +++ b/iree/tools/init_iree_passes.h
@@ -27,7 +27,9 @@ #include "iree/compiler/Dialect/VM/Transforms/Passes.h" #include "iree/compiler/InputConversion/Common/Passes.h" #include "iree/compiler/InputConversion/MHLO/Passes.h" +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS #include "iree/compiler/InputConversion/TMTensor/Passes.h" +#endif #include "iree/compiler/InputConversion/TOSA/Passes.h" #include "iree/compiler/Translation/IREEVM.h" @@ -44,7 +46,9 @@ registerCommonInputConversionPasses(); MHLO::registerMHLOConversionPasses(); +#ifdef IREE_HAVE_TORCH_MLIR_DIALECTS TMTensor::registerTMTensorConversionPasses(); +#endif registerTOSAConversionPasses(); ConstEval::registerConstEvalPasses();
diff --git a/llvm-external-projects/iree-compiler-api/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/CMakeLists.txt index 95cae0e..86cdf40 100644 --- a/llvm-external-projects/iree-compiler-api/CMakeLists.txt +++ b/llvm-external-projects/iree-compiler-api/CMakeLists.txt
@@ -42,7 +42,9 @@ set(LLVM_MAIN_SRC_DIR "${IREE_COMPILER_API_SOURCE_DIR}/../../third_party/llvm-project/llvm") set(LLVM_MAIN_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/llvm") set(LLVM_EXTERNAL_MLIR_HLO_SOURCE_DIR "${IREE_COMPILER_API_SOURCE_DIR}/../../third_party/mlir-hlo") - set(LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR "${IREE_COMPILER_API_SOURCE_DIR}/../../third_party/torch-mlir-dialects") + if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + set(LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR "${IREE_COMPILER_API_SOURCE_DIR}/../../third_party/torch-mlir-dialects") + endif() # Resources generated on windows must have valid version numbers. # See set_windows_version_resource_properties. set(LLVM_VERSION_MAJOR 0) @@ -101,8 +103,10 @@ add_system_include_hack(${LLVM_MAIN_BINARY_DIR}/tools/iree-dialects/include) add_system_include_hack(${LLVM_EXTERNAL_MLIR_HLO_SOURCE_DIR}/include) add_system_include_hack(${LLVM_MAIN_BINARY_DIR}/tools/mlir-hlo/include) -add_system_include_hack(${LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR}/include) -add_system_include_hack(${LLVM_MAIN_BINARY_DIR}/tools/torch-mlir-dialects/include) +if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + add_system_include_hack(${LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR}/include) + add_system_include_hack(${LLVM_MAIN_BINARY_DIR}/tools/torch-mlir-dialects/include) +endif() add_system_include_hack(${IREE_COMPILER_API_SOURCE_DIR}/include) function(iree_compiler_target_includes target) @@ -127,7 +131,9 @@ iree_set_llvm_cmake_options() iree_add_llvm_external_project(mlir-iree-dialects MLIR_IREE_DIALECTS ${LLVM_EXTERNAL_MLIR_IREE_DIALECTS_SOURCE_DIR}) iree_add_llvm_external_project(mlir-hlo MLIR_HLO ${LLVM_EXTERNAL_MLIR_HLO_SOURCE_DIR}) - iree_add_llvm_external_project(torch-mlir-dialects TORCH_MLIR_DIALECTS ${LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR}) + if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + iree_add_llvm_external_project(torch-mlir-dialects TORCH_MLIR_DIALECTS ${LLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR}) + endif() set(MLIR_ENABLE_BINDINGS_PYTHON ON) # Required mlir-hlo settings.
diff --git a/llvm-external-projects/iree-compiler-api/lib/CAPI/CMakeLists.txt b/llvm-external-projects/iree-compiler-api/lib/CAPI/CMakeLists.txt index 7a103fc..412ce7a 100644 --- a/llvm-external-projects/iree-compiler-api/lib/CAPI/CMakeLists.txt +++ b/llvm-external-projects/iree-compiler-api/lib/CAPI/CMakeLists.txt
@@ -45,6 +45,13 @@ ${_OPTIONAL_LINK_LIBS} ) +if(${IREE_BUILD_TORCH_MLIR_SUPPORT}) + target_compile_definitions(obj.IREECompilerAPICompilerCAPI + PUBLIC + -DIREE_HAVE_TORCH_MLIR_DIALECTS=1 + ) +endif() + # TODO: Fix upstream so there is a way to know what the actual compile target # is (versus prefixing with "obj." which is conditional). iree_compiler_target_includes(obj.IREECompilerAPICompilerCAPI)