提交 f9109bed 编写于 作者: M Megvii Engine Team

feat(mge): add mlir cmake

GitOrigin-RevId: e86db34b18d75c85f7ef120481f79359dc3e1da2
上级 a51d5b4c
......@@ -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)
......
# - 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
"$<BUILD_INTERFACE:${_RULE_INCLUDES}>")
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)
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 $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
......@@ -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})
......
......@@ -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"
......
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
)
......@@ -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}}}
......@@ -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}}}
......@@ -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
......@@ -154,6 +154,6 @@ std::unique_ptr<mlir::Pass> mgb::jit::create_lower_to_affine_pass() {
return std::make_unique<MgbToAffineLoweringPass>();
}
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -206,6 +206,6 @@ std::unique_ptr<mlir::Pass> mgb::jit::create_lower_to_gpu_pass() {
return std::make_unique<MgbToGpuLoweringPass>();
}
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -51,6 +51,6 @@ std::unique_ptr<mlir::Pass> mgb::jit::create_lower_to_llvm_pass() {
return std::make_unique<AffineToLLVMLoweringPass>();
}
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -95,6 +95,6 @@ std::unique_ptr<mlir::Pass> mgb::jit::create_shape_inference_pass() {
return std::make_unique<ShapeInferencePass>();
}
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -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<llvm::StringRef, mlir::OwningModuleRef> 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
......@@ -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<llvm::StringRef, mlir::OwningModuleRef> mlir_gen(
}
} // namespace mgb
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -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
......@@ -11,6 +11,7 @@
*/
#pragma once
#include "megbrain_build_config.h"
#include <mlir/IR/Module.h>
#include "megbrain_build_config.h"
......@@ -38,6 +39,6 @@ std::unique_ptr<mlir::Pass> create_lower_to_gpu_pass();
} // namespace jit
} // namespace mgb
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -28,6 +28,6 @@ namespace jit {
#endif // MGB_JIT_MLIR
#endif // MGB_JIT && MGB_JIT_MLIR
// vim: syntax=cpp.doxygen
......@@ -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();
......
......@@ -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
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册