From 7a708d68435dfb393eb569267c06a938a25e9ecd Mon Sep 17 00:00:00 2001
From: Yimin Jiang <jiangyimin@bytedance.com>
Date: Wed, 15 Jan 2020 20:00:40 +0800
Subject: [PATCH] rdma: only reg_mr for vals (#18)

---
 src/rdma_transport.h |  8 ++++----
 src/rdma_van.h       | 24 +++++++++++++-----------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/rdma_transport.h b/src/rdma_transport.h
index fec1c70d3..9135d2576 100644
--- a/src/rdma_transport.h
+++ b/src/rdma_transport.h
@@ -305,7 +305,7 @@ class RDMATransport : public Transport {
   }
 
   void SendPushRequest(Message &msg, MessageBuffer *msg_buf, RemoteTuple remote_tuple) {
-    CHECK_EQ(msg_buf->mrs.size(), 3);
+    CHECK_EQ(msg_buf->mrs.size(), 1);
     auto raddr = std::get<0>(remote_tuple);
     auto rkey = std::get<1>(remote_tuple);
     auto idx = std::get<2>(remote_tuple);
@@ -313,9 +313,9 @@ class RDMATransport : public Transport {
     // push request, split the meta and data into two writes
     // further, it does not send keys and lens since these meta already carries these info 
     struct ibv_sge my_sge;
-    my_sge.addr = reinterpret_cast<uint64_t>(msg_buf->mrs[1].first->addr);
-    my_sge.length = msg_buf->mrs[1].second;
-    my_sge.lkey = msg_buf->mrs[1].first->lkey;
+    my_sge.addr = reinterpret_cast<uint64_t>(msg_buf->mrs[0].first->addr);
+    my_sge.length = msg_buf->mrs[0].second;
+    my_sge.lkey = msg_buf->mrs[0].first->lkey;
 
     // this rdma-write will not trigger any signal both remotely and locally
     struct ibv_send_wr wr, *bad_wr = nullptr;
diff --git a/src/rdma_van.h b/src/rdma_van.h
index 945b205ff..684ed6d3c 100755
--- a/src/rdma_van.h
+++ b/src/rdma_van.h
@@ -366,7 +366,7 @@ class RDMAVan : public Van {
       LOG(INFO) << "Send Push Request: key=" << msg.meta.key
           << "\t timestamp=" << msg.meta.timestamp 
           << "\t recver=" << msg.meta.recver
-          << "\t tensor_len=" << msg_buf->mrs[1].second
+          << "\t tensor_len=" << msg_buf->mrs[0].second
           << "\t remote_idx=" << std::get<2>(remote_tuple)
           << "\t remote_addr=" << std::get<0>(remote_tuple) 
           << std::flush;
@@ -532,11 +532,12 @@ class RDMAVan : public Van {
     return msg_buf;
   }
 
-  void RegisterMemory(Message &msg) {
+  void RegisterMemory(Message &msg) { 
+    size_t sa_cnt = 0;
     for (auto& sa : msg.data) {
       if (sa.size() == 0) continue;
       std::lock_guard<std::mutex> lock(map_mu_);
-      if (mem_mr_.find(sa.data()) == mem_mr_.end()) {
+      if ((mem_mr_.find(sa.data()) == mem_mr_.end()) && (sa_cnt==1)) { // only vals register memory
         struct ibv_mr *temp_mr;
         CHECK (temp_mr = ibv_reg_mr(mem_allocator_->GetPD(), sa.data(), sa.size(),
             IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE))
@@ -544,19 +545,20 @@ class RDMAVan : public Van {
             << ", sa.size()=" << sa.size();
         mem_mr_[sa.data()] = temp_mr;
       }
+      ++sa_cnt;
     }
   }
 
   void PrepareData(Message &msg, MessageBuffer *msg_buf) {
     if (!(msg.meta.push && msg.meta.request)) return; // only push request
-    for (auto &sa : msg_buf->data) {
-      if (sa.size() == 0) continue;
-      std::lock_guard<std::mutex> lock(map_mu_);
-      auto it = mem_mr_.find(sa.data());
-      MRPtr ptr(it->second, [](struct ibv_mr *mr) {});
-      CHECK(ptr.get()) << strerror(errno);
-      msg_buf->mrs.push_back(std::make_pair(std::move(ptr), sa.size()));
-    }
+    auto &sa = msg_buf->data[1];
+    if (sa.size() == 0) return;
+    std::lock_guard<std::mutex> lock(map_mu_);
+    auto it = mem_mr_.find(sa.data());
+    CHECK_NE(it, mem_mr_.end());
+    MRPtr ptr(it->second, [](struct ibv_mr *mr) {});
+    CHECK(ptr.get()) << strerror(errno);
+    msg_buf->mrs.push_back(std::make_pair(std::move(ptr), sa.size()));
   }
 
   void AddMeta(Message &msg) {