UCP/UCT/IB: Add per-endpoint traffic class (DSCP) for RC/RoCE (QOS)

AUTHORS

@@ -129,6 +129,8 @@ Xu Yifeng <yifeng.xyf@alibaba-inc.com>
 Yaser Afshar <yaser.afshar@intel.com>
 Yihua Xu <yihua.xu@intel.com>
 Yiltan Hassan Temucin <yiltan.temucin@amd.com>
+Yoel Benabou <ybenabou@nvidia.com>
+Yoel Benabou <ybenabou@nvidia.com>
 Yossi Itigin <yosefe@nvidia.com>
 Yuriy Shestakov <yuriis@mellanox.com>
 Zhenlong Ma <zhenlongm@nvidia.com>

bindings/java/src/main/java/org/openucx/jucx/ucp/UcpConstants.java

@@ -86,6 +86,7 @@ public class UcpConstants {
     static long UCP_EP_PARAM_FIELD_LOCAL_SOCK_ADDR;
     static long UCP_EP_PARAM_FIELD_CONN_REQUEST;
     static long UCP_EP_PARAM_FIELD_NAME;
+    static long UCP_EP_PARAM_FIELD_EP_TRAFFIC_CLASS;
 
     /**
      * UCP error handling mode.

bindings/java/src/main/native/ucp_constants.cc

@@ -65,6 +65,7 @@ Java_org_openucx_jucx_ucp_UcpConstants_loadConstants(JNIEnv *env, jclass cls)
     JUCX_DEFINE_LONG_CONSTANT(UCP_EP_PARAM_FIELD_LOCAL_SOCK_ADDR);
     JUCX_DEFINE_LONG_CONSTANT(UCP_EP_PARAM_FIELD_CONN_REQUEST);
     JUCX_DEFINE_LONG_CONSTANT(UCP_EP_PARAM_FIELD_NAME);
+    JUCX_DEFINE_LONG_CONSTANT(UCP_EP_PARAM_FIELD_EP_TRAFFIC_CLASS);
 
     // UCP error handling mode
     JUCX_DEFINE_INT_CONSTANT(UCP_ERR_HANDLING_MODE_PEER);

src/ucp/api/ucp.h

@@ -268,7 +268,10 @@ enum ucp_ep_params_field {
     /**< Connection request field */
     UCP_EP_PARAM_FIELD_CONN_REQUEST      = UCS_BIT(6),
     UCP_EP_PARAM_FIELD_NAME              = UCS_BIT(7), /**< Endpoint name */
-    UCP_EP_PARAM_FIELD_LOCAL_SOCK_ADDR   = UCS_BIT(8)  /**< Local socket Address */
+    UCP_EP_PARAM_FIELD_LOCAL_SOCK_ADDR   = UCS_BIT(
+            8), /**< Local socket Address */
+    UCP_EP_PARAM_FIELD_EP_TRAFFIC_CLASS  = UCS_BIT(
+            9) /**< Endpoint traffic class */
 };
 
 

src/ucp/api/ucp_def.h

@@ -779,6 +779,13 @@ typedef struct ucp_ep_params {
      */
     const char              *name;
 
+    /**
+     * Traffic class for this endpoint.
+     * This setting is optional. To enable it, the corresponding - @ref
+     * UCP_EP_PARAM_FIELD_EP_TRAFFIC_CLASS bit in the field mask must be set.
+     */
+    uint8_t                 ep_traffic_class;
+
     /**
      * The sockaddr to bind locally. Specifies the associated network device
      * to bind locally to establish new connections.
@@ -787,7 +794,6 @@ typedef struct ucp_ep_params {
      * UCP_EP_PARAM_FIELD_LOCAL_SOCK_ADDR bit in the field mask must be set.
      */
     ucs_sock_addr_t         local_sockaddr;
-
 } ucp_ep_params_t;
 
 

src/ucp/core/ucp_ep.c

@@ -667,6 +667,13 @@ ucp_ep_adjust_params(ucp_ep_h ep, const ucp_ep_params_t *params)
         ep->flags |= UCP_EP_FLAG_USER_DATA_PARAM;
     }
 
+    if (params->field_mask & UCP_EP_PARAM_FIELD_EP_TRAFFIC_CLASS) {
+        ep->ep_traffic_class = params->ep_traffic_class;
+        ep->flags           |= UCP_EP_FLAG_EP_TRAFFIC_CLASS;
+    } else {
+        ep->ep_traffic_class = UCP_EP_NO_TCLASS;
+    }
+
     return UCS_OK;
 }
 
@@ -1987,7 +1994,8 @@ ucp_ep_recovery_rebuild_p2p_lane(
     }
 
     status = ucp_wireup_ep_connect_to_ep_v2(ucp_ep_get_lane(ep, lane),
-                                            address_entry, ep_entry);
+                                            address_entry, ep_entry,
+                                            ep->ep_traffic_class, ep->flags);
     if (status != UCS_OK) {
         ucs_diag("ep %p: connect_to_ep_v2 failed for recovery p2p lane %d: %s",
                  ep, lane, ucs_status_string(status));

src/ucp/core/ucp_ep.h

@@ -104,7 +104,8 @@ enum {
                                                         while merging pending queues */
     UCP_EP_FLAG_CONNECT_PRE_REQ_QUEUED = UCS_BIT(9), /* Pre-Connection request was queued */
     UCP_EP_FLAG_CLOSED                 = UCS_BIT(10),/* EP was closed */
-    /* 11 bit is vacant for a flag */
+    UCP_EP_FLAG_EP_TRAFFIC_CLASS = UCS_BIT(11), /* Endpoint traffic class */
+    /* 12 bit is vacant for a flag */
     UCP_EP_FLAG_ERR_HANDLER_INVOKED    = UCS_BIT(12),/* error handler was called */
     UCP_EP_FLAG_INTERNAL               = UCS_BIT(13),/* the internal EP which holds
                                                         temporary wireup configuration or
@@ -581,6 +582,7 @@ typedef struct ucp_ep {
     /* Transports for every lane */
     uct_ep_h                      uct_eps[UCP_MAX_FAST_PATH_LANES];
     ucp_ep_ext_t                  *ext;                   /* Endpoint extension */
+    uint8_t                       ep_traffic_class; /* Endpoint traffic class */
 
 #if ENABLE_DEBUG_DATA
     char                          peer_name[UCP_WORKER_ADDRESS_NAME_MAX];

src/ucp/core/ucp_types.h

@@ -60,6 +60,8 @@ typedef uint16_t                     ucp_worker_cfg_index_t;
 #define UCP_WORKER_MAX_RKEY_CONFIG   UINT16_MAX
 #define UCP_WORKER_CFG_INDEX_NULL    UINT16_MAX
 
+#define UCP_EP_NO_TCLASS ((uint8_t)-1)
+
 
 /* Forward declarations */
 typedef struct ucp_request            ucp_request_t;

src/ucp/wireup/wireup.c

@@ -480,7 +480,9 @@ ucp_wireup_connect_local(ucp_ep_h ep,
         }
 
         status = ucp_wireup_ep_connect_to_ep_v2(ucp_ep_get_lane(ep, lane),
-                                                address_entry, ep_entry);
+                                                address_entry, ep_entry,
+                                                ep->ep_traffic_class,
+                                                ep->flags);
         if (status != UCS_OK) {
             goto out;
         }

src/ucp/wireup/wireup_ep.c

@@ -513,6 +513,14 @@ ucs_status_t ucp_wireup_ep_connect(uct_ep_h uct_ep, unsigned ep_init_flags,
                                UCT_EP_PARAM_FIELD_PATH_INDEX;
     uct_ep_params.path_index = path_index;
     uct_ep_params.iface      = ucp_worker_iface(worker, rsc_index)->iface;
+
+    if (ucp_ep->flags & UCP_EP_FLAG_EP_TRAFFIC_CLASS) {
+        uct_ep_params.field_mask      |= UCT_EP_PARAM_FIELD_EP_TRAFFIC_CLASS;
+        uct_ep_params.ep_traffic_class = ucp_ep->ep_traffic_class;
+    } else {
+        uct_ep_params.ep_traffic_class = UCT_EP_NO_TCLASS;
+    }
+
     status = uct_ep_create(&uct_ep_params, &next_ep);
     if (status != UCS_OK) {
         /* make Coverity happy */
@@ -695,13 +703,19 @@ void ucp_wireup_eps_pending_extract(ucp_ep_t *ucp_ep, ucs_queue_head_t *queue)
 ucs_status_t
 ucp_wireup_ep_connect_to_ep_v2(uct_ep_h tl_ep,
                                const ucp_address_entry_t *address_entry,
-                               const ucp_address_entry_ep_addr_t *ep_entry)
+                               const ucp_address_entry_ep_addr_t *ep_entry,
+                               uint8_t ep_traffic_class, ucp_ep_flags_t flags)
 {
     const uct_ep_connect_to_ep_params_t param = {
-        .field_mask = UCT_EP_CONNECT_TO_EP_PARAM_FIELD_DEVICE_ADDR_LENGTH |
-                      UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_ADDR_LENGTH,
+        .field_mask =
+                UCT_EP_CONNECT_TO_EP_PARAM_FIELD_DEVICE_ADDR_LENGTH |
+                UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_ADDR_LENGTH |
+                ((flags & UCP_EP_FLAG_EP_TRAFFIC_CLASS) ?
+                         UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_TRAFFIC_CLASS :
+                         0),
         .device_addr_length = address_entry->dev_addr_len,
-        .ep_addr_length     = ep_entry->len
+        .ep_addr_length     = ep_entry->len,
+        .ep_traffic_class   = ep_traffic_class
     };
     ucp_wireup_ep_t *wireup_ep                = ucp_wireup_ep(tl_ep);
     ucs_status_t status;

src/ucp/wireup/wireup_ep.h

@@ -141,6 +141,7 @@ void ucp_wireup_eps_pending_extract(ucp_ep_t *ucp_ep, ucs_queue_head_t *queue);
 ucs_status_t
 ucp_wireup_ep_connect_to_ep_v2(uct_ep_h tl_ep,
                                const ucp_address_entry_t *address_entry,
-                               const ucp_address_entry_ep_addr_t *ep_entry);
+                               const ucp_address_entry_ep_addr_t *ep_entry,
+                               uint8_t ep_traffic_class, ucp_ep_flags_t flags);
 
 #endif

src/uct/api/uct.h

@@ -990,7 +990,10 @@ enum uct_ep_params_field {
     UCT_EP_PARAM_FIELD_DEV_ADDR_LENGTH            = UCS_BIT(17),
 
     /** Enables @ref uct_ep_params::iface_addr_length */
-    UCT_EP_PARAM_FIELD_IFACE_ADDR_LENGTH          = UCS_BIT(18)
+    UCT_EP_PARAM_FIELD_IFACE_ADDR_LENGTH          = UCS_BIT(18),
+
+    /** Enables @ref uct_ep_params::ep_traffic_class */
+    UCT_EP_PARAM_FIELD_EP_TRAFFIC_CLASS           = UCS_BIT(19)
 };
 
 
@@ -1464,6 +1467,8 @@ struct uct_ep_params {
      * default minimum length according to the address buffer contents.
      */
     size_t                              iface_addr_length;
+
+    uint8_t                                  ep_traffic_class;
 };
 
 

src/uct/api/uct_def.h

@@ -28,6 +28,7 @@
 #define UCT_INVALID_RKEY           ((uintptr_t)(-1))
 #define UCT_INLINE_API             static UCS_F_ALWAYS_INLINE
 #define UCT_DMABUF_FD_INVALID      -1
+#define UCT_EP_NO_TCLASS           ((uint8_t)-1)
 
 
 /**

src/uct/api/v2/uct_v2.h

@@ -374,7 +374,10 @@ typedef enum {
     UCT_EP_CONNECT_TO_EP_PARAM_FIELD_DEVICE_ADDR_LENGTH = UCS_BIT(0),
 
     /** Endpoint address length */
-    UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_ADDR_LENGTH     = UCS_BIT(1)
+    UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_ADDR_LENGTH     = UCS_BIT(1),
+
+    /** Endpoint traffic class */
+    UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_TRAFFIC_CLASS   = UCS_BIT(2)
 } uct_ep_connect_to_ep_param_field_t;
 
 
@@ -730,6 +733,11 @@ typedef struct uct_ep_connect_to_ep_params {
      * default minimal length according to the address buffer contents.
      */
     size_t                        ep_addr_length;
+
+    /**
+     * Endpoint traffic class.
+     */
+    uint8_t                       ep_traffic_class;
 } uct_ep_connect_to_ep_params_t;
 
 

src/uct/ib/base/ib_md.c

@@ -491,6 +491,11 @@ ucs_status_t uct_ib_reg_mr(const uct_ib_md_t *md, void *address, size_t length,
     uint64_t flags;
     int dmabuf_fd;
 
+    ucs_info("inside uct_ib_reg_mr");
+    ucs_info("start time : %lu", start_time);
+    ucs_log(UCS_LOG_LEVEL_INFO, "inside uct_ib_reg_mr");
+    ucs_log(UCS_LOG_LEVEL_INFO, "start time : %lu", start_time);
+
     flags         = UCT_MD_MEM_REG_FIELD_VALUE(params, flags, FIELD_FLAGS, 0);
     dmabuf_fd     = UCS_PARAM_VALUE(UCT_MD_MEM_REG_FIELD, params, dmabuf_fd,
                                     DMABUF_FD, UCT_DMABUF_FD_INVALID);

src/uct/ib/mlx5/gdaki/gdaki.c

@@ -854,7 +854,8 @@ uct_rc_gdaki_ep_connect_to_ep_v2(uct_ep_h tl_ep,
                 *ucs_serialize_next(&ep_addr, uct_ib_uint24_t));
         status      = uct_rc_mlx5_iface_common_devx_connect_qp(
                 &iface->super, &channels[i].qp.super, dest_qp_num, &ah_attr,
-                path_mtu, path_index, iface->super.super.config.max_rd_atomic);
+                path_mtu, path_index, iface->super.super.config.max_rd_atomic,
+                params);
         if (status != UCS_OK) {
             return status;
         }

src/uct/ib/mlx5/gga/gga_mlx5.c

@@ -524,7 +524,8 @@ uct_gga_mlx5_ep_connect_to_ep_v2(uct_ep_h tl_ep,
     qp_num = uct_ib_unpack_uint24(gga_ep_addr->super.qp_num);
     status = uct_rc_mlx5_iface_common_devx_connect_qp(
             iface, &ep->super.tx.wq.super, qp_num, &ah_attr, path_mtu,
-            ep->super.super.path_index, iface->super.config.max_rd_atomic);
+            ep->super.super.path_index, iface->super.config.max_rd_atomic,
+            params);
     if (status != UCS_OK) {
         return status;
     }

src/uct/ib/mlx5/rc/rc_mlx5.h

@@ -209,7 +209,8 @@ ucs_status_t
 uct_rc_mlx5_ep_connect_qp(uct_rc_mlx5_iface_common_t *iface,
                           uct_ib_mlx5_qp_t *qp, uint32_t qp_num,
                           struct ibv_ah_attr *ah_attr, enum ibv_mtu path_mtu,
-                          uint8_t path_index);
+                          uint8_t path_index,
+                          const uct_ep_connect_to_ep_params_t *params);
 
 int uct_rc_mlx5_base_ep_is_connected(
         const uct_ep_h tl_ep, const uct_ep_is_connected_params_t *params);

src/uct/ib/mlx5/rc/rc_mlx5_common.c

@@ -306,7 +306,7 @@ ucs_status_t uct_rc_mlx5_devx_create_cmd_qp(uct_rc_mlx5_iface_common_t *iface)
     status            = uct_rc_mlx5_iface_common_devx_connect_qp(
             iface, &iface->tm.cmd_wq.super.super,
             iface->tm.cmd_wq.super.super.qp_num, &ah_attr,
-            iface->super.super.config.path_mtu, 0, 1);
+            iface->super.super.config.path_mtu, 0, 1, NULL);
     if (status != UCS_OK) {
         goto err_destroy_qp;
     }

src/uct/ib/mlx5/rc/rc_mlx5_common.h

@@ -719,12 +719,14 @@ void uct_rc_mlx5_common_fill_dv_qp_attr(uct_rc_mlx5_iface_common_t *iface,
 ucs_status_t uct_rc_mlx5_iface_common_devx_connect_qp(
         uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_qp_t *qp,
         uint32_t dest_qp_num, struct ibv_ah_attr *ah_attr,
-        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic);
+        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic,
+        const uct_ep_connect_to_ep_params_t *params);
 #else
 static UCS_F_MAYBE_UNUSED ucs_status_t uct_rc_mlx5_iface_common_devx_connect_qp(
         uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_qp_t *qp,
         uint32_t dest_qp_num, struct ibv_ah_attr *ah_attr,
-        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic)
+        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic,
+        const uct_ep_connect_to_ep_params_t *params)
 {
     return UCS_ERR_UNSUPPORTED;
 }

src/uct/ib/mlx5/rc/rc_mlx5_devx.c

@@ -384,7 +384,8 @@ void uct_rc_mlx5_devx_cleanup_srq(uct_ib_mlx5_md_t *md, uct_ib_mlx5_srq_t *srq)
 ucs_status_t uct_rc_mlx5_iface_common_devx_connect_qp(
         uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_qp_t *qp,
         uint32_t dest_qp_num, struct ibv_ah_attr *ah_attr,
-        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic)
+        enum ibv_mtu path_mtu, uint8_t path_index, unsigned max_rd_atomic,
+        const uct_ep_connect_to_ep_params_t *params)
 {
     uct_ib_mlx5_md_t *md = uct_ib_mlx5_iface_md(&iface->super.super);
     char in_2rtr[UCT_IB_MLX5DV_ST_SZ_BYTES(init2rtr_qp_in)]   = {};
@@ -433,8 +434,15 @@ ucs_status_t uct_rc_mlx5_iface_common_devx_connect_qp(
             ucs_assert(ah_attr->dlid >= UCT_IB_ROCE_UDP_SRC_PORT_BASE);
             UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.udp_sport,
                               ah_attr->dlid);
-            UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.dscp,
-                              uct_ib_iface_roce_dscp(&iface->super.super));
+
+            if (params && (params->field_mask &
+                           UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_TRAFFIC_CLASS)) {
+                UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.dscp,
+                                  params->ep_traffic_class);
+            } else {
+                UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.dscp,
+                                  uct_ib_iface_roce_dscp(&iface->super.super));
+            }
         }
 
         uct_ib_mlx5_devx_set_qpc_port_affinity(md, path_index, qpc,
@@ -456,8 +464,14 @@ ucs_status_t uct_rc_mlx5_iface_common_devx_connect_qp(
                    &ah_attr->grh.dgid,
                    UCT_IB_MLX5DV_FLD_SZ_BYTES(qpc, primary_address_path.rgid_rip));
             /* TODO add flow_label support */
-            UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.tclass,
-                              iface->super.super.config.traffic_class);
+            if (params && (params->field_mask &
+                           UCT_EP_CONNECT_TO_EP_PARAM_FIELD_EP_TRAFFIC_CLASS)) {
+                UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.tclass,
+                                  params->ep_traffic_class);
+            } else {
+                UCT_IB_MLX5DV_SET(qpc, qpc, primary_address_path.tclass,
+                                  iface->super.super.config.traffic_class);
+            }
         }
     }
 

src/uct/ib/mlx5/rc/rc_mlx5_ep.c

@@ -746,15 +746,16 @@ ucs_status_t
 uct_rc_mlx5_ep_connect_qp(uct_rc_mlx5_iface_common_t *iface,
                           uct_ib_mlx5_qp_t *qp, uint32_t qp_num,
                           struct ibv_ah_attr *ah_attr, enum ibv_mtu path_mtu,
-                          uint8_t path_index)
+                          uint8_t path_index,
+                          const uct_ep_connect_to_ep_params_t *params)
 {
     uct_ib_mlx5_md_t *md = ucs_derived_of(iface->super.super.super.md, uct_ib_mlx5_md_t);
 
     ucs_assert(path_mtu != UCT_IB_ADDRESS_INVALID_PATH_MTU);
     if (md->flags & UCT_IB_MLX5_MD_FLAG_DEVX) {
         return uct_rc_mlx5_iface_common_devx_connect_qp(
                 iface, qp, qp_num, ah_attr, path_mtu, path_index,
-                iface->super.config.max_rd_atomic);
+                iface->super.config.max_rd_atomic, params);
     } else {
         return uct_rc_iface_qp_connect(&iface->super, qp->verbs.qp, qp_num,
                                        ah_attr, path_mtu);
@@ -820,7 +821,7 @@ uct_rc_mlx5_ep_connect_to_ep_v2(uct_ep_h tl_ep,
          * should be posted to the send side of the QP which is owned by device. */
         status = uct_rc_mlx5_ep_connect_qp(
                 iface, &ep->tm_qp, uct_ib_unpack_uint24(rc_addr->super.qp_num),
-                &ah_attr, path_mtu, ep->super.super.path_index);
+                &ah_attr, path_mtu, ep->super.super.path_index, params);
         if (status != UCS_OK) {
             return status;
         }
@@ -834,7 +835,7 @@ uct_rc_mlx5_ep_connect_to_ep_v2(uct_ep_h tl_ep,
 
     status = uct_rc_mlx5_ep_connect_qp(iface, &ep->super.tx.wq.super, qp_num,
                                        &ah_attr, path_mtu,
-                                       ep->super.super.path_index);
+                                       ep->super.super.path_index, params);
     if (status != UCS_OK) {
         return status;
     }

test/gtest/common/test_obj_size.cc

@@ -49,7 +49,7 @@ UCS_TEST_F(test_obj_size, size) {
 #elif UCS_ENABLE_ASSERT
     UCS_TEST_SKIP_R("Assert enabled");
 #else
-    EXPECTED_SIZE(ucp_ep_t, 64);
+    EXPECTED_SIZE(ucp_ep_t, 72);
     EXPECTED_SIZE(ucp_ep_ext_t, 216);
 #if ENABLE_PARAMS_CHECK
     EXPECTED_SIZE(ucp_rkey_t, 24 + sizeof(ucp_ep_h));

test/gtest/uct/ib/test_ib_event.cc

@@ -352,8 +352,8 @@ class uct_qp_test_event : public uct_test_event_base {
             }
 
             struct ibv_ah_attr ah = ah_attr();
-            status = uct_rc_mlx5_ep_connect_qp(m_iface, &m_txwq.super,
-                                               qp_num(), &ah, path_mtu(), 0);
+            status = uct_rc_mlx5_ep_connect_qp(m_iface, &m_txwq.super, qp_num(),
+                                               &ah, path_mtu(), 0, NULL);
             ASSERT_UCS_OK(status);
         }