1 files changed, 9 insertions, 5 deletions

diff --git a/drivers/infiniband/hw/mlx5/mr.c b/drivers/infiniband/hw/mlx5/mr.c
index a5fad3e87ff7..95a36e9ea552 100644
--- a/drivers/infiniband/hw/mlx5/mr.c
+++ b/drivers/infiniband/hw/mlx5/mr.c
@@ -839,7 +839,8 @@ static int mr_umem_get(struct ib_pd *pd, u64 start, u64 length,
         *umem = ib_umem_get(pd->uobject->context, start, length,
                             access_flags, 0);
         err = PTR_ERR_OR_ZERO(*umem);
-        if (err < 0) {
+        if (err) {
+                *umem = NULL;
                 mlx5_ib_err(dev, "umem get failed (%d)\n", err);
                 return err;
         }
@@ -1416,6 +1417,7 @@ int mlx5_ib_rereg_user_mr(struct ib_mr *ib_mr, int flags, u64 start,
                 if (err) {
                         mlx5_ib_warn(dev, "Failed to rereg UMR\n");
                         ib_umem_release(mr->umem);
+                        mr->umem = NULL;
                         clean_mr(dev, mr);
                         return err;
                 }
@@ -1499,14 +1501,11 @@ static int clean_mr(struct mlx5_ib_dev *dev, struct mlx5_ib_mr *mr)
                 u32 key = mr->mmkey.key;
 
                 err = destroy_mkey(dev, mr);
-                kfree(mr);
                 if (err) {
                         mlx5_ib_warn(dev, "failed to destroy mkey 0x%x (%d)\n",
                                      key, err);
                         return err;
                 }
-        } else {
-                mlx5_mr_cache_free(dev, mr);
         }
 
         return 0;
@@ -1549,6 +1548,11 @@ static int dereg_mr(struct mlx5_ib_dev *dev, struct mlx5_ib_mr *mr)
                 atomic_sub(npages, &dev->mdev->priv.reg_pages);
         }
 
+        if (!mr->allocated_from_cache)
+                kfree(mr);
+        else
+                mlx5_mr_cache_free(dev, mr);
+
         return 0;
 }
 
@@ -1817,7 +1821,6 @@ mlx5_ib_sg_to_klms(struct mlx5_ib_mr *mr,
 
         mr->ibmr.iova = sg_dma_address(sg) + sg_offset;
         mr->ibmr.length = 0;
-        mr->ndescs = sg_nents;
 
         for_each_sg(sgl, sg, sg_nents, i) {
                 if (unlikely(i >= mr->max_descs))
@@ -1829,6 +1832,7 @@ mlx5_ib_sg_to_klms(struct mlx5_ib_mr *mr,
 
                 sg_offset = 0;
         }
+        mr->ndescs = i;
 
         if (sg_offset_p)
                 *sg_offset_p = sg_offset;