RDMA/mlx5: Move mlx5_ib_cont_pages() to the creation of the mlx5_ib_mr

For the user MR path, instead of calling this after getting the umem, call
it as part of creating the struct mlx5_ib_mr and distill its output to a
single page_shift stored inside the mr.

This avoids passing around the tuple of its output. Based on the umem and
page_shift, the output arguments can be computed using:

  count == ib_umem_num_pages(mr->umem)
  shift == mr->page_shift
  ncont == ib_umem_num_blocks(mr->umem, mr->page_shift)
  order == order_base_2(ncont)

And since mr->page_shift == umem_odp->page_shift then ncont ==
ib_umem_odp_num_pages() for ODP umems.

Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>

drivers/infiniband/hw/mlx5/mem.c

@@ -57,6 +57,17 @@ void mlx5_ib_cont_pages(struct ib_umem *umem, u64 addr,
 	struct scatterlist *sg;
 	int entry;
 
+	if (umem->is_odp) {
+		struct ib_umem_odp *odp = to_ib_umem_odp(umem);
+
+		*shift = odp->page_shift;
+		*ncont = ib_umem_odp_num_pages(odp);
+		*count = *ncont << (*shift - PAGE_SHIFT);
+		if (order)
+			*order = ilog2(roundup_pow_of_two(*count));
+		return;
+	}
+
 	addr = addr >> PAGE_SHIFT;
 	tmp = (unsigned long)addr;
 	m = find_first_bit(&tmp, BITS_PER_LONG);

drivers/infiniband/hw/mlx5/mlx5_ib.h

@@ -580,6 +580,7 @@ struct mlx5_ib_mr {
 	int			max_descs;
 	int			desc_size;
 	int			access_mode;
+	unsigned int		page_shift;
 	struct mlx5_core_mkey	mmkey;
 	struct ib_umem	       *umem;
 	struct mlx5_shared_mr_info	*smr_info;
@@ -614,6 +615,12 @@ struct mlx5_ib_mr {
 	struct mlx5_async_work  cb_work;
 };
 
+/* Returns the number of HCA pages inside the MR, ie the size of the MTT */
+static inline size_t mlx5_mr_num_pages(struct mlx5_ib_mr *mr)
+{
+	return ib_umem_num_blocks(mr->umem, mr->page_shift);
+}
+
 static inline bool is_odp_mr(struct mlx5_ib_mr *mr)
 {
 	return IS_ENABLED(CONFIG_INFINIBAND_ON_DEMAND_PAGING) && mr->umem &&

drivers/infiniband/hw/mlx5/mr.c

@@ -862,14 +862,11 @@ static int mr_cache_max_order(struct mlx5_ib_dev *dev)
 	return MLX5_MAX_UMR_SHIFT;
 }
 
-static int mr_umem_get(struct mlx5_ib_dev *dev, u64 start, u64 length,
-		       int access_flags, struct ib_umem **umem, int *npages,
-		       int *page_shift, int *ncont, int *order)
+static struct ib_umem *mr_umem_get(struct mlx5_ib_dev *dev, u64 start,
+				   u64 length, int access_flags)
 {
 	struct ib_umem *u;
 
-	*umem = NULL;
-
 	if (access_flags & IB_ACCESS_ON_DEMAND) {
 		struct ib_umem_odp *odp;
 
@@ -878,39 +875,17 @@ static int mr_umem_get(struct mlx5_ib_dev *dev, u64 start, u64 length,
 		if (IS_ERR(odp)) {
 			mlx5_ib_dbg(dev, "umem get failed (%ld)\n",
 				    PTR_ERR(odp));
-			return PTR_ERR(odp);
-		}
-
-		u = &odp->umem;
-
-		*page_shift = odp->page_shift;
-		*ncont = ib_umem_odp_num_pages(odp);
-		*npages = *ncont << (*page_shift - PAGE_SHIFT);
-		if (order)
-			*order = ilog2(roundup_pow_of_two(*ncont));
-	} else {
-		u = ib_umem_get(&dev->ib_dev, start, length, access_flags);
-		if (IS_ERR(u)) {
-			mlx5_ib_dbg(dev, "umem get failed (%ld)\n", PTR_ERR(u));
-			return PTR_ERR(u);
+			return ERR_CAST(odp);
 		}
-
-		mlx5_ib_cont_pages(u, start, MLX5_MKEY_PAGE_SHIFT_MASK, npages,
-				   page_shift, ncont, order);
+		return &odp->umem;
 	}
 
-	if (!*npages) {
-		mlx5_ib_warn(dev, "avoid zero region\n");
-		ib_umem_release(u);
-		return -EINVAL;
+	u = ib_umem_get(&dev->ib_dev, start, length, access_flags);
+	if (IS_ERR(u)) {
+		mlx5_ib_dbg(dev, "umem get failed (%ld)\n", PTR_ERR(u));
+		return u;
 	}
-
-	*umem = u;
-
-	mlx5_ib_dbg(dev, "npages %d, ncont %d, order %d, page_shift %d\n",
-		    *npages, *ncont, *order, *page_shift);
-
-	return 0;
+	return u;
 }
 
 static void mlx5_ib_umr_done(struct ib_cq *cq, struct ib_wc *wc)
@@ -969,15 +944,21 @@ static struct mlx5_cache_ent *mr_cache_ent_from_order(struct mlx5_ib_dev *dev,
 	return &cache->ent[order];
 }
 
-static struct mlx5_ib_mr *
-alloc_mr_from_cache(struct ib_pd *pd, struct ib_umem *umem, u64 virt_addr,
-		    u64 len, int npages, int page_shift, unsigned int order,
-		    int access_flags)
+static struct mlx5_ib_mr *alloc_mr_from_cache(struct ib_pd *pd,
+					      struct ib_umem *umem, u64 iova,
+					      int access_flags)
 {
 	struct mlx5_ib_dev *dev = to_mdev(pd->device);
-	struct mlx5_cache_ent *ent = mr_cache_ent_from_order(dev, order);
+	struct mlx5_cache_ent *ent;
 	struct mlx5_ib_mr *mr;
+	int npages;
+	int page_shift;
+	int ncont;
+	int order;
 
+	mlx5_ib_cont_pages(umem, iova, MLX5_MKEY_PAGE_SHIFT_MASK, &npages,
+			   &page_shift, &ncont, &order);
+	ent = mr_cache_ent_from_order(dev, order);
 	if (!ent)
 		return ERR_PTR(-E2BIG);
 	mr = get_cache_mr(ent);
@@ -991,9 +972,10 @@ alloc_mr_from_cache(struct ib_pd *pd, struct ib_umem *umem, u64 virt_addr,
 	mr->umem = umem;
 	mr->access_flags = access_flags;
 	mr->desc_size = sizeof(struct mlx5_mtt);
-	mr->mmkey.iova = virt_addr;
-	mr->mmkey.size = len;
+	mr->mmkey.iova = iova;
+	mr->mmkey.size = umem->length;
 	mr->mmkey.pd = to_mpd(pd)->pdn;
+	mr->page_shift = page_shift;
 
 	return mr;
 }
@@ -1152,14 +1134,15 @@ int mlx5_ib_update_xlt(struct mlx5_ib_mr *mr, u64 idx, int npages,
  * Else, the given ibmr will be used.
  */
 static struct mlx5_ib_mr *reg_create(struct ib_mr *ibmr, struct ib_pd *pd,
-				     u64 virt_addr, u64 length,
-				     struct ib_umem *umem, int npages,
-				     int page_shift, int access_flags,
-				     bool populate)
+				     struct ib_umem *umem, u64 iova,
+				     int access_flags, bool populate)
 {
 	struct mlx5_ib_dev *dev = to_mdev(pd->device);
 	struct mlx5_ib_mr *mr;
+	int page_shift;
 	__be64 *pas;
+	int npages;
+	int ncont;
 	void *mkc;
 	int inlen;
 	u32 *in;
@@ -1170,6 +1153,10 @@ static struct mlx5_ib_mr *reg_create(struct ib_mr *ibmr, struct ib_pd *pd,
 	if (!mr)
 		return ERR_PTR(-ENOMEM);
 
+	mlx5_ib_cont_pages(umem, iova, MLX5_MKEY_PAGE_SHIFT_MASK, &npages,
+			   &page_shift, &ncont, NULL);
+
+	mr->page_shift = page_shift;
 	mr->ibmr.pd = pd;
 	mr->access_flags = access_flags;
 
@@ -1206,17 +1193,17 @@ static struct mlx5_ib_mr *reg_create(struct ib_mr *ibmr, struct ib_pd *pd,
 	MLX5_SET(mkc, mkc, lr, 1);
 	MLX5_SET(mkc, mkc, umr_en, 1);
 
-	MLX5_SET64(mkc, mkc, start_addr, virt_addr);
-	MLX5_SET64(mkc, mkc, len, length);
+	MLX5_SET64(mkc, mkc, start_addr, iova);
+	MLX5_SET64(mkc, mkc, len, umem->length);
 	MLX5_SET(mkc, mkc, pd, to_mpd(pd)->pdn);
 	MLX5_SET(mkc, mkc, bsf_octword_size, 0);
 	MLX5_SET(mkc, mkc, translations_octword_size,
-		 get_octo_len(virt_addr, length, page_shift));
+		 get_octo_len(iova, umem->length, page_shift));
 	MLX5_SET(mkc, mkc, log_page_size, page_shift);
 	MLX5_SET(mkc, mkc, qpn, 0xffffff);
 	if (populate) {
 		MLX5_SET(create_mkey_in, in, translations_octword_actual_size,
-			 get_octo_len(virt_addr, length, page_shift));
+			 get_octo_len(iova, umem->length, page_shift));
 	}
 
 	err = mlx5_ib_create_mkey(dev, &mr->mmkey, in, inlen);
@@ -1353,10 +1340,6 @@ struct ib_mr *mlx5_ib_reg_user_mr(struct ib_pd *pd, u64 start, u64 length,
 	struct mlx5_ib_mr *mr = NULL;
 	bool use_umr;
 	struct ib_umem *umem;
-	int page_shift;
-	int npages;
-	int ncont;
-	int order;
 	int err;
 
 	if (!IS_ENABLED(CONFIG_INFINIBAND_USER_MEM))
@@ -1379,28 +1362,25 @@ struct ib_mr *mlx5_ib_reg_user_mr(struct ib_pd *pd, u64 start, u64 length,
 		return &mr->ibmr;
 	}
 
-	err = mr_umem_get(dev, start, length, access_flags, &umem,
-			  &npages, &page_shift, &ncont, &order);
-
-	if (err < 0)
-		return ERR_PTR(err);
+	umem = mr_umem_get(dev, start, length, access_flags);
+	if (IS_ERR(umem))
+		return ERR_CAST(umem);
 
 	use_umr = mlx5_ib_can_use_umr(dev, umem->length, true, access_flags);
 	/* ODP always uses UMR */
 	if (!use_umr && (access_flags & IB_ACCESS_ON_DEMAND))
 		return ERR_PTR(-EINVAL);
 
 	if (use_umr) {
-		mr = alloc_mr_from_cache(pd, umem, virt_addr, length, ncont,
-					 page_shift, order, access_flags);
+		mr = alloc_mr_from_cache(pd, umem, virt_addr, access_flags);
 		if (IS_ERR(mr))
 			mr = NULL;
 	}
 
 	if (!mr) {
 		mutex_lock(&dev->slow_path_mutex);
-		mr = reg_create(NULL, pd, virt_addr, length, umem, ncont,
-				page_shift, access_flags, !use_umr);
+		mr = reg_create(NULL, pd, umem, virt_addr, access_flags,
+				!use_umr);
 		mutex_unlock(&dev->slow_path_mutex);
 	}
 
@@ -1421,8 +1401,8 @@ struct ib_mr *mlx5_ib_reg_user_mr(struct ib_pd *pd, u64 start, u64 length,
 		if (access_flags & IB_ACCESS_ON_DEMAND)
 			update_xlt_flags |= MLX5_IB_UPD_XLT_ZAP;
 
-		err = mlx5_ib_update_xlt(mr, 0, ncont, page_shift,
-					 update_xlt_flags);
+		err = mlx5_ib_update_xlt(mr, 0, mlx5_mr_num_pages(mr),
+					 mr->page_shift, update_xlt_flags);
 
 		if (err) {
 			dereg_mr(dev, mr);
@@ -1507,11 +1487,7 @@ int mlx5_ib_rereg_user_mr(struct ib_mr *ib_mr, int flags, u64 start,
 	int access_flags = flags & IB_MR_REREG_ACCESS ?
 			    new_access_flags :
 			    mr->access_flags;
-	int page_shift = 0;
 	int upd_flags = 0;
-	int npages = 0;
-	int ncont = 0;
-	int order = 0;
 	u64 addr, len;
 	int err;
 
@@ -1541,12 +1517,11 @@ int mlx5_ib_rereg_user_mr(struct ib_mr *ib_mr, int flags, u64 start,
 		atomic_sub(ib_umem_num_pages(mr->umem),
 			   &dev->mdev->priv.reg_pages);
 		ib_umem_release(mr->umem);
-		mr->umem = NULL;
-
-		err = mr_umem_get(dev, addr, len, access_flags, &mr->umem,
-				  &npages, &page_shift, &ncont, &order);
-		if (err)
+		mr->umem = mr_umem_get(dev, addr, len, access_flags);
+		if (IS_ERR(mr->umem)) {
+			mr->umem = NULL;
 			goto err;
+		}
 		atomic_add(ib_umem_num_pages(mr->umem),
 			   &dev->mdev->priv.reg_pages);
 	}
@@ -1562,9 +1537,7 @@ int mlx5_ib_rereg_user_mr(struct ib_mr *ib_mr, int flags, u64 start,
 		if (err)
 			goto err;
 
-		mr = reg_create(ib_mr, pd, addr, len, mr->umem, ncont,
-				page_shift, access_flags, true);
-
+		mr = reg_create(ib_mr, pd, mr->umem, addr, access_flags, true);
 		if (IS_ERR(mr)) {
 			err = PTR_ERR(mr);
 			mr = to_mmr(ib_mr);
@@ -1586,8 +1559,8 @@ int mlx5_ib_rereg_user_mr(struct ib_mr *ib_mr, int flags, u64 start,
 				upd_flags |= MLX5_IB_UPD_XLT_PD;
 			if (flags & IB_MR_REREG_ACCESS)
 				upd_flags |= MLX5_IB_UPD_XLT_ACCESS;
-			err = mlx5_ib_update_xlt(mr, 0, npages, page_shift,
-						 upd_flags);
+			err = mlx5_ib_update_xlt(mr, 0, mlx5_mr_num_pages(mr),
+						 mr->page_shift, upd_flags);
 		} else {
 			err = rereg_umr(pd, mr, access_flags, flags);
 		}

include/rdma/ib_umem.h

@@ -33,11 +33,16 @@ static inline int ib_umem_offset(struct ib_umem *umem)
 	return umem->address & ~PAGE_MASK;
 }
 
+static inline size_t ib_umem_num_blocks(struct ib_umem *umem,
+					unsigned int block_shift)
+{
+	return (ALIGN(umem->address + umem->length, 1UL << block_shift) -
+		ALIGN_DOWN(umem->address, 1UL << block_shift)) >> block_shift;
+}
+
 static inline size_t ib_umem_num_pages(struct ib_umem *umem)
 {
-	return (ALIGN(umem->address + umem->length, PAGE_SIZE) -
-		ALIGN_DOWN(umem->address, PAGE_SIZE)) >>
-	       PAGE_SHIFT;
+	return ib_umem_num_blocks(umem, PAGE_SHIFT);
 }
 
 #ifdef CONFIG_INFINIBAND_USER_MEM