qp->sumem = umem;
        qplib_qp->sq.sg_info.sghead = umem->sg_head.sgl;
-       qplib_qp->sq.sg_info.npages = ib_umem_num_pages(umem);
+       qplib_qp->sq.sg_info.npages = ib_umem_num_dma_blocks(umem, PAGE_SIZE);
        qplib_qp->sq.sg_info.nmap = umem->nmap;
        qplib_qp->sq.sg_info.pgsize = PAGE_SIZE;
        qplib_qp->sq.sg_info.pgshft = PAGE_SHIFT;
                        goto rqfail;
                qp->rumem = umem;
                qplib_qp->rq.sg_info.sghead = umem->sg_head.sgl;
-               qplib_qp->rq.sg_info.npages = ib_umem_num_pages(umem);
+               qplib_qp->rq.sg_info.npages =
+                       ib_umem_num_dma_blocks(umem, PAGE_SIZE);
                qplib_qp->rq.sg_info.nmap = umem->nmap;
                qplib_qp->rq.sg_info.pgsize = PAGE_SIZE;
                qplib_qp->rq.sg_info.pgshft = PAGE_SHIFT;
 
        srq->umem = umem;
        qplib_srq->sg_info.sghead = umem->sg_head.sgl;
-       qplib_srq->sg_info.npages = ib_umem_num_pages(umem);
+       qplib_srq->sg_info.npages = ib_umem_num_dma_blocks(umem, PAGE_SIZE);
        qplib_srq->sg_info.nmap = umem->nmap;
        qplib_srq->sg_info.pgsize = PAGE_SIZE;
        qplib_srq->sg_info.pgshft = PAGE_SHIFT;
                        goto fail;
                }
                cq->qplib_cq.sg_info.sghead = cq->umem->sg_head.sgl;
-               cq->qplib_cq.sg_info.npages = ib_umem_num_pages(cq->umem);
+               cq->qplib_cq.sg_info.npages =
+                       ib_umem_num_dma_blocks(cq->umem, PAGE_SIZE);
                cq->qplib_cq.sg_info.nmap = cq->umem->nmap;
                cq->qplib_cq.dpi = &uctx->dpi;
        } else {
        return rc;
 }
 
-static int bnxt_re_page_size_ok(int page_shift)
-{
-       switch (page_shift) {
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_4K:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_8K:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_64K:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_2M:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_256K:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_1M:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_4M:
-       case CMDQ_REGISTER_MR_LOG2_PBL_PG_SIZE_PG_1G:
-               return 1;
-       default:
-               return 0;
-       }
-}
-
 static int fill_umem_pbl_tbl(struct ib_umem *umem, u64 *pbl_tbl_orig,
                             int page_shift)
 {
        struct bnxt_re_mr *mr;
        struct ib_umem *umem;
        u64 *pbl_tbl = NULL;
-       int umem_pgs, page_shift, rc;
+       unsigned long page_size;
+       int umem_pgs, rc;
 
        if (length > BNXT_RE_MAX_MR_SIZE) {
                ibdev_err(&rdev->ibdev, "MR Size: %lld > Max supported:%lld\n",
        mr->ib_umem = umem;
 
        mr->qplib_mr.va = virt_addr;
-       umem_pgs = ib_umem_page_count(umem);
-       if (!umem_pgs) {
-               ibdev_err(&rdev->ibdev, "umem is invalid!");
-               rc = -EINVAL;
-               goto free_umem;
-       }
-       mr->qplib_mr.total_size = length;
-
-       pbl_tbl = kcalloc(umem_pgs, sizeof(u64 *), GFP_KERNEL);
-       if (!pbl_tbl) {
-               rc = -ENOMEM;
-               goto free_umem;
-       }
-
-       page_shift = __ffs(ib_umem_find_best_pgsz(umem,
-                               BNXT_RE_PAGE_SIZE_4K | BNXT_RE_PAGE_SIZE_2M,
-                               virt_addr));
-
-       if (!bnxt_re_page_size_ok(page_shift)) {
+       page_size = ib_umem_find_best_pgsz(
+               umem, BNXT_RE_PAGE_SIZE_4K | BNXT_RE_PAGE_SIZE_2M, virt_addr);
+       if (!page_size) {
                ibdev_err(&rdev->ibdev, "umem page size unsupported!");
                rc = -EFAULT;
-               goto fail;
+               goto free_umem;
        }
+       mr->qplib_mr.total_size = length;
 
-       if (page_shift == BNXT_RE_PAGE_SHIFT_4K &&
+       if (page_size == BNXT_RE_PAGE_SIZE_4K &&
            length > BNXT_RE_MAX_MR_SIZE_LOW) {
                ibdev_err(&rdev->ibdev, "Requested MR Sz:%llu Max sup:%llu",
                          length, (u64)BNXT_RE_MAX_MR_SIZE_LOW);
                rc = -EINVAL;
-               goto fail;
+               goto free_umem;
+       }
+
+       umem_pgs = ib_umem_num_dma_blocks(umem, page_size);
+       pbl_tbl = kcalloc(umem_pgs, sizeof(u64 *), GFP_KERNEL);
+       if (!pbl_tbl) {
+               rc = -ENOMEM;
+               goto free_umem;
        }
 
        /* Map umem buf ptrs to the PBL */
-       umem_pgs = fill_umem_pbl_tbl(umem, pbl_tbl, page_shift);
+       umem_pgs = fill_umem_pbl_tbl(umem, pbl_tbl, order_base_2(page_size));
        rc = bnxt_qplib_reg_mr(&rdev->qplib_res, &mr->qplib_mr, pbl_tbl,
-                              umem_pgs, false, 1 << page_shift);
+                              umem_pgs, false, page_size);
        if (rc) {
                ibdev_err(&rdev->ibdev, "Failed to register user MR");
                goto fail;