From be684eeb4ce80cee51b200cadf4175745a6b3824 Mon Sep 17 00:00:00 2001 From: bdf <36697723+defei-coder@users.noreply.github.com> Date: Tue, 18 Oct 2022 14:24:27 +0800 Subject: [PATCH] [Feature] Add getJobLimitCapability interface and use it in nms (#2337) Co-authored-by: budefei --- mmcv/ops/csrc/common/pytorch_mlu_helper.hpp | 10 ++++++++++ mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp | 3 +-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/mmcv/ops/csrc/common/pytorch_mlu_helper.hpp b/mmcv/ops/csrc/common/pytorch_mlu_helper.hpp index 65bf3856a8..69f926f4ba 100644 --- a/mmcv/ops/csrc/common/pytorch_mlu_helper.hpp +++ b/mmcv/ops/csrc/common/pytorch_mlu_helper.hpp @@ -25,6 +25,16 @@ #define CEIL_ALIGN(x, y) (((x) + (y)-1) / (y) * (y)) +inline int32_t getJobLimitCapability() { + CNcontext drv_ctx; + CNctxConfigParam ctx_conf_param; + TORCH_CHECK( + CN_SUCCESS == cnGetCtxConfigParam(drv_ctx, CN_CTX_CONFIG_UNION_LIMIT, + &ctx_conf_param), + "cnGetCtxConfigParam fails."); + return (int32_t)ctx_conf_param.unionLimit; +} + #endif // MMCV_WITH_MLU #endif // PYTORCH_MLU_HELPER_HPP_ diff --git a/mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp b/mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp index 51a3003812..a45a510e89 100644 --- a/mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp +++ b/mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp @@ -34,8 +34,7 @@ static cnnlStatus_t policyFunc(cnrtDim3_t *k_dim, cnrtFunctionType_t *k_type, int &core_num_per_class, const int input_box_num) { uint32_t core_dim = torch_mlu::getDeviceAttr(cnrtAttrMcorePerCluster); - uint32_t cluster_number = torch_mlu::getDeviceAttr(cnrtAttrClusterCount); - uint32_t job_limit = cluster_number * core_dim; + uint32_t job_limit = getJobLimitCapability(); uint32_t core_number = job_limit; int box_num_per_core = (input_box_num + core_number - 1) / core_number;