Support cuda pin memory in vineyard.

Signed-off-by: Ye Cao <caoye.cao@alibaba-inc.com>
dashanji · Nov 19, 2024 · ba07da0 · ba07da0
1 parent 0f78867
commit ba07da0
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/server/memory/malloc.cc b/src/server/memory/malloc.cc
@@ -40,6 +40,10 @@
 
 #include "common/util/logging.h"  // IWYU pragma: keep
 
+#if defined(ENABLE_CUDA)
+#include <cuda_runtime.h>
+#endif
+
 namespace vineyard {
 
 namespace memory {
@@ -60,6 +64,8 @@ constexpr int64_t kMmapRegionsGap = sizeof(size_t);
 DEFINE_bool(reserve_memory, false,
             "Reserving enough physical memory pages for vineyardd");
 
+DEFINE_bool(cuda_pin_memory, false, "Pin memory for CUDA");
+
 std::unordered_map<void*, MmapRecord> mmap_records;
 
 static void* pointer_advance(void* p, ptrdiff_t n) {
@@ -266,6 +272,16 @@ void* mmap_buffer(int fd, int64_t size, bool gap, bool* is_committed,
     return pointer;
   }
 
+  if (FLAGS_cuda_pin_memory) {
+#ifdef ENABLE_CUDA
+    cudaError_t status = cudaHostRegister(pointer, size, cudaHostRegisterDefault);
+    if (status != cudaSuccess) {
+      LOG(ERROR) << "cudaHostRegister failed with error: "
+                 << cudaGetErrorString(status);
+    }
+#endif
+  }
+
   MmapRecord& record = mmap_records[pointer];
   record.fd = fd;
   record.size = size;