roi/psroi fixed

Author: yury-intel

src/plugins/intel_cpu/src/nodes/psroi_pooling.cpp

@@ -85,6 +85,7 @@ PSROIPooling::PSROIPooling(const std::shared_ptr<ov::Node>& op, const GraphConte
     const auto defPsroi = ov::as_type_ptr<const ov::op::v1::DeformablePSROIPooling>(op);
 
     noTrans = op->get_input_size() == 2;
+    inBatchNum = op->get_input_shape(0)[0];
     CPU_NODE_ASSERT(op->get_input_shape(0).size() == 4,
                     "has first input with incorrect rank: " + std::to_string(op->get_input_shape(0).size()));
     CPU_NODE_ASSERT(op->get_input_shape(1).size() == 2,
@@ -619,6 +620,7 @@ void PSROIPooling::executeSpecified() {
     cpu_parallel->parallel_for(realRois, [&](int currentRoi) {
         const float* bottomRois = bottomRoisBeginning + currentRoi * 5;
         auto roiBatchInd = static_cast<int>(bottomRois[0]);
+        OPENVINO_ASSERT(roiBatchInd <= inBatchNum, "required batch index > batch amount");
         if (getAlgorithm() == Algorithm::PSROIPoolingAverage) {
             executeAverage(srcData, dstData, bottomRois, currentRoi, roiBatchInd, *srcDesc, *dstDesc);
         } else if (getAlgorithm() == Algorithm::PSROIPoolingBilinear) {

src/plugins/intel_cpu/src/nodes/psroi_pooling.h

@@ -48,6 +48,8 @@ class PSROIPooling : public Node {
     int nh = 0;
     int nw = 0;
 
+    int inBatchNum = 0;
+
     // for Deformable PSROIPolling
     bool noTrans;
     int partSize = 1;

src/plugins/intel_cpu/src/nodes/roi_pooling.cpp

@@ -377,6 +377,7 @@ size_t RoiPoolingKey::hash() const {
     seed = hash_combine(seed, refParams.mb);
     seed = hash_combine(seed, refParams.c);
     seed = hash_combine(seed, refParams.nb_c);
+    seed = hash_combine(seed, refParams.b_num);
     seed = hash_combine(seed, refParams.c_block);
     seed = hash_combine(seed, refParams.nb_c_blocking);
     seed = hash_combine(seed, refParams.ih);
@@ -546,9 +547,12 @@ void ROIPooling::prepareParams() {
     const auto& inDims = getParentEdgeAt(0)->getMemory().getStaticDims();
     const auto& outDims = getChildEdgeAt(0)->getMemory().getStaticDims();
 
+    const auto& featureShape = getParentEdgeAt(0)->getMemory().getStaticDims();
+
     refParams.mb = outDims[0];
     refParams.c = rnd_up(inDims[1], refParams.c_block);
     refParams.nb_c = refParams.c / refParams.c_block;
+    refParams.b_num = featureShape[0];
     refParams.ih = inDims[2];
     refParams.iw = inDims[3];
     refParams.oh = outDims[2];
@@ -619,6 +623,9 @@ class ROIPooling::ROIPoolingJitExecutor : public ROIPooling::ROIPoolingExecutor
             if (roi_batch_ind == -1) {
                 break;
             }
+            OPENVINO_ASSERT(0 <= roi_batch_ind && roi_batch_ind <= jpp.b_num,
+                            "takes incorrect roi_ind, max roi_ind = ",
+                            jpp.b_num);
         }
 
         cpuParallel->parallel_for4d(MB, cb_work, jpp.oh, jpp.ow, [&](int n, int cbb, int oh, int ow) {
@@ -636,6 +643,9 @@ class ROIPooling::ROIPoolingJitExecutor : public ROIPooling::ROIPoolingExecutor
                 const auto* src_roi_ptr = &src_roi[roi_off];
 
                 auto roi_batch_ind = static_cast<int>(src_roi_ptr[0]);
+                OPENVINO_ASSERT(0 <= roi_batch_ind && roi_batch_ind <= jpp.b_num,
+                                "takes incorrect roi_ind, max roi_ind = ",
+                                jpp.b_num);
 
                 if (jpp.alg == Algorithm::ROIPoolingMax) {
                     auto roi_start_w = static_cast<int>(round(src_roi_ptr[1] * jpp.spatial_scale));
@@ -758,6 +768,9 @@ class ROIPooling::ROIPoolingRefExecutor : public ROIPooling::ROIPoolingExecutor
             if (roi_batch_ind == -1) {
                 break;
             }
+            OPENVINO_ASSERT(0 <= roi_batch_ind && roi_batch_ind <= jpp.b_num,
+                            "takes incorrect roi_ind, max roi_ind = ",
+                            jpp.b_num);
         }
 
         cpuParallel->parallel_for4d(MB, cb_work, jpp.oh, jpp.ow, [&](int n, int cbb, int oh, int ow) {
@@ -780,6 +793,9 @@ class ROIPooling::ROIPoolingRefExecutor : public ROIPooling::ROIPoolingExecutor
                 const auto* src_roi_ptr = &src_roi[roi_off];
 
                 auto roi_batch_ind = static_cast<int>(src_roi_ptr[0]);
+                OPENVINO_ASSERT(0 <= roi_batch_ind && roi_batch_ind <= jpp.b_num,
+                                "takes incorrect roi_ind, max roi_ind = ",
+                                jpp.b_num);
 
                 if (jpp.alg == Algorithm::ROIPoolingMax) {
                     auto roi_start_w = static_cast<int>(round(src_roi_ptr[1] * jpp.spatial_scale));

src/plugins/intel_cpu/src/nodes/roi_pooling.h

@@ -27,6 +27,8 @@ struct jit_roi_pooling_params {
 
     int c_block, nb_c, nb_c_blocking;
 
+    int b_num;
+
     double spatial_scale;
     int pooled_h;
     int pooled_w;