From f9109bed8b8cd169a980331dbabe30ab1d7752c0 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Fri, 14 Aug 2020 13:11:03 +0800 Subject: [PATCH] feat(mge): add mlir cmake GitOrigin-RevId: e86db34b18d75c85f7ef120481f79359dc3e1da2 --- CMakeLists.txt | 38 +++++- cmake/llvm-project.cmake | 123 ++++++++++++++++++ src/CMakeLists.txt | 18 ++- src/jit/impl/mlir/compiler.cpp | 1 - src/jit/impl/mlir/ir/CMakeLists.txt | 31 +++++ src/jit/impl/mlir/ir/common.cpp | 2 +- src/jit/impl/mlir/ir/common.h | 2 +- src/jit/impl/mlir/ir/dialect.cpp | 2 +- src/jit/impl/mlir/ir/lower_to_affine_pass.cpp | 2 +- src/jit/impl/mlir/ir/lower_to_gpu_pass.cpp | 2 +- src/jit/impl/mlir/ir/lower_to_llvm_pass.cpp | 2 +- src/jit/impl/mlir/ir/shape_inference_pass.cpp | 2 +- src/jit/impl/mlir/mlir_gen.cpp | 3 +- src/jit/impl/mlir/mlir_gen.h | 3 +- .../include/megbrain/jit/mlir/ir/dialect.h | 2 +- src/jit/include/megbrain/jit/mlir/ir/passes.h | 3 +- .../jit/mlir/ir/shape_inference_interface.h | 2 +- src/jit/test/fusion.cpp | 2 + src/megbrain_build_config.h.in | 1 + 19 files changed, 219 insertions(+), 22 deletions(-) create mode 100644 cmake/llvm-project.cmake create mode 100644 src/jit/impl/mlir/ir/CMakeLists.txt diff --git a/CMakeLists.txt b/CMakeLists.txt index 1404220a2..b1b42d614 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -29,6 +29,7 @@ set_property(CACHE MGE_ARCH PROPERTY STRINGS AUTO set (MGE_EXPORT_TARGETS MegEngine-targets) option(MGE_WITH_JIT "Build MegEngine with JIT." ON) +option(MGE_WITH_JIT_MLIR "Build MegEngine with MLIR JIT." OFF) option(MGE_WITH_HALIDE "Build MegEngine with Halide JIT" ON) option(MGE_WITH_MIDOUT_PROFILE "Build MegEngine with Midout profile." OFF) option(MGE_WITH_MINIMUM_SIZE "Swith off MGE_ENABLE_RTTI、MGE_ENABLE_EXCEPTIONS、MGE_ENABLE_LOGGING and switch on MGE_INFERENCE_ONLY so that compile minimum load_and_run. Take effect only when MGE_BIN_REDUCE was set" OFF) @@ -176,9 +177,10 @@ if(MSVC OR WIN32) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${WIN_FLAGS}") #FIXME: fix halide JIT on windows - message("-- disable jit and halide on windows host build...") + message("-- disable jit, halide and mlir on windows host build...") set(MGE_WITH_HALIDE OFF) set(MGE_WITH_JIT OFF) + set(MGE_WITH_JIT_MLIR OFF) #FIXME: fix MegRay on windows message("-- Disable distributed build on windows host build...") set(MGE_WITH_DISTRIBUTED OFF) @@ -209,6 +211,19 @@ if(NOT MGE_WITH_JIT) message(WARNING "MGE_WITH_HALIDE is set to OFF with MGE_WITH_JIT disabled") set(MGE_WITH_HALIDE OFF) endif() + if(MGE_WITH_JIT_MLIR) + message(WARNING "MGE_WITH_JIT_MLIR is set to OFF with MGE_WITH_JIT disabled") + set(MGE_WITH_JIT_MLIR OFF) + endif() +endif() + +# FIXME At present, there are some conflicts between the LLVM that halide +# depends on and the LLVM that MLIR depends on. Should be fixed in subsequent +# versions. +if(MGE_WITH_JIT_MLIR) + if(MGE_WITH_HALIDE) + message(FATAL_ERROR "please set MGE_WITH_HALIDE to OFF with MGE_WITH_JIT_MLIR enabled") + endif() endif() if(MGE_WITH_CUDA) @@ -223,8 +238,11 @@ if(MGE_WITH_CUDA) endif() if(NOT MGE_WITH_CUDA) - message("-- Disable JIT support, as CUDA is not enabled.") - set(MGE_WITH_JIT OFF) + if(NOT MGE_ARCH STREQUAL "x86_64" AND NOT MGE_ARCH STREQUAL "i386") + message("-- Disable JIT support, as the MGE_ARCH is not X86 and CUDA is not enabled.") + set(MGE_WITH_JIT OFF) + set(MGE_WITH_JIT_MLIR OFF) + endif() set(MGE_WITH_HALIDE OFF) message("-- Disable TensorRT support, as CUDA is not enabled.") set(MGE_WITH_TRT OFF) @@ -283,6 +301,10 @@ if(MGE_INFERENCE_ONLY) set(MGE_BUILD_IMPERATIVE_RT OFF) endif() +if(MGE_WITH_JIT_MLIR) + include(cmake/llvm-project.cmake) +endif() + if(MGE_WITH_DISTRIBUTED) include(cmake/protobuf.cmake) include(cmake/zmq.cmake) @@ -532,8 +554,6 @@ if(MGE_WITH_JIT AND MGE_WITH_HALIDE) set(HALIDE_SHARED_LIBRARY OFF CACHE BOOL "Build as a shared library") include(cmake/Halide.cmake) endif() -set(MGB_JIT ${MGE_WITH_JIT}) -set(MGB_JIT_HALIDE ${MGE_WITH_HALIDE}) # Thread IF(APPLE) @@ -542,8 +562,16 @@ IF(APPLE) set(CMAKE_USE_WIN32_THREADS_INIT 0) set(CMAKE_USE_PTHREADS_INIT 1) set(THREADS_PREFER_PTHREAD_FLAG ON) + message("-- disable jit, halide and mlir on macos host build...") + set(MGE_WITH_HALIDE OFF) + set(MGE_WITH_JIT OFF) + set(MGE_WITH_JIT_MLIR OFF) ENDIF() +set(MGB_JIT ${MGE_WITH_JIT}) +set(MGB_JIT_MLIR ${MGE_WITH_JIT_MLIR}) +set(MGB_JIT_HALIDE ${MGE_WITH_HALIDE}) + if(MSVC OR WIN32) set(CMAKE_HAVE_THREADS_LIBRARY 1) set(CMAKE_USE_WIN32_THREADS_INIT 1) diff --git a/cmake/llvm-project.cmake b/cmake/llvm-project.cmake new file mode 100644 index 000000000..9bbb3c230 --- /dev/null +++ b/cmake/llvm-project.cmake @@ -0,0 +1,123 @@ +# - Find the llvm/mlir libraries +# This module finds if llvm/mlir is installed, or build llvm/mlir from source. +# This module sets the following variables. +# +# MLIR_LLVM_INCLUDE_DIR - path to the LLVM/MLIR include files +# MLIR_LLVM_LIBS - path to the LLVM/MLIR libraries +# +# This module define the following functions. +# +# external_tablegen_library - created interface library which depends on tablegen outputs + +include(CMakeParseArguments) + +function(external_tablegen_library) + cmake_parse_arguments( + _RULE + "TESTONLY" + "NAME;TBLGEN" + "SRCS;INCLUDES;OUTS" + ${ARGN} + ) + + if(_RULE_TESTONLY AND NOT MGE_WITH_TEST) + return() + endif() + + set(_NAME ${_RULE_NAME}) + + set(LLVM_TARGET_DEFINITIONS ${_RULE_SRCS}) + set(_INCLUDE_DIRS ${_RULE_INCLUDES}) + list(TRANSFORM _INCLUDE_DIRS PREPEND "-I") + set(_OUTPUTS) + while(_RULE_OUTS) + list(GET _RULE_OUTS 0 _COMMAND) + list(REMOVE_AT _RULE_OUTS 0) + list(GET _RULE_OUTS 0 _FILE) + list(REMOVE_AT _RULE_OUTS 0) + tablegen(${_RULE_TBLGEN} ${_FILE} ${_COMMAND} ${_INCLUDE_DIRS}) + list(APPEND _OUTPUTS ${CMAKE_CURRENT_BINARY_DIR}/${_FILE}) + endwhile() + add_custom_target(${_NAME}_target DEPENDS ${_OUTPUTS}) + + add_library(${_NAME} INTERFACE) + add_dependencies(${_NAME} ${_NAME}_target) + + target_include_directories(${_NAME} INTERFACE + "$") + + install(TARGETS ${_NAME} EXPORT ${MGE_EXPORT_TARGETS}) +endfunction() + +if (MGE_USE_SYSTEM_LIB) + find_package(ZLIB) + find_package(MLIR REQUIRED CONFIG) + message(STATUS "Using MLIRConfig.cmake in: ${MLIR_DIR}") + message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}") + list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}") + list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}") + include(TableGen) + + set(MLIR_LLVM_INCLUDE_DIR ${LLVM_INCLUDE_DIRS} ${MLIR_INCLUDE_DIRS}) + + set(MLIR_LLVM_COMPONENTS Core;Support;X86CodeGen;OrcJIT;NVPTX) + llvm_map_components_to_libnames(MLIR_LLVM_LIBS ${MLIR_LLVM_COMPONENTS}) + set(MLIR_LLVM_LIB_DIR ${MLIR_INSTALL_PREFIX}/lib) + + function(find_mlir_llvm_lib lib) + find_library(${lib} + NAMES ${lib} + PATHS ${MLIR_LLVM_LIB_DIR} + NO_DEFAULT_PATH) + if(${${lib}} STREQUAL ${lib}-NOTFOUND) + message(FATAL_ERROR "${lib} not found, did you forget to build llvm-project?") + else() + list(APPEND MLIR_LLVM_LIBS ${lib}) + set(MLIR_LLVM_LIBS "${MLIR_LLVM_LIBS}" PARENT_SCOPE) + endif() + endfunction(find_mlir_llvm_lib) + + set(MLIR_COMPONENTS MLIRAnalysis;MLIRExecutionEngine;MLIRIR;MLIRParser;MLIRPass;MLIRSideEffectInterfaces;MLIRTargetLLVMIR;MLIRTransforms;MLIRAffineToStandard;MLIRSCFToStandard;MLIRAVX512ToLLVM;MLIRAVX512;MLIRLLVMAVX512;MLIRSDBM;MLIRROCDLIR;MLIRGPU;MLIRQuant;MLIRSPIRV;MLIRNVVMIR;MLIRShape;MLIRGPUToNVVMTransforms;MLIRTargetNVVMIR;MLIRGPUToGPURuntimeTransforms) + + foreach(c ${MLIR_COMPONENTS}) + find_mlir_llvm_lib(${c}) + endforeach() + return() +endif() + +function(add_mge_mlir_src_dep llvm_monorepo_path) + set(_CMAKE_BUILD_TYPE "${CMAKE_BUILD_TYPE}") + string(TOUPPER "${CMAKE_BUILD_TYPE}" uppercase_CMAKE_BUILD_TYPE) + if(NOT uppercase_CMAKE_BUILD_TYPE MATCHES "^(DEBUG|RELEASE|RELWITHDEBINFO|MINSIZEREL)$") + set(CMAKE_BUILD_TYPE "Debug") + endif() + set(_CMAKE_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS}) + set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE) + + add_subdirectory("${llvm_monorepo_path}/llvm" ${LLVM_BUILD_DIR} EXCLUDE_FROM_ALL) + + # Reset CMAKE_BUILD_TYPE to its previous setting + set(CMAKE_BUILD_TYPE "${_CMAKE_BUILD_TYPE}" CACHE STRING "Build type" FORCE) + # Reset BUILD_SHARED_LIBS to its previous setting + set(BUILD_SHARED_LIBS ${_CMAKE_BUILD_SHARED_LIBS} CACHE BOOL "Build shared libraries" FORCE) +endfunction() + +set(LLVM_INCLUDE_EXAMPLES OFF CACHE BOOL "" FORCE) +set(LLVM_INCLUDE_TESTS OFF CACHE BOOL "" FORCE) +set(LLVM_INCLUDE_BENCHMARKS OFF CACHE BOOL "" FORCE) +set(LLVM_ENABLE_RTTI ${MGE_ENABLE_RTTI} CACHE BOOL "" FORCE) +set(LLVM_TARGETS_TO_BUILD "X86;NVPTX;AMDGPU;AArch64;ARM;PowerPC;SystemZ" CACHE STRING "" FORCE) +set(LLVM_ENABLE_PROJECTS "mlir" CACHE STRING "" FORCE) +set(LLVM_BUILD_DIR ${PROJECT_BINARY_DIR}/third_party/llvm-project/llvm) + +add_mge_mlir_src_dep("third_party/llvm-project") + +set(MLIR_LLVM_INCLUDE_DIR + ${PROJECT_SOURCE_DIR}/third_party/llvm-project/llvm/include + ${PROJECT_BINARY_DIR}/third_party/llvm-project/llvm/include + ${PROJECT_SOURCE_DIR}/third_party/llvm-project/mlir/include + ${PROJECT_BINARY_DIR}/third_party/llvm-project/llvm/tools/mlir/include + ) +set(MLIR_TABLEGEN_EXE mlir-tblgen) + +set(MLIR_LLVM_LIBS LLVMCore;LLVMSupport;LLVMX86CodeGen;LLVMOrcJIT;LLVMNVPTXCodeGen;LLVMNVPTXDesc;LLVMNVPTXInfo;MLIRAnalysis;MLIRExecutionEngine;MLIRIR;MLIRParser;MLIRPass;MLIRSideEffectInterfaces;MLIRTargetLLVMIR;MLIRTransforms;MLIRAffineToStandard;MLIRSCFToStandard;MLIRAVX512ToLLVM;MLIRAVX512;MLIRLLVMAVX512;MLIRSDBM;MLIRROCDLIR;MLIRGPU;MLIRQuant;MLIRSPIRV;MLIRNVVMIR;MLIRGPUToNVVMTransforms;MLIRShape;MLIRTargetNVVMIR;MLIRGPUToGPURuntimeTransforms) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d3f9c19e7..733a8cee3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,7 +1,15 @@ +if(MGE_WITH_JIT_MLIR) + add_subdirectory(jit/impl/mlir/ir) +endif() + file(GLOB_RECURSE SOURCES core/impl/*.cpp gopt/impl/*.cpp opr/impl/*.cpp plugin/impl/*.cpp serialization/impl/*.cpp core/impl/*.inl gopt/impl/*.inl opr/impl/*.inl plugin/impl/*.inl serialization/impl/*.inl) if(MGE_WITH_JIT) file(GLOB_RECURSE SOURCES_ jit/impl/*.cpp jit/impl/*.inl) + if(MGE_WITH_JIT_MLIR) + file(GLOB_RECURSE MLIR_SOURCES_ jit/impl/mlir/ir/*.cpp jit/impl/mlir/*.cpp) + list(APPEND SOURCES_ ${MLIR_SOURCES_}) + endif() list(APPEND SOURCES ${SOURCES_}) endif() @@ -47,7 +55,7 @@ if(MGE_WITH_CUDA) list(APPEND SOURCES ${SOURCES_}) endif() -add_library(megbrain OBJECT EXCLUDE_FROM_ALL ${SOURCES}) +add_library(megbrain OBJECT ${SOURCES}) target_link_libraries(megbrain PUBLIC mgb_opr_param_defs) target_include_directories(megbrain PUBLIC $ @@ -91,6 +99,11 @@ if(MGE_WITH_JIT AND MGE_WITH_HALIDE) target_link_libraries(megbrain PRIVATE libhalide) target_link_libraries(megbrain PRIVATE ${HALIDE_LLVM_LIBS}) endif() +if(MGE_WITH_JIT_MLIR) + target_link_libraries(megbrain PRIVATE mlir_op_def) + target_link_libraries(megbrain PRIVATE mlir_shape_inference) + target_link_libraries(megbrain PRIVATE ${MLIR_LLVM_LIBS}) +endif() if (MGB_WITH_FLATBUFFERS) set (GEN_FLATBUFFERS_SCHEMA_PY ${PROJECT_SOURCE_DIR}/dnn/scripts/gen_flatbuffers_schema.py) set (OPR_PARAM_DEFS_PY ${PROJECT_SOURCE_DIR}/dnn/scripts/opr_param_defs.py) @@ -203,8 +216,7 @@ install(TARGETS megengine if (NOT MGE_WITH_DISTRIBUTED) install(TARGETS megbrain EXPORT ${MGE_EXPORT_TARGETS} - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}) + ) endif() foreach(_PATH ${MGB_INC}) diff --git a/src/jit/impl/mlir/compiler.cpp b/src/jit/impl/mlir/compiler.cpp index 4742c6851..6e3ae0235 100644 --- a/src/jit/impl/mlir/compiler.cpp +++ b/src/jit/impl/mlir/compiler.cpp @@ -17,7 +17,6 @@ #include "./executable_cpu.h" #include "./executable_cuda.h" #include "./mlir_gen.h" - #include "megbrain/common.h" #include "megbrain/comp_node_env.h" #include "megbrain/jit/mlir/ir/dialect.h" diff --git a/src/jit/impl/mlir/ir/CMakeLists.txt b/src/jit/impl/mlir/ir/CMakeLists.txt new file mode 100644 index 000000000..405a53410 --- /dev/null +++ b/src/jit/impl/mlir/ir/CMakeLists.txt @@ -0,0 +1,31 @@ +set(MGB_MLIR_TABLEGEN_INC_BASE ${CMAKE_CURRENT_BINARY_DIR}/include/) +file(MAKE_DIRECTORY ${MGB_MLIR_TABLEGEN_INC_BASE}/megbrain/jit/mlir/ir/) +list(APPEND MGB_MLIR_TABLEGEN_INC ${MGB_MLIR_TABLEGEN_INC_BASE}) + +external_tablegen_library( + NAME + mlir_shape_inference + TBLGEN + MLIR + SRCS + "shape_inference_interface.td" + INCLUDES + ${MGB_MLIR_TABLEGEN_INC} ${MLIR_LLVM_INCLUDE_DIR} + OUTS + -gen-op-interface-decls include/megbrain/jit/mlir/ir/shape_inference_interface.h.inc + -gen-op-interface-defs include/megbrain/jit/mlir/ir/shape_inference_interface.cpp.inc +) + +external_tablegen_library( + NAME + mlir_op_def + TBLGEN + MLIR + SRCS + "ops.td" + INCLUDES + ${MGB_MLIR_TABLEGEN_INC} ${MLIR_LLVM_INCLUDE_DIR} + OUTS + -gen-op-decls include/megbrain/jit/mlir/ir/ops.h.inc + -gen-op-defs include/megbrain/jit/mlir/ir/ops.cpp.inc +) diff --git a/src/jit/impl/mlir/ir/common.cpp b/src/jit/impl/mlir/ir/common.cpp index bddaeb6bb..cd22eda60 100644 --- a/src/jit/impl/mlir/ir/common.cpp +++ b/src/jit/impl/mlir/ir/common.cpp @@ -36,6 +36,6 @@ mlir::Value jit::insert_alloc_and_dealloc(mlir::MemRefType type, return alloc; } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}} diff --git a/src/jit/impl/mlir/ir/common.h b/src/jit/impl/mlir/ir/common.h index 251c569e5..ee01ff6c5 100644 --- a/src/jit/impl/mlir/ir/common.h +++ b/src/jit/impl/mlir/ir/common.h @@ -27,6 +27,6 @@ mlir::Value insert_alloc_and_dealloc(mlir::MemRefType type, mlir::Location loc, } // namespace jit } // namespace mgb -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}} diff --git a/src/jit/impl/mlir/ir/dialect.cpp b/src/jit/impl/mlir/ir/dialect.cpp index dcff5900d..35576109d 100644 --- a/src/jit/impl/mlir/ir/dialect.cpp +++ b/src/jit/impl/mlir/ir/dialect.cpp @@ -86,6 +86,6 @@ void AddOp::infer_shapes() { getResult().setType(getOperand(0).getType()); } #define GET_OP_CLASSES #include "megbrain/jit/mlir/ir/ops.cpp.inc" -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/ir/lower_to_affine_pass.cpp b/src/jit/impl/mlir/ir/lower_to_affine_pass.cpp index dbfbdeb37..41c20cdc6 100644 --- a/src/jit/impl/mlir/ir/lower_to_affine_pass.cpp +++ b/src/jit/impl/mlir/ir/lower_to_affine_pass.cpp @@ -154,6 +154,6 @@ std::unique_ptr mgb::jit::create_lower_to_affine_pass() { return std::make_unique(); } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/ir/lower_to_gpu_pass.cpp b/src/jit/impl/mlir/ir/lower_to_gpu_pass.cpp index ad72760ff..0e0c4ed0f 100644 --- a/src/jit/impl/mlir/ir/lower_to_gpu_pass.cpp +++ b/src/jit/impl/mlir/ir/lower_to_gpu_pass.cpp @@ -206,6 +206,6 @@ std::unique_ptr mgb::jit::create_lower_to_gpu_pass() { return std::make_unique(); } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/ir/lower_to_llvm_pass.cpp b/src/jit/impl/mlir/ir/lower_to_llvm_pass.cpp index ac0b22436..47785c040 100644 --- a/src/jit/impl/mlir/ir/lower_to_llvm_pass.cpp +++ b/src/jit/impl/mlir/ir/lower_to_llvm_pass.cpp @@ -51,6 +51,6 @@ std::unique_ptr mgb::jit::create_lower_to_llvm_pass() { return std::make_unique(); } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/ir/shape_inference_pass.cpp b/src/jit/impl/mlir/ir/shape_inference_pass.cpp index 12f3174cb..c00d5c837 100644 --- a/src/jit/impl/mlir/ir/shape_inference_pass.cpp +++ b/src/jit/impl/mlir/ir/shape_inference_pass.cpp @@ -95,6 +95,6 @@ std::unique_ptr mgb::jit::create_shape_inference_pass() { return std::make_unique(); } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/mlir_gen.cpp b/src/jit/impl/mlir/mlir_gen.cpp index aeb9f1426..346a53550 100644 --- a/src/jit/impl/mlir/mlir_gen.cpp +++ b/src/jit/impl/mlir/mlir_gen.cpp @@ -14,7 +14,6 @@ #if MGB_JIT && MGB_JIT_MLIR #include "./mlir_gen.h" - #include "./utils.h" #include "megbrain/jit/mlir/ir/dialect.h" #include "megbrain/opr/basic_arith.h" @@ -202,6 +201,6 @@ std::pair mgb::jit::mlir_gen( return MLIRGenImpl(context).gen(internal_graph, args); } -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/impl/mlir/mlir_gen.h b/src/jit/impl/mlir/mlir_gen.h index 25b79488f..21cd8a2a8 100644 --- a/src/jit/impl/mlir/mlir_gen.h +++ b/src/jit/impl/mlir/mlir_gen.h @@ -11,6 +11,7 @@ */ #pragma once +#include "megbrain_build_config.h" #include "megbrain_build_config.h" #if MGB_JIT && MGB_JIT_MLIR @@ -37,6 +38,6 @@ std::pair mlir_gen( } } // namespace mgb -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/include/megbrain/jit/mlir/ir/dialect.h b/src/jit/include/megbrain/jit/mlir/ir/dialect.h index 706587e2e..fd0b668e6 100644 --- a/src/jit/include/megbrain/jit/mlir/ir/dialect.h +++ b/src/jit/include/megbrain/jit/mlir/ir/dialect.h @@ -40,6 +40,6 @@ using namespace mlir; } // namespace jit } // namespace mgb -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/include/megbrain/jit/mlir/ir/passes.h b/src/jit/include/megbrain/jit/mlir/ir/passes.h index dc7eaaa74..9ca399080 100644 --- a/src/jit/include/megbrain/jit/mlir/ir/passes.h +++ b/src/jit/include/megbrain/jit/mlir/ir/passes.h @@ -11,6 +11,7 @@ */ #pragma once +#include "megbrain_build_config.h" #include #include "megbrain_build_config.h" @@ -38,6 +39,6 @@ std::unique_ptr create_lower_to_gpu_pass(); } // namespace jit } // namespace mgb -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/include/megbrain/jit/mlir/ir/shape_inference_interface.h b/src/jit/include/megbrain/jit/mlir/ir/shape_inference_interface.h index 594e72ba2..ea5ba9e40 100644 --- a/src/jit/include/megbrain/jit/mlir/ir/shape_inference_interface.h +++ b/src/jit/include/megbrain/jit/mlir/ir/shape_inference_interface.h @@ -28,6 +28,6 @@ namespace jit { -#endif // MGB_JIT_MLIR +#endif // MGB_JIT && MGB_JIT_MLIR // vim: syntax=cpp.doxygen diff --git a/src/jit/test/fusion.cpp b/src/jit/test/fusion.cpp index ec7f71d4b..94ad73561 100644 --- a/src/jit/test/fusion.cpp +++ b/src/jit/test/fusion.cpp @@ -1500,6 +1500,7 @@ TEST(TestJITExecutor, GradBehavior) { ASSERT_EQ(nr_ph, 2); ASSERT_EQ(nr_mul, 1); } +#if MGB_JIT_HALIDE { set_backend(Backend::HALIDE); auto graph = ComputingGraph::make(); @@ -1529,6 +1530,7 @@ TEST(TestJITExecutor, GradBehavior) { // GetVarShape(a) and broadcast ASSERT_EQ(nr_ops, 4); } +#endif // MGB_JIT_HALIDE { set_backend(Backend::NVRTC); auto graph = ComputingGraph::make(); diff --git a/src/megbrain_build_config.h.in b/src/megbrain_build_config.h.in index 5db6df2e1..30daf3e31 100644 --- a/src/megbrain_build_config.h.in +++ b/src/megbrain_build_config.h.in @@ -25,6 +25,7 @@ #cmakedefine01 MGB_BUILD_SLIM_SERVING #cmakedefine01 MGB_ENABLE_EXCEPTION #cmakedefine01 MGB_JIT +#cmakedefine01 MGB_JIT_MLIR #cmakedefine01 MGB_JIT_HALIDE #cmakedefine01 MGB_ENABLE_TENSOR_RT #cmakedefine01 MGB_ENABLE_JSON -- GitLab