rdma: use page aligned memory (#12)

* use page aligned memory * fix page aligned memory * add check * nit
bytedance · Nov 28, 2019 · b1b5efb · b1b5efb
1 parent 7672fdc
commit b1b5efb
Showing 1 changed file with 27 additions and 6 deletions.
diff --git a/src/rdma_van.h b/src/rdma_van.h
@@ -34,6 +34,9 @@
 
 namespace ps {
 
+#define DIVUP(x, y) (((x)+(y)-1)/(y))
+#define ROUNDUP(x, y) (DIVUP((x), (y))*(y))
+
 static const int kStartDepth = 128;
 static const int kWriteDepth = kStartDepth;
 
@@ -52,6 +55,7 @@ static const size_t kAlignment = 8;
 static const int kMaxResolveRetry = 50000;
 static const int kBasePort = 9010;
 
+
 template <typename T>
 static inline T align_floor(T v, T align) {
   return v - (v % align);
@@ -62,6 +66,17 @@ static inline T align_ceil(T v, T align) {
   return align_floor(v + align - 1, align);
 }
 
+static inline void ib_malloc(void** ptr, size_t size) {
+  size_t page_size = sysconf(_SC_PAGESIZE);
+  void* p;
+  int size_aligned = ROUNDUP(size, page_size);
+  int ret = posix_memalign(&p, page_size, size_aligned);
+  CHECK_EQ(ret, 0) << "posix_memalign error: " << strerror(ret);
+  CHECK(p);
+  memset(p, 0, size);
+  *ptr = p;
+}
+
 class SimpleMempool {
  public:
   explicit SimpleMempool(struct ibv_pd *pd, size_t size = 0x10000000) {
@@ -78,7 +93,8 @@ class SimpleMempool {
                << ", mempool num set to " << mempool_num;
 
     for (size_t i = 0; i < mempool_num; ++i) {
-      char *p = reinterpret_cast<char *>(aligned_alloc(kAlignment, size));
+      char *p;
+      ib_malloc((void**) &p, size);
       total_allocated_size += size;
       CHECK(p);
       CHECK(mr = ibv_reg_mr(pd, p, size,
@@ -112,7 +128,8 @@ class SimpleMempool {
       while (proper_size > new_mem_size) {
         new_mem_size *= 2;
       }
-      char *p = reinterpret_cast<char *>(aligned_alloc(kAlignment, new_mem_size));
+      char *p;
+      ib_malloc((void**) &p, new_mem_size);
       CHECK(p);
       struct ibv_mr *mr;
       CHECK(mr = ibv_reg_mr(pd_, p, new_mem_size, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE));
@@ -385,7 +402,8 @@ struct Endpoint {
                              ThreadsafeQueue<WRContext *> *queue, size_t num,
                              WRContextType type) {
     for (size_t i = 0; i < num; ++i) {
-      void *buf = aligned_alloc(kAlignment, kMempoolChunkSize);
+      void *buf;
+      ib_malloc((void**) &buf, kMempoolChunkSize);
       CHECK(buf);
       struct ibv_mr *mr = ibv_reg_mr(pd, buf, kMempoolChunkSize, 0);
       CHECK(mr);
@@ -420,7 +438,8 @@ struct Endpoint {
                           kWriteContext);
 
     for (size_t i = 0; i < kRxDepth; ++i) {
-      void *buf = aligned_alloc(kAlignment, kMempoolChunkSize);
+      void *buf;
+      ib_malloc((void**) &buf, kMempoolChunkSize);
       CHECK(buf);
       struct ibv_mr *mr =
           ibv_reg_mr(pd, buf, kMempoolChunkSize, IBV_ACCESS_LOCAL_WRITE);
@@ -455,7 +474,9 @@ struct Endpoint {
 
 class RDMAVan : public Van {
  public:
-  RDMAVan() {}
+  RDMAVan() {
+    CHECK_EQ(ibv_fork_init(), 0) << strerror(errno);
+  }
   ~RDMAVan() {}
 
  protected:
@@ -1392,4 +1413,4 @@ class RDMAVan : public Van {
 };  // namespace ps
 
 #endif  // DMLC_USE_RDMA
-#endif  // PS_RDMA_VAN_H_
+#endif  // PS_RDMA_VAN_H_