diff --git a/CMakeLists.txt b/CMakeLists.txt index bbfb90b7..3d03d3be 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -79,6 +79,7 @@ option(MGE_WITH_MKLDNN "Enable Intel MKL_DNN support," ON) option(MGE_WITH_ROCM "Enable ROCM support" OFF) option(MGE_WITH_LARGE_ARCHIVE "Enable big archive link support" OFF) option(MGE_BUILD_WITH_ASAN "Enable build with ASAN, need compiler support" OFF) +option(MGE_WITH_CUSTOM_OP "Build with Custom op" OFF) if(MSVC OR WIN32) option(MGE_DEPLOY_INFERENCE_ON_WINDOWS_XP "Enable deploy inference on Windows xp" OFF) # special MGE_DEPLOY_INFERENCE_ON_WINDOWS_XP_SP2 for Windows XP sp2(32bit) @@ -874,6 +875,8 @@ if(NOT "${CUSTOM_C_OPR_INIT_FUNC}" STREQUAL "") message(STATUS "override MGB_C_OPR_INIT_FUNC to ${CUSTOM_C_OPR_INIT_FUNC}") endif() +set(MGB_CUSTOM_OP ${MGE_WITH_CUSTOM_OP}) + if(MSVC OR WIN32) set(CMAKE_HAVE_THREADS_LIBRARY 1) set(CMAKE_USE_WIN32_THREADS_INIT 1) diff --git a/imperative/python/src/ops.cpp b/imperative/python/src/ops.cpp index 60127c4f..50cc2186 100644 --- a/imperative/python/src/ops.cpp +++ b/imperative/python/src/ops.cpp @@ -632,6 +632,7 @@ void init_ops(py::module m) { } PyObject *make_custom_op(PyObject *self, PyObject **args, Py_ssize_t nargs) { +#if MGB_CUSTOM_OP auto op_name = py::handle(args[0]).cast(); auto kwargs = py::handle(args[1]).cast(); @@ -673,31 +674,52 @@ PyObject *make_custom_op(PyObject *self, PyObject **args, Py_ssize_t nargs) { reinterpret_cast(obj)->op = opdef; return obj; +#else + mgb_assert(false, "Custom Op is disabled now, please build megengine with Custom Op open"); + return nullptr; +#endif } #undef CUSTOM_CASE_TO_PARSE_LIST #undef CUSTOM_CASE_TO_PARSE_NON_LIST py::list install_custom(const std::string &name, const std::string &path) { +#if MGB_CUSTOM_OP py::list ret; const auto &ops_in_lib = custom::LibManager::inst()->install(name, path); for (const auto &op: ops_in_lib) { ret.append(op); } return ret; +#else + mgb_assert(false, "Custom Op is disabled now, please build megengine with Custom Op open"); + py::list ret; + return ret; +#endif } bool uninstall_custom(const std::string &name) { +#if MGB_CUSTOM_OP return custom::LibManager::inst()->uninstall(name); +#else + mgb_assert(false, "Custom Op is disabled now, please build megengine with Custom Op open"); + return false; +#endif } py::list get_custom_op_list(void) { +#if MGB_CUSTOM_OP std::vector all_ops = CustomOpDefFactory::inst()->op_list(); py::list ret; for (auto &op: all_ops) { ret.append(op); } return ret; +#else + mgb_assert(false, "Custom Op is disabled now, please build megengine with Custom Op open"); + py::list ret; + return ret; +#endif } #ifndef METH_FASTCALL diff --git a/imperative/src/impl/ops/custom_opdef.cpp b/imperative/src/impl/ops/custom_opdef.cpp index b59ff282..287b8536 100644 --- a/imperative/src/impl/ops/custom_opdef.cpp +++ b/imperative/src/impl/ops/custom_opdef.cpp @@ -10,6 +10,9 @@ */ #include "megbrain/imperative/ops/custom_opdef.h" + +#if MGB_CUSTOM_OP + #include "megbrain/opr/custom_opnode.h" #include "megbrain/custom/data_adaptor.h" #include "../op_trait.h" @@ -293,3 +296,5 @@ OP_TRAIT_REG(CustomOpDef, CustomOpDef) } // imperative } // mgb + +#endif diff --git a/imperative/src/include/megbrain/imperative/ops/custom_opdef.h b/imperative/src/include/megbrain/imperative/ops/custom_opdef.h index 82568a82..b8ed645d 100644 --- a/imperative/src/include/megbrain/imperative/ops/custom_opdef.h +++ b/imperative/src/include/megbrain/imperative/ops/custom_opdef.h @@ -11,6 +11,10 @@ #pragma once +#include "megbrain/common.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/custom.h" #include "megbrain/custom/manager.h" #include "megbrain/imperative/op_def.h" @@ -62,3 +66,5 @@ public: } // imperative } // mgb + +#endif diff --git a/scripts/whl/macos/macos_build_whl.sh b/scripts/whl/macos/macos_build_whl.sh index b7f7da8b..1993b481 100755 --- a/scripts/whl/macos/macos_build_whl.sh +++ b/scripts/whl/macos/macos_build_whl.sh @@ -130,6 +130,7 @@ function do_build() { echo "PYTHON_INCLUDE_DIR: ${PYTHON_INCLUDE_DIR}" #config build type to RelWithDebInfo to enable MGB_ENABLE_DEBUG_UTIL etc export EXTRA_CMAKE_ARGS="${ORG_EXTRA_CMAKE_FLAG} -DCMAKE_BUILD_TYPE=RelWithDebInfo" + export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DMGE_WITH_CUSTOM_OP=ON" #append cmake args for config python export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_EXECUTABLE=${PYTHON_DIR}/bin/python3" export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_LIBRARY=${PYTHON_LIBRARY}" diff --git a/scripts/whl/manylinux2014/do_build_common.sh b/scripts/whl/manylinux2014/do_build_common.sh index 31c85c9a..5f6b1087 100755 --- a/scripts/whl/manylinux2014/do_build_common.sh +++ b/scripts/whl/manylinux2014/do_build_common.sh @@ -114,6 +114,7 @@ do MINOR=${ver:1} PYTHON_DIR=/opt/python/cp${python_ver}-cp${ver}/ export EXTRA_CMAKE_ARGS="${ORG_EXTRA_CMAKE_FLAG} -DCMAKE_BUILD_TYPE=RelWithDebInfo" + export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DMGE_WITH_CUSTOM_OP=ON" export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_EXECUTABLE=${PYTHON_DIR}/bin/python3" export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_LIBRARY=${PYTHON_DIR}lib/" export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python${MAJOR}.${MINOR}" diff --git a/scripts/whl/windows/windows_build_whl.sh b/scripts/whl/windows/windows_build_whl.sh index 1faf81fa..e61ef3c8 100755 --- a/scripts/whl/windows/windows_build_whl.sh +++ b/scripts/whl/windows/windows_build_whl.sh @@ -151,6 +151,7 @@ function do_build() { echo "PYTHON_INCLUDE_DIR: ${PYTHON_INCLUDE_DIR}" #config build type to RelWithDebInfo to enable MGB_ENABLE_DEBUG_UTIL etc export EXTRA_CMAKE_ARGS="${ORG_EXTRA_CMAKE_FLAG} -DCMAKE_BUILD_TYPE=RelWithDebInfo " + export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DMGE_WITH_CUSTOM_OP=ON" #call build and install HOST_BUILD_ARGS=" -t -s" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d6615ace..be7caf74 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -2,7 +2,8 @@ if(MGE_WITH_JIT_MLIR) add_subdirectory(jit/include/megbrain/jit/mlir/ir) endif() -file(GLOB_RECURSE SOURCES core/impl/*.cpp gopt/impl/*.cpp opr/impl/*.cpp opr/impl/nvof/*.cpp plugin/impl/*.cpp serialization/impl/*.cpp custom/impl/*.cpp core/impl/*.inl gopt/impl/*.inl opr/impl/*.inl plugin/impl/*.inl serialization/impl/*.inl) +file(GLOB_RECURSE SOURCES core/impl/*.cpp gopt/impl/*.cpp opr/impl/*.cpp opr/impl/nvof/*.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) @@ -22,7 +23,7 @@ if(MGE_WITH_DISTRIBUTED) list(APPEND SOURCES ${GRPC_SRCS}) endif() -set(MGB_INC ${PROJECT_BINARY_DIR}/genfiles ${CMAKE_CURRENT_LIST_DIR}/core/include ${CMAKE_CURRENT_LIST_DIR}/gopt/include ${CMAKE_CURRENT_LIST_DIR}/opr/include ${CMAKE_CURRENT_LIST_DIR}/plugin/include ${CMAKE_CURRENT_LIST_DIR}/serialization/include ${CMAKE_CURRENT_LIST_DIR}/custom/include) +set(MGB_INC ${PROJECT_BINARY_DIR}/genfiles ${CMAKE_CURRENT_LIST_DIR}/core/include ${CMAKE_CURRENT_LIST_DIR}/gopt/include ${CMAKE_CURRENT_LIST_DIR}/opr/include ${CMAKE_CURRENT_LIST_DIR}/plugin/include ${CMAKE_CURRENT_LIST_DIR}/serialization/include) if(MGE_WITH_JIT) list(APPEND MGB_INC ${CMAKE_CURRENT_LIST_DIR}/jit/include) @@ -55,6 +56,12 @@ if(MGE_WITH_CUDA) list(APPEND SOURCES ${SOURCES_}) endif() +if(MGE_WITH_CUSTOM_OP) + list(APPEND MGB_INC ${CMAKE_CURRENT_LIST_DIR}/custom/include) + file(GLOB_RECURSE SOURCES_ custom/impl/*.cpp) + list(APPEND SOURCES ${SOURCES_}) +endif() + add_library(megbrain OBJECT ${SOURCES}) target_link_libraries(megbrain PUBLIC mgb_opr_param_defs) if(MGE_WITH_CUDA) diff --git a/src/custom/impl/manager.cpp b/src/custom/impl/manager.cpp index ddef40c8..f87abacd 100644 --- a/src/custom/impl/manager.cpp +++ b/src/custom/impl/manager.cpp @@ -9,8 +9,11 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ -#include "megbrain/custom/manager.h" #include "megbrain/common.h" + +#if MGB_CUSTOM_OP + +#include "megbrain/custom/manager.h" #include #ifndef _WIN32 @@ -179,3 +182,5 @@ std::shared_ptr op_insert(std::string opname, uint32_t version) { } } + +#endif diff --git a/src/custom/impl/op.cpp b/src/custom/impl/op.cpp index 156ee988..7204111a 100644 --- a/src/custom/impl/op.cpp +++ b/src/custom/impl/op.cpp @@ -10,6 +10,9 @@ */ #include "megbrain/common.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/op.h" #include "megbrain/custom/utils.h" #include @@ -529,3 +532,5 @@ void CustomOp::compute(const std::vector &inputs, const Param ¶m, st } } + +#endif diff --git a/src/custom/impl/param.cpp b/src/custom/impl/param.cpp index 5d790b1d..8f434cb4 100644 --- a/src/custom/impl/param.cpp +++ b/src/custom/impl/param.cpp @@ -9,8 +9,11 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ -#include "megbrain/custom/param.h" #include "megbrain/common.h" + +#if MGB_CUSTOM_OP + +#include "megbrain/custom/param.h" #include "megbrain/utils/hash.h" #include #include @@ -177,3 +180,5 @@ bool operator==(const Param &lhs, const Param &rhs) { } } + +#endif diff --git a/src/custom/impl/param_val.cpp b/src/custom/impl/param_val.cpp index 0ff1662f..3004ba97 100644 --- a/src/custom/impl/param_val.cpp +++ b/src/custom/impl/param_val.cpp @@ -9,9 +9,12 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ -#include "megbrain/custom/param_val.h" #include "megbrain/common.h" +#if MGB_CUSTOM_OP + +#include "megbrain/custom/param_val.h" + #pragma GCC diagnostic ignored "-Wsign-compare" using namespace mgb; @@ -398,3 +401,5 @@ CUSTOM_DEFINE_BINARY_OP_FOR_BASIC_AND_STRING_AND_LIST(>, bool) CUSTOM_DEFINE_BINARY_OP_FOR_BASIC_AND_STRING_AND_LIST(<, bool) } + +#endif diff --git a/src/custom/impl/tensor.cpp b/src/custom/impl/tensor.cpp index 88491e11..31b68925 100644 --- a/src/custom/impl/tensor.cpp +++ b/src/custom/impl/tensor.cpp @@ -9,9 +9,12 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "megbrain/common.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/tensor.h" #include "megbrain/comp_node.h" -#include "megbrain/common.h" #include "megbrain/tensor.h" #include #include @@ -484,3 +487,5 @@ const void *Tensor::data(void) const { } } // namespace custom + +#endif diff --git a/src/custom/impl/utils.cpp b/src/custom/impl/utils.cpp index 1428c7c9..73810916 100644 --- a/src/custom/impl/utils.cpp +++ b/src/custom/impl/utils.cpp @@ -9,8 +9,11 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ -#include "megbrain/custom/utils.h" #include "megbrain/common.h" + +#if MGB_CUSTOM_OP + +#include "megbrain/custom/utils.h" #include using namespace mgb; @@ -39,3 +42,5 @@ UnImpleWarnLog::UnImpleWarnLog(const std::string &func, const std::string &attr, } } + +#endif diff --git a/src/custom/test/manager.cpp b/src/custom/test/manager.cpp index d4389237..aee8e1ff 100644 --- a/src/custom/test/manager.cpp +++ b/src/custom/test/manager.cpp @@ -9,6 +9,10 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "megbrain_build_config.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/manager.h" #include "megbrain/custom/custom.h" #include "gtest/gtest.h" @@ -94,3 +98,5 @@ TEST(TestOpManager, TestOpReg) { } } + +#endif diff --git a/src/custom/test/op.cpp b/src/custom/test/op.cpp index 20449d52..b42c6c1c 100644 --- a/src/custom/test/op.cpp +++ b/src/custom/test/op.cpp @@ -9,6 +9,10 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "megbrain_build_config.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/op.h" #include "megbrain/comp_node.h" #include "megbrain/tensor.h" @@ -203,3 +207,5 @@ TEST(TestCustomOp, TestCustomOpFuncSetter) { } } + +#endif diff --git a/src/custom/test/param.cpp b/src/custom/test/param.cpp index 7f22ead8..4beacb06 100644 --- a/src/custom/test/param.cpp +++ b/src/custom/test/param.cpp @@ -9,6 +9,10 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "megbrain_build_config.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/param.h" #include "gtest/gtest.h" #include @@ -206,3 +210,5 @@ TEST(TestParam, TestParam) { } } + +#endif diff --git a/src/custom/test/tensor.cpp b/src/custom/test/tensor.cpp index d9d4366c..007b0b14 100644 --- a/src/custom/test/tensor.cpp +++ b/src/custom/test/tensor.cpp @@ -9,6 +9,10 @@ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "megbrain_build_config.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/tensor.h" #include "megbrain/custom/data_adaptor.h" #include "megbrain/comp_node.h" @@ -323,3 +327,5 @@ TEST(TestTensor, TestTensorAccessor1D) { } } + +#endif diff --git a/src/megbrain_build_config.h.in b/src/megbrain_build_config.h.in index 49de6b1d..0753327d 100644 --- a/src/megbrain_build_config.h.in +++ b/src/megbrain_build_config.h.in @@ -34,6 +34,7 @@ #cmakedefine01 MGB_ENABLE_OPR_MM #cmakedefine01 MGB_ENABLE_FBS_SERIALIZATION #cmakedefine01 MGB_IS_DEV +#cmakedefine01 MGB_CUSTOM_OP // DNN related flags // Platform macro's #cmakedefine01 MEGDNN_WITH_CUDA @@ -161,6 +162,10 @@ #define MGB_JIT_HALIDE 0 #endif +#ifndef MGB_CUSTOM_OP +#define MGB_CUSTOM_OP 0 +#endif + #ifndef MEGDNN_WITH_CAMBRICON #define MEGDNN_WITH_CAMBRICON 0 #endif diff --git a/src/opr/impl/custom_opnode.cpp b/src/opr/impl/custom_opnode.cpp index 68a627b1..1cce0572 100644 --- a/src/opr/impl/custom_opnode.cpp +++ b/src/opr/impl/custom_opnode.cpp @@ -11,6 +11,8 @@ #include "megbrain/opr/custom_opnode.h" +#if MGB_CUSTOM_OP + namespace mgb { namespace opr { @@ -324,3 +326,5 @@ custom::ArgInfo CustomOpNode::output_info(size_t idx) const { } } + +#endif diff --git a/src/opr/include/megbrain/opr/custom_opnode.h b/src/opr/include/megbrain/opr/custom_opnode.h index 8385626d..a0a9b2a0 100644 --- a/src/opr/include/megbrain/opr/custom_opnode.h +++ b/src/opr/include/megbrain/opr/custom_opnode.h @@ -11,6 +11,10 @@ #pragma once +#include "megbrain/common.h" + +#if MGB_CUSTOM_OP + #include "megbrain/custom/custom.h" #include "megbrain/custom/manager.h" #include "megbrain/custom/data_adaptor.h" @@ -101,3 +105,5 @@ public: } // namespace opr } + +#endif diff --git a/src/serialization/impl/sereg_caller.cpp b/src/serialization/impl/sereg_caller.cpp index 08a5c96c..f304baed 100644 --- a/src/serialization/impl/sereg_caller.cpp +++ b/src/serialization/impl/sereg_caller.cpp @@ -29,7 +29,7 @@ namespace mgb{void call_sereg(){}} #include "../../opr/impl/tensor_gen.sereg.h" #include "../../opr/impl/tensor_manip.sereg.h" #include "../../opr/impl/utility.sereg.h" -#include "../../opr/impl/custom_opnode.sereg.h" + #if MGB_ENABLE_TENSOR_RT #include "../../tensorrt/impl/tensorrt_opr.sereg.h" #endif @@ -42,3 +42,7 @@ namespace mgb{void call_sereg(){}} #if MGB_CAMBRICON #include "../../cambricon/impl/cambricon_runtime_opr.sereg.h" #endif + +#if MGB_CUSTOM_OP +#include "../../opr/impl/custom_opnode.sereg.h" +#endif diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index fbea06a3..2b13c9b4 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -1,7 +1,7 @@ include_directories("./src/include") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter") -file(GLOB_RECURSE SOURCES ./*.cpp ../src/core/test/*.cpp ../src/gopt/test/*.cpp ../src/opr/test/*.cpp ../src/plugin/test/*.cpp ../src/serialization/test/*.cpp ../src/custom/test/*.cpp) +file(GLOB_RECURSE SOURCES ./*.cpp ../src/core/test/*.cpp ../src/gopt/test/*.cpp ../src/opr/test/*.cpp ../src/plugin/test/*.cpp ../src/serialization/test/*.cpp) if(MGE_WITH_JIT) file(GLOB_RECURSE SOURCES_ ../src/jit/test/*.cpp) list(APPEND SOURCES ${SOURCES_})