[Feature] Add getJobLimitCapability interface and use it in nms (#2337)

Co-authored-by: budefei <budefei@cambricon.com>
open-mmlab · Oct 18, 2022 · be684ee · be684ee
1 parent bd1da5a
commit be684ee
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git 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
@@ -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;