diff --git a/drivers/infiniband/hw/hns/hns_roce_hw_v2.c b/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
index a017ca173deb75fc9a96efeba0977137e7c1dc12..0eb614d9388ea852c3afa14f51e239754a7c8c9b 100644
--- a/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
+++ b/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
@@ -3039,15 +3039,10 @@ static void *get_srq_wqe(struct hns_roce_srq *srq, int n)
 
 static void hns_roce_free_srq_wqe(struct hns_roce_srq *srq, int wqe_index)
 {
-	u32 bitmap_num;
-	int bit_num;
-
 	/* always called with interrupts disabled. */
 	spin_lock(&srq->lock);
 
-	bitmap_num = wqe_index / BITS_PER_LONG_LONG;
-	bit_num = wqe_index % BITS_PER_LONG_LONG;
-	srq->idx_que.bitmap[bitmap_num] |= (1ULL << bit_num);
+	bitmap_clear(srq->idx_que.bitmap, wqe_index, 1);
 	srq->idx_que.tail++;
 
 	spin_unlock(&srq->lock);
@@ -7043,18 +7038,19 @@ int hns_roce_srqwq_overflow(struct hns_roce_srq *srq, int nreq)
 	return cur + nreq >= srq->max - 1;
 }
 
-static int find_empty_entry(struct hns_roce_idx_que *idx_que)
+static int find_empty_entry(struct hns_roce_idx_que *idx_que,
+	unsigned long size)
 {
-	int bit_num;
-	int i;
+	int wqe_idx;
 
-	/* bitmap[i] is set zero if all bits are allocated */
-	for (i = 0; idx_que->bitmap[i] == 0; ++i)
-		;
-	bit_num = __ffs64(idx_que->bitmap[i]) + 1;
-	idx_que->bitmap[i] &= ~(1ULL << (bit_num - 1));
+	if (unlikely(bitmap_full(idx_que->bitmap, size)))
+		return -ENOSPC;
+
+	wqe_idx = find_first_zero_bit(idx_que->bitmap, size);
+
+	bitmap_set(idx_que->bitmap, wqe_idx, 1);
 
-	return i * BITS_PER_LONG_LONG + (bit_num - 1);
+	return wqe_idx;
 }
 
 static void fill_idx_queue(struct hns_roce_idx_que *idx_que,
@@ -7112,7 +7108,13 @@ static int hns_roce_v2_post_srq_recv(struct ib_srq *ibsrq,
 			break;
 		}
 
-		wqe_idx = find_empty_entry(&srq->idx_que);
+		wqe_idx = find_empty_entry(&srq->idx_que, srq->max);
+		if (wqe_idx < 0) {
+			ret = -ENOMEM;
+			*bad_wr = wr;
+			break;
+		}
+
 		fill_idx_queue(&srq->idx_que, ind, wqe_idx);
 		wqe = get_srq_wqe(srq, wqe_idx);
 		dseg = (struct hns_roce_v2_wqe_data_seg *)wqe;
diff --git a/drivers/infiniband/hw/hns/hns_roce_srq.c b/drivers/infiniband/hw/hns/hns_roce_srq.c
index 40ce6abdcd3744e5c9d893a96d02d05dc129b2b0..52a9b23c13e0bc18c9228a6af37c6081f2848bc8 100644
--- a/drivers/infiniband/hw/hns/hns_roce_srq.c
+++ b/drivers/infiniband/hw/hns/hns_roce_srq.c
@@ -312,29 +312,22 @@ static int hns_roce_create_idx_que(struct ib_pd *pd, struct hns_roce_srq *srq,
 	struct hns_roce_dev *hr_dev = to_hr_dev(pd->device);
 	struct hns_roce_idx_que *idx_que = &srq->idx_que;
 	struct hns_roce_buf *kbuf;
-	u32 bitmap_num;
-	int i;
 
 	idx_que->entry_sz = HNS_ROCE_IDX_QUE_ENTRY_SZ;
-	bitmap_num = HNS_ROCE_ALIGN_UP(srq->max, BITS_PER_LONG_LONG);
 
-	idx_que->bitmap = kcalloc(1, bitmap_num / BITS_PER_BYTE, GFP_KERNEL);
+	idx_que->bitmap = bitmap_zalloc(srq->max, GFP_KERNEL);
 	if (!idx_que->bitmap)
 		return -ENOMEM;
 
-	bitmap_num = bitmap_num / BITS_PER_LONG_LONG;
-
 	idx_que->buf_size = srq->max * idx_que->entry_sz;
 
 	kbuf = hns_roce_buf_alloc(hr_dev, idx_que->buf_size, page_shift, 0);
 	if (IS_ERR(kbuf)) {
-		kfree(idx_que->bitmap);
+		bitmap_free(idx_que->bitmap);
 		return -ENOMEM;
 	}
 
 	idx_que->idx_buf = kbuf;
-	for (i = 0; i < bitmap_num; i++)
-		idx_que->bitmap[i] = ~(0UL);
 
 	idx_que->head = 0;
 	idx_que->tail = 0;
@@ -405,7 +398,7 @@ static int create_kernel_srq(struct ib_pd *pd, struct hns_roce_srq *srq,
 
 err_kernel_create_idx:
 	hns_roce_buf_free(hr_dev, srq->idx_que.idx_buf);
-	kfree(srq->idx_que.bitmap);
+	bitmap_free(srq->idx_que.bitmap);
 
 err_kernel_srq_mtt:
 	hns_roce_mtt_cleanup(hr_dev, &srq->mtt);