2 files changed, 18 insertions, 1 deletions

diff --git a/drivers/infiniband/hw/mlx5/main.c b/drivers/infiniband/hw/mlx5/main.c
index 829c37c95d60..a55bf05c8522 100644
--- a/drivers/infiniband/hw/mlx5/main.c
+++ b/drivers/infiniband/hw/mlx5/main.c
@@ -942,6 +942,13 @@ static struct ib_ucontext *mlx5_ib_alloc_ucontext(struct ib_device *ibdev,
         context->ibucontext.invalidate_range = &mlx5_ib_invalidate_range;
 #endif
 
+        if (MLX5_CAP_GEN(dev->mdev, log_max_transport_domain)) {
+                err = mlx5_core_alloc_transport_domain(dev->mdev,
+                                                       &context->tdn);
+                if (err)
+                        goto out_uars;
+        }
+
         INIT_LIST_HEAD(&context->db_page_list);
         mutex_init(&context->db_page_mutex);
 
@@ -964,7 +971,7 @@ static struct ib_ucontext *mlx5_ib_alloc_ucontext(struct ib_device *ibdev,
 
         err = ib_copy_to_udata(udata, &resp, resp.response_length);
         if (err)
-                goto out_uars;
+                goto out_td;
 
         uuari->ver = ver;
         uuari->num_low_latency_uuars = req.num_low_latency_uuars;
@@ -974,6 +981,10 @@ static struct ib_ucontext *mlx5_ib_alloc_ucontext(struct ib_device *ibdev,
 
         return &context->ibucontext;
 
+out_td:
+        if (MLX5_CAP_GEN(dev->mdev, log_max_transport_domain))
+                mlx5_core_dealloc_transport_domain(dev->mdev, context->tdn);
+
 out_uars:
         for (i--; i >= 0; i--)
                 mlx5_cmd_free_uar(dev->mdev, uars[i].index);
@@ -998,6 +1009,9 @@ static int mlx5_ib_dealloc_ucontext(struct ib_ucontext *ibcontext)
         struct mlx5_uuar_info *uuari = &context->uuari;
         int i;
 
+        if (MLX5_CAP_GEN(dev->mdev, log_max_transport_domain))
+                mlx5_core_dealloc_transport_domain(dev->mdev, context->tdn);
+
         for (i = 0; i < uuari->num_uars; i++) {
                 if (mlx5_cmd_free_uar(dev->mdev, uuari->uars[i].index))
                         mlx5_ib_warn(dev, "failed to free UAR 0x%x\n", uuari->uars[i].index);
diff --git a/drivers/infiniband/hw/mlx5/mlx5_ib.h b/drivers/infiniband/hw/mlx5/mlx5_ib.h
index a6a57de278b4..2ccc3d95d574 100644
--- a/drivers/infiniband/hw/mlx5/mlx5_ib.h
+++ b/drivers/infiniband/hw/mlx5/mlx5_ib.h
@@ -42,6 +42,7 @@
 #include <linux/mlx5/qp.h>
 #include <linux/mlx5/srq.h>
 #include <linux/types.h>
+#include <linux/mlx5/transobj.h>
 
 #define mlx5_ib_dbg(dev, format, arg...)                                \
 pr_debug("%s:%s:%d:(pid %d): " format, (dev)->ib_dev.name, __func__,    \
@@ -110,6 +111,8 @@ struct mlx5_ib_ucontext {
         struct mutex            db_page_mutex;
         struct mlx5_uuar_info   uuari;
         u8                      cqe_version;
+        /* Transport Domain number */
+        u32                     tdn;
 };
 
 static inline struct mlx5_ib_ucontext *to_mucontext(struct ib_ucontext *ibucontext)