feat(src): add MGB_CUSTOM_OP macro

GitOrigin-RevId: beddcf974576af51da24b34b5d8bba4753a2e4f6

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)

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<std::string>();
     auto kwargs = py::handle(args[1]).cast<py::dict>();
 
@@ -673,31 +674,52 @@ PyObject *make_custom_op(PyObject *self, PyObject **args, Py_ssize_t nargs) {
     reinterpret_cast<PyOp(OpDef)*>(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<std::string> 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

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

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

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}"

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}"

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"

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)

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 <unordered_set>
 
 #ifndef _WIN32
@@ -179,3 +182,5 @@ std::shared_ptr<CustomOp> op_insert(std::string opname, uint32_t version) {
 }
 
 }
+
+#endif

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 <unordered_set>
@@ -529,3 +532,5 @@ void CustomOp::compute(const std::vector<Tensor> &inputs, const Param &param, st
 }
 
 }
+
+#endif

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 <limits>
 #include <sstream>
@@ -177,3 +180,5 @@ bool operator==(const Param &lhs, const Param &rhs) {
 }
 
 }
+
+#endif

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

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 <cctype>
 #include <algorithm>
@@ -484,3 +487,5 @@ const void *Tensor::data(void) const {
 }
 
 } // namespace custom
+
+#endif

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 <sstream>
 
 using namespace mgb;
@@ -39,3 +42,5 @@ UnImpleWarnLog::UnImpleWarnLog(const std::string &func, const std::string &attr,
 }
 
 }
+
+#endif

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

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

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 <iostream>
@@ -206,3 +210,5 @@ TEST(TestParam, TestParam) {
 }
 
 }
+
+#endif

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

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

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

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

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

test/CMakeLists.txt

 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_})