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)
