feat: add PCIe Relaxed Ordering (RO) support and RDMA traffic class (…

mooncake-transfer-engine/include/config.h

@@ -57,6 +57,8 @@ struct GlobalConfig {
     size_t fragment_limit = 16384;
     bool enable_dest_device_affinity = false;
     EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE;
+    int traffic_class = -1;  // -1 means not set, will be read from MC_IB_TC
+    int ib_pci_relaxed_ordering_mode = 2;  // 0: off, 1: on if supported, 2: auto
 };
 
 void loadGlobalConfig(GlobalConfig &config);

mooncake-transfer-engine/src/config.cpp

@@ -293,6 +293,32 @@ void loadGlobalConfig(GlobalConfig &config) {
                             "MC_ENDPOINT_STORE_TYPE, it should be FIFO|SIEVE";
         }
     }
+
+    const char *traffic_class_env = std::getenv("MC_IB_TC");
+    if (traffic_class_env) {
+        int val = atoi(traffic_class_env);
+        if (val >= 0 && val <= 255)
+            config.traffic_class = val;
+        else
+            LOG(WARNING) << "Ignore value from environment variable "
+                            "MC_IB_TC, it should be 0-255";
+    }
+
+    const char *ib_relaxed_ordering_env =
+        std::getenv("MC_IB_PCI_RELAXED_ORDERING");
+    if (ib_relaxed_ordering_env) {
+        int val = atoi(ib_relaxed_ordering_env);
+        if (val >= 0 && val <= 2)
+            config.ib_pci_relaxed_ordering_mode = val;
+        else
+            LOG(WARNING)
+                << "Ignore value from environment variable "
+                   "MC_IB_PCI_RELAXED_ORDERING, it should be 0|1|2";
+    }
+    if (traffic_class_env) {
+        LOG(INFO) << "[Config] traffic_class set to " << config.traffic_class
+                  << " via MC_IB_TC";
+    }
 }
 
 std::string mtuLengthToString(ibv_mtu mtu) {

mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp

@@ -405,6 +405,9 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
     // TODO gidIndex and portNum must fetch from REMOTE
     attr.ah_attr.grh.sgid_index = context_.gidIndex();
     attr.ah_attr.grh.hop_limit = MAX_HOP_LIMIT;
+    if (globalConfig().traffic_class >= 0) {
+        attr.ah_attr.grh.traffic_class = globalConfig().traffic_class;
+    }
     attr.ah_attr.dlid = peer_lid;
     attr.ah_attr.sl = 0;
     attr.ah_attr.src_path_bits = 0;

mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp

@@ -24,6 +24,8 @@
 #include <set>
 #include <thread>
 
+#include <dlfcn.h>
+
 #include "common.h"
 #include "config.h"
 #include "memory_location.h"
@@ -32,7 +34,56 @@
 #include "transport/rdma_transport/rdma_endpoint.h"
 
 namespace mooncake {
-RdmaTransport::RdmaTransport() {}
+
+static bool MCIbRelaxedOrderingEnabled = false;
+static int MCIbRelaxedOrderingMode = 2;
+
+// Mode definition for MC_IB_PCI_RELAXED_ORDERING env.
+// 0 - disabled, 1 - enabled if supported, 2 - auto (default, same as 1 today).
+static int getIbRelaxedOrderingMode() {
+    int val = globalConfig().ib_pci_relaxed_ordering_mode;
+    if (val < 0 || val > 2) {
+        return 2;
+    }
+    return val;
+}
+
+// Determine whether RELAXED_ORDERING is enabled and possible
+// This function checks for ibv_reg_mr_iova2 symbol which is available
+// in IBVERBS_1.8 and above. The feature is only supported in IBVERBS_1.8+.
+bool has_ibv_reg_mr_iova2(void) {
+    void *handle = dlopen("libibverbs.so", RTLD_NOW);
+    if (!handle) {
+        handle = dlopen("libibverbs.so.1", RTLD_NOW);
+        if (!handle) return false;
+    }
+
+    void *sym = dlsym(handle, "ibv_reg_mr_iova2");
+    dlclose(handle);
+    return sym != NULL;
+}
+
+RdmaTransport::RdmaTransport() {
+    MCIbRelaxedOrderingMode = getIbRelaxedOrderingMode();
+    if (MCIbRelaxedOrderingMode == 0) {
+        LOG(INFO) << "[RDMA] Relaxed ordering disabled via "
+                  << "MC_IB_PCI_RELAXED_ORDERING=0. "
+                  << "Falling back to strict ordering.";
+        MCIbRelaxedOrderingEnabled = false;
+        return;
+    }
+
+    MCIbRelaxedOrderingEnabled = has_ibv_reg_mr_iova2();
+    if (MCIbRelaxedOrderingEnabled) {
+        LOG(INFO) << "[RDMA] Relaxed ordering is supported on this host; "
+                     "IBV_ACCESS_RELAXED_ORDERING will be requested for "
+                     "registered memory regions.";
+    } else {
+        LOG(INFO) << "[RDMA] Relaxed ordering is NOT supported ("
+                  << "ibv_reg_mr_iova2 missing or unavailable). "
+                  << "Falling back to strict ordering.";
+    }
+}
 
 RdmaTransport::~RdmaTransport() {
 #ifdef CONFIG_USE_BATCH_DESC_SET
@@ -132,10 +183,17 @@ int RdmaTransport::registerLocalMemory(void *addr, size_t length,
                                        bool update_metadata) {
     (void)remote_accessible;
     BufferDesc buffer_desc;
-    const static int access_rights = IBV_ACCESS_LOCAL_WRITE |
-                                     IBV_ACCESS_REMOTE_WRITE |
-                                     IBV_ACCESS_REMOTE_READ;
-
+    const int kBaseAccessRights = IBV_ACCESS_LOCAL_WRITE |
+                                  IBV_ACCESS_REMOTE_WRITE |
+                                  IBV_ACCESS_REMOTE_READ;
+
+    static int access_rights = 0;
+    if (access_rights == 0) {
+        access_rights = kBaseAccessRights;
+        if (MCIbRelaxedOrderingEnabled) {
+            access_rights |= IBV_ACCESS_RELAXED_ORDERING;
+        }
+    }
     bool do_pre_touch = context_list_.size() > 0 &&
                         std::thread::hardware_concurrency() >= 4 &&
                         length >= (size_t)4 * 1024 * 1024 * 1024;