diff --git a/python_module/megengine/_internal/__init__.py b/python_module/megengine/_internal/__init__.py index e46ec7cd4f6f67ae7bbbad04b75c0f509e99f951..684aa51e94e40c7fa318e466315dd51da92f3764 100644 --- a/python_module/megengine/_internal/__init__.py +++ b/python_module/megengine/_internal/__init__.py @@ -25,7 +25,7 @@ from . import config, craniotome, dtype from . import global_init as _global_init from . import helper as _helper from . import mgb as _detail -from . import opr, opr_param_defs, plugin +from . import opr, opr_extra, opr_param_defs, plugin from .exc import MegBrainError from .logconf import get_logger from .mgb import ( diff --git a/python_module/megengine/_internal/opr_extra.py b/python_module/megengine/_internal/opr_extra.py new file mode 100644 index 0000000000000000000000000000000000000000..7a59a2cfc0e933c1c744874791387748bbda5467 --- /dev/null +++ b/python_module/megengine/_internal/opr_extra.py @@ -0,0 +1,3 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2015-2019 Megvii Inc. All rights reserved. + diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 20ddcbbcb338119e58236882920493ee308dfe93..0a9263ac8f8a2bc1c689778b9a24b2771fa5f4b3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -34,6 +34,11 @@ if(MGE_WITH_CUDA AND MGE_WITH_TRT) endif() +if(MGE_WITH_CUDA) + file(GLOB_RECURSE SOURCES_ opr/impl/standalone/*.cu) + list(APPEND SOURCES ${SOURCES_}) +endif() + add_library(megbrain OBJECT EXCLUDE_FROM_ALL ${SOURCES}) target_link_libraries(megbrain PUBLIC mgb_opr_param_defs) target_include_directories(megbrain diff --git a/src/opr/impl/mgb_cpp_opr.fbs b/src/opr/impl/mgb_cpp_opr.fbs index 75b76795c285e4d6b286e7824e836d7422e0d6d2..48dfe44c99d3c21e35b94cdb42d9fdd90f67900e 100644 --- a/src/opr/impl/mgb_cpp_opr.fbs +++ b/src/opr/impl/mgb_cpp_opr.fbs @@ -55,3 +55,8 @@ struct IndexDescMaskItem { table IndexDescMaskDump { items:[IndexDescMaskItem]; } + +table NMSKeep { + iou_thresh:float; + max_output:uint; +} diff --git a/src/serialization/impl/schema.fbs b/src/serialization/impl/schema.fbs index d0a85bfdbf259810b335a1b54ab177d1a9b00ec7..623f5a704c0ff006c55f70b5949cce959af8d972 100644 --- a/src/serialization/impl/schema.fbs +++ b/src/serialization/impl/schema.fbs @@ -98,6 +98,7 @@ union OperatorParam { param.IndexDescMaskDump = 66, DType = 67, param.Remap = 68, + param.NMSKeep = 69, } table Operator {