Merge pull request #9299 from ivankochin/uct/ib/avoid-using-ibv-mr-addr-and-fix-relaxed-order-memh-sizev1.15

UCT/IB: avoid using ibv mr addr and fix relaxed order memh size - v1.15

Author: yosefe

src/uct/ib/base/ib_md.c

@@ -1258,7 +1258,7 @@ static void uct_ib_md_release_device_config(uct_ib_md_t *md)
     ucs_free(md->custom_devices.specs);
 }
 
-static ucs_status_t
+ucs_status_t
 uct_ib_md_parse_reg_methods(uct_ib_md_t *md,
                             const uct_ib_md_config_t *md_config)
 {
@@ -1524,9 +1524,20 @@ ucs_status_t uct_ib_md_open(uct_component_t *component, const char *md_name,
     return status;
 }
 
+static void uct_ib_md_init_memh_size(uct_ib_md_t *md, size_t memh_base_size,
+                                     size_t mr_size)
+{
+    int num_mrs = md->relaxed_order ?
+                  2 /* UCT_IB_MR_DEFAULT and UCT_IB_MR_STRICT_ORDER */ :
+                  1 /* UCT_IB_MR_DEFAULT */;
+
+    md->memh_struct_size = memh_base_size + (mr_size * num_mrs);
+}
+
 void uct_ib_md_parse_relaxed_order(uct_ib_md_t *md,
                                    const uct_ib_md_config_t *md_config,
-                                   int is_supported)
+                                   int is_supported,
+                                   size_t memh_base_size, size_t mr_size)
 {
     int have_relaxed_order = (IBV_ACCESS_RELAXED_ORDERING != 0) && is_supported;
 
@@ -1546,20 +1557,13 @@ void uct_ib_md_parse_relaxed_order(uct_ib_md_t *md,
                             ucs_cpu_prefer_relaxed_order();
     }
 
+    /* Reset the memory handler size */
+    uct_ib_md_init_memh_size(md, memh_base_size, mr_size);
+
     ucs_debug("%s: relaxed order memory access is %sabled",
               uct_ib_device_name(&md->dev), md->relaxed_order ? "en" : "dis");
 }
 
-static void uct_ib_md_init_memh_size(uct_ib_md_t *md, size_t memh_base_size,
-                                     size_t mr_size)
-{
-    int num_mrs = md->relaxed_order ?
-                  2 /* UCT_IB_MR_DEFAULT and UCT_IB_MR_STRICT_ORDER */ :
-                  1 /* UCT_IB_MR_DEFAULT */;
-
-    md->memh_struct_size = memh_base_size + (mr_size * num_mrs);
-}
-
 static void uct_ib_check_gpudirect_driver(uct_ib_md_t *md, const char *file,
                                           ucs_memory_type_t mem_type)
 {
@@ -1601,8 +1605,7 @@ static void uct_ib_md_check_dmabuf(uct_ib_md_t *md)
 
 ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
                                    struct ibv_device *ib_device,
-                                   const uct_ib_md_config_t *md_config,
-                                   size_t memh_base_size, size_t mr_size)
+                                   const uct_ib_md_config_t *md_config)
 {
     ucs_status_t status;
 
@@ -1615,8 +1618,6 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
                           UCT_MD_FLAG_ADVISE;
     md->relaxed_order   = 0;
 
-    uct_ib_md_init_memh_size(md, memh_base_size, mr_size);
-
     /* Create statistics */
     status = UCS_STATS_NODE_ALLOC(&md->stats, &uct_ib_md_stats_class,
                                   ucs_stats_get_root(), "%s-%p",
@@ -1668,11 +1669,6 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
         goto err_cleanup_device;
     }
 
-    status = uct_ib_md_parse_reg_methods(md, md_config);
-    if (status != UCS_OK) {
-        goto err_cleanup_device;
-    }
-
     md->dev.max_zcopy_log_sge = INT_MAX;
     if (md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) {
         md->dev.max_zcopy_log_sge = 1;
@@ -1843,9 +1839,7 @@ static ucs_status_t uct_ib_verbs_md_open(struct ibv_device *ibv_device,
         goto err_device_config_release;
     }
 
-    status  = uct_ib_md_open_common(md, ibv_device, md_config,
-                                    sizeof(uct_ib_verbs_mem_t),
-                                    sizeof(uct_ib_mr_t));
+    status  = uct_ib_md_open_common(md, ibv_device, md_config);
     if (status != UCS_OK) {
         goto err_md_free;
     }
@@ -1855,7 +1849,14 @@ static ucs_status_t uct_ib_verbs_md_open(struct ibv_device *ibv_device,
     md->flush_rkey = UCT_IB_MD_INVALID_FLUSH_RKEY;
 
     uct_ib_md_ece_check(md);
-    uct_ib_md_parse_relaxed_order(md, md_config, 0);
+    uct_ib_md_parse_relaxed_order(md, md_config, 0,
+                                  sizeof(uct_ib_verbs_mem_t),
+                                  sizeof(uct_ib_mr_t));
+
+    status = uct_ib_md_parse_reg_methods(md, md_config);
+    if (status != UCS_OK) {
+        goto err_md_free;
+    }
 
     *p_md = md;
     return UCS_OK;

src/uct/ib/base/ib_md.h

@@ -630,14 +630,18 @@ ucs_status_t uct_ib_md_open(uct_component_t *component, const char *md_name,
 
 void uct_ib_md_parse_relaxed_order(uct_ib_md_t *md,
                                    const uct_ib_md_config_t *md_config,
-                                   int is_supported);
+                                   int is_supported,
+                                   size_t memh_base_size, size_t mr_size);
 
 int uct_ib_device_is_accessible(struct ibv_device *device);
 
+ucs_status_t
+uct_ib_md_parse_reg_methods(uct_ib_md_t *md,
+                            const uct_ib_md_config_t *md_config);
+
 ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
                                    struct ibv_device *ib_device,
-                                   const uct_ib_md_config_t *md_config,
-                                   size_t mem_size, size_t mr_size);
+                                   const uct_ib_md_config_t *md_config);
 
 void uct_ib_md_close_common(uct_ib_md_t *md);
 

src/uct/ib/mlx5/dv/ib_mlx5dv_md.c

@@ -34,6 +34,7 @@ typedef union uct_ib_mlx5_mr {
 
 typedef struct uct_ib_mlx5_mem {
     uct_ib_mem_t               super;
+    void                       *address;
 #if HAVE_DEVX
     struct mlx5dv_devx_obj     *atomic_dvmr;
     struct mlx5dv_devx_obj     *indirect_dvmr;
@@ -51,6 +52,8 @@ uct_ib_mlx5_reg_key(uct_ib_md_t *md, void *address, size_t length,
 {
     uct_ib_mlx5_mem_t *memh = ucs_derived_of(ib_memh, uct_ib_mlx5_mem_t);
 
+    memh->address = address;
+
     return uct_ib_reg_key_impl(md, address, length, access_flags, dmabuf_fd,
                                dmabuf_offset, ib_memh,
                                &memh->mrs[mr_type].super, mr_type, silent);
@@ -208,12 +211,13 @@ static ucs_status_t uct_ib_mlx5_devx_reg_ksm(uct_ib_mlx5_md_t *md, int atomic,
 }
 
 static ucs_status_t
-uct_ib_mlx5_devx_reg_ksm_data(uct_ib_mlx5_md_t *md, int atomic,
+uct_ib_mlx5_devx_reg_ksm_data(uct_ib_mlx5_md_t *md, int atomic, void *address,
                               uct_ib_mlx5_ksm_data_t *ksm_data,
                               size_t length, off_t off,
                               struct mlx5dv_devx_obj **mr_p,
                               uint32_t *mkey)
 {
+    void *mr_address = address;
     ucs_status_t status;
     char *in;
     void *klm;
@@ -226,14 +230,15 @@ uct_ib_mlx5_devx_reg_ksm_data(uct_ib_mlx5_md_t *md, int atomic,
 
     klm = UCT_IB_MLX5DV_ADDR_OF(create_mkey_in, in, klm_pas_mtt);
     for (i = 0; i < ksm_data->mr_num; i++) {
-        UCT_IB_MLX5DV_SET64(klm, klm, address, (intptr_t)ksm_data->mrs[i]->addr);
+        UCT_IB_MLX5DV_SET64(klm, klm, address, (uintptr_t)mr_address);
         UCT_IB_MLX5DV_SET(klm, klm, byte_count, ksm_data->mrs[i]->length);
         UCT_IB_MLX5DV_SET(klm, klm, mkey, ksm_data->mrs[i]->lkey);
-        klm = UCS_PTR_BYTE_OFFSET(klm, UCT_IB_MLX5DV_ST_SZ_BYTES(klm));
+        klm        = UCS_PTR_BYTE_OFFSET(klm, UCT_IB_MLX5DV_ST_SZ_BYTES(klm));
+        mr_address = UCS_PTR_BYTE_OFFSET(mr_address, ksm_data->mrs[i]->length);
     }
 
     status = uct_ib_mlx5_devx_reg_ksm(md, atomic,
-                                      (intptr_t)ksm_data->mrs[0]->addr + off,
+                                      (uintptr_t)address + off,
                                       length, ksm_data->mr_num,
                                       ksm_data->mrs[0]->length, in, mr_p,
                                       mkey);
@@ -276,14 +281,15 @@ uct_ib_mlx5_devx_reg_ksm_data_addr(uct_ib_mlx5_md_t *md, struct ibv_mr *mr,
 static ucs_status_t
 uct_ib_mlx5_devx_reg_ksm_data_contig(uct_ib_mlx5_md_t *md,
                                      uct_ib_mlx5_mr_t *mr, off_t off,
-                                     int atomic, struct mlx5dv_devx_obj **mr_p,
+                                     void *address, int atomic,
+                                     struct mlx5dv_devx_obj **mr_p,
                                      uint32_t *mkey)
 {
-    intptr_t addr = (intptr_t)mr->super.ib->addr & ~(UCT_IB_MD_MAX_MR_SIZE - 1);
+    intptr_t addr = (intptr_t)address & ~(UCT_IB_MD_MAX_MR_SIZE - 1);
     /* FW requires indirect atomic MR addr and length to be aligned
      * to max supported atomic argument size */
     size_t length = ucs_align_up(mr->super.ib->length +
-                                 (intptr_t)mr->super.ib->addr - addr,
+                                 (intptr_t)address - addr,
                                  md->super.dev.atomic_align);
     /* add off to workaround CREATE_MKEY range check issue */
     int list_size = ucs_div_round_up(length + off, UCT_IB_MD_MAX_MR_SIZE);
@@ -445,8 +451,8 @@ static ucs_status_t uct_ib_mlx5_devx_reg_indirect_key(uct_ib_md_t *ibmd,
 
     do {
         status = uct_ib_mlx5_devx_reg_ksm_data_contig(
-                md, &memh->mrs[UCT_IB_MR_DEFAULT], 0, 0, &memh->indirect_dvmr,
-                &memh->super.indirect_rkey);
+                md, &memh->mrs[UCT_IB_MR_DEFAULT], 0, memh->address, 0,
+                &memh->indirect_dvmr, &memh->super.indirect_rkey);
         if (status != UCS_OK) {
             break;
         }
@@ -549,23 +555,24 @@ static ucs_status_t uct_ib_mlx5_devx_reg_atomic_key(uct_ib_md_t *ibmd,
     is_atomic = memh->super.flags & UCT_IB_MEM_ACCESS_REMOTE_ATOMIC;
 
     if (memh->super.flags & UCT_IB_MEM_MULTITHREADED) {
-        return uct_ib_mlx5_devx_reg_ksm_data(md, is_atomic, mr->ksm_data,
+        return uct_ib_mlx5_devx_reg_ksm_data(md, is_atomic, memh->address,
+                                             mr->ksm_data,
                                              mr->ksm_data->length,
                                              uct_ib_md_atomic_offset(mr_id),
                                              &memh->atomic_dvmr,
                                              &memh->super.atomic_rkey);
     }
 
     status = uct_ib_mlx5_devx_reg_ksm_data_contig(
-            md, mr, uct_ib_md_atomic_offset(mr_id), is_atomic,
+            md, mr, uct_ib_md_atomic_offset(mr_id), memh->address, is_atomic,
             &memh->atomic_dvmr, &memh->super.atomic_rkey);
     if (status != UCS_OK) {
         return status;
     }
 
     ucs_debug("KSM registered memory %p..%p offset 0x%x%s on %s rkey 0x%x",
-              mr->super.ib->addr,
-              UCS_PTR_BYTE_OFFSET(mr->super.ib->addr, mr->super.ib->length),
+              memh->address,
+              UCS_PTR_BYTE_OFFSET(memh->address, mr->super.ib->length),
               uct_ib_md_atomic_offset(mr_id), is_atomic ? " atomic" : "",
               uct_ib_device_name(&md->super.dev), memh->super.atomic_rkey);
     return status;
@@ -627,8 +634,8 @@ static ucs_status_t uct_ib_mlx5_devx_reg_multithreaded(uct_ib_md_t *ibmd,
         goto err;
     }
 
-    status = uct_ib_mlx5_devx_reg_ksm_data(md, is_atomic, ksm_data, length, 0,
-                                           &ksm_data->dvmr, &mkey);
+    status = uct_ib_mlx5_devx_reg_ksm_data(md, is_atomic, address, ksm_data,
+                                           length, 0, &ksm_data->dvmr, &mkey);
     if (status != UCS_OK) {
         goto err_dereg;
     }
@@ -639,6 +646,8 @@ static ucs_status_t uct_ib_mlx5_devx_reg_multithreaded(uct_ib_md_t *ibmd,
     if (mr_type == UCT_IB_MR_DEFAULT) {
         uct_ib_memh_init_keys(ib_memh, mkey, mkey);
     }
+
+    memh->address = address;
     return UCS_OK;
 
 err_dereg:
@@ -1225,9 +1234,8 @@ static ucs_status_t uct_ib_mlx5_devx_md_open(struct ibv_device *ibv_device,
     }
 
     md->super.ops = &uct_ib_mlx5_devx_md_ops;
-    status        = uct_ib_md_open_common(&md->super, ibv_device, md_config,
-                                          sizeof(uct_ib_mlx5_mem_t),
-                                          sizeof(uct_ib_mlx5_mr_t));
+
+    status = uct_ib_md_open_common(&md->super, ibv_device, md_config);
     if (status != UCS_OK) {
         goto err_lru_cleanup;
     }
@@ -1274,7 +1282,14 @@ static ucs_status_t uct_ib_mlx5_devx_md_open(struct ibv_device *ibv_device,
 
     /* Enable relaxed order only if we would be able to create an indirect key
        (with offset) for strict order access */
-    uct_ib_md_parse_relaxed_order(&md->super, md_config, ksm_atomic);
+    uct_ib_md_parse_relaxed_order(&md->super, md_config, ksm_atomic,
+                                  sizeof(uct_ib_mlx5_mem_t),
+                                  sizeof(uct_ib_mlx5_mr_t));
+
+    status = uct_ib_md_parse_reg_methods(&md->super, md_config);
+    if (status != UCS_OK) {
+        goto err_dbrec_mpool_cleanup;
+    }
 
     uct_ib_mlx5_devx_init_flush_mr(md);
 
@@ -1454,18 +1469,18 @@ uct_ib_mlx5_devx_reg_exported_key(uct_ib_md_t *ib_md, uct_ib_mem_t *ib_memh)
     struct mlx5dv_devx_umem_in umem_in;
     ucs_status_t status;
     void *access_key;
-    void *address, *aligned_address;
+    void *aligned_address;
     size_t length;
     void *mkc;
 
-    address = memh->mrs[UCT_IB_MR_DEFAULT].super.ib->addr;
     length  = memh->mrs[UCT_IB_MR_DEFAULT].super.ib->length;
 
     /* register umem */
-    umem_in.addr        = address;
+    umem_in.addr        = memh->address;
     umem_in.size        = length;
     umem_in.access      = UCT_IB_MLX5_MD_UMEM_ACCESS;
-    aligned_address     = ucs_align_down_pow2_ptr(address, ucs_get_page_size());
+    aligned_address     = ucs_align_down_pow2_ptr(memh->address,
+                                                  ucs_get_page_size());
     umem_in.pgsz_bitmap = UCS_MASK(ucs_ffs64((uint64_t)aligned_address) + 1);
     umem_in.comp_mask   = 0;
 
@@ -1496,7 +1511,7 @@ uct_ib_mlx5_devx_reg_exported_key(uct_ib_md_t *ib_md, uct_ib_mem_t *ib_memh)
     UCT_IB_MLX5DV_SET(mkc, mkc, qpn, 0xffffff);
     UCT_IB_MLX5DV_SET(mkc, mkc, pd, uct_ib_mlx5_devx_md_get_pdn(md));
     UCT_IB_MLX5DV_SET(mkc, mkc, mkey_7_0, md->mkey_tag);
-    UCT_IB_MLX5DV_SET64(mkc, mkc, start_addr, (intptr_t)address);
+    UCT_IB_MLX5DV_SET64(mkc, mkc, start_addr, (intptr_t)memh->address);
     UCT_IB_MLX5DV_SET64(mkc, mkc, len, length);
 
     ucs_assertv(memh->cross_mr == NULL, "memh %p cross_mr %p", memh,
@@ -1775,19 +1790,24 @@ static ucs_status_t uct_ib_mlx5dv_md_open(struct ibv_device *ibv_device,
     md->super.ops        = &uct_ib_mlx5_md_ops;
     md->max_rd_atomic_dc = IBV_DEV_ATTR(dev, max_qp_rd_atom);
     status               = uct_ib_md_open_common(&md->super, ibv_device,
-                                                 md_config,
-                                                 sizeof(uct_ib_mlx5_mem_t),
-                                                 sizeof(uct_ib_mlx5_mr_t));
+                                                 md_config);
     if (status != UCS_OK) {
         goto err_md_free;
     }
 
     dev->flags    |= UCT_IB_DEVICE_FLAG_MLX5_PRM;
     md->super.name = UCT_IB_MD_NAME(mlx5);
 
-    uct_ib_md_parse_relaxed_order(&md->super, md_config, 0);
+    uct_ib_md_parse_relaxed_order(&md->super, md_config, 0,
+                                  sizeof(uct_ib_mlx5_mem_t),
+                                  sizeof(uct_ib_mlx5_mr_t));
     uct_ib_md_ece_check(&md->super);
 
+    status = uct_ib_md_parse_reg_methods(&md->super, md_config);
+    if (status != UCS_OK) {
+        goto err_md_free;
+    }
+
     md->super.flush_rkey = uct_ib_mlx5_flush_rkey_make();
 
     /* cppcheck-suppress autoVariables */