feat(flex): Implemented basic support for hugepages. (#3489)

alibaba · Jan 17, 2024 · cf654e1 · cf654e1
1 parent b8974be
commit cf654e1
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 7 deletions.
diff --git a/flex/CMakeLists.txt b/flex/CMakeLists.txt
@@ -15,6 +15,7 @@ option(BUILD_TEST "Whether to build test" ON)
 option(BUILD_DOC "Whether to build doc" ON)
 option(BUILD_ODPS_FRAGMENT_LOADER "Whether to build odps fragment loader" ON)
 option(MONITOR_SESSIONS "Whether monitor sessions" OFF)
+option(ENABLE_HUGEPAGE "Whether to use hugepages when open mmap array in memory" OFF)
 
 # ------------------------------------------------------------------------------
 # cmake configs
@@ -35,6 +36,12 @@ if (MONITOR_SESSIONS)
     add_definitions(-DMONITOR_SESSIONS)
 endif ()
 
+
+if (ENABLE_HUGEPAGE)
+    message("Hugepage is enabled")
+    add_definitions(-DHUGEPAGE)
+endif ()
+
 execute_process(COMMAND uname -r OUTPUT_VARIABLE LINUX_KERNEL_VERSION)
 string(STRIP ${LINUX_KERNEL_VERSION} LINUX_KERNEL_VERSION)
 message(${LINUX_KERNEL_VERSION})

diff --git a/flex/utils/allocators.h b/flex/utils/allocators.h
@@ -41,9 +41,6 @@ class ArenaAllocator {
   {
   }
   ~ArenaAllocator() {
-    for (auto ptr : malloc_buffers_) {
-      free(ptr);
-    }
     for (auto ptr : mmap_buffers_) {
       delete ptr;
     }
@@ -85,9 +82,11 @@ class ArenaAllocator {
  private:
   void* allocate_batch(size_t size) {
     if (prefix_.empty()) {
-      void* ret = malloc(size);
-      malloc_buffers_.push_back(ret);
-      return ret;
+      mmap_array<char>* buf = new mmap_array<char>();
+      buf->open_in_memory("");
+      buf->resize(size);
+      mmap_buffers_.push_back(buf);
+      return static_cast<void*>(buf->data());
     } else {
       mmap_array<char>* buf = new mmap_array<char>();
       buf->open(prefix_ + std::to_string(mmap_buffers_.size()), false);
@@ -99,7 +98,6 @@ class ArenaAllocator {
 
   std::string prefix_;
   std::vector<mmap_array<char>*> mmap_buffers_;
-  std::vector<void*> malloc_buffers_;
 
   void* cur_buffer_;
   size_t cur_loc_;

diff --git a/flex/utils/mmap_array.h b/flex/utils/mmap_array.h
@@ -119,8 +119,25 @@ class mmap_array {
     madvise(data_, size_ * sizeof(T), MADV_RANDOM | MADV_WILLNEED);
   }
 
+#ifdef HUGEPAGE
+#ifdef __ia64__
+#define ADDR (void*) (0x8000000000000000UL)
+#define FLAGS (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB | MAP_FIXED)
+#else
+#define ADDR (void*) (0x0UL)
+#define FLAGS (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB)
+#endif
+
+#define PROTECTION (PROT_READ | PROT_WRITE)
+
+#define HUGEPAGE_SIZE (2UL * 1024 * 1024)
+#define HUGEPAGE_MASK (2UL * 1024 * 1024 - 1UL)
+#define ROUND_UP(size) (((size) + HUGEPAGE_MASK) & (~HUGEPAGE_MASK))
+#endif
+
   void open_in_memory(const std::string& filename) {
     reset();
+#ifndef HUGEPAGE
     filename_ = filename;
     if (!filename_.empty() && std::filesystem::exists(filename_)) {
       size_t file_size = std::filesystem::file_size(filename_);
@@ -137,6 +154,24 @@ class mmap_array {
         }
       }
     }
+#else
+    if (!filename.empty() && std::filesystem::exists(filename)) {
+      LOG(INFO) << "open in memory with hugepage...";
+      size_t file_size = std::filesystem::file_size(filename);
+      size_ = file_size / sizeof(T);
+      if (size_ != 0) {
+        size_t size_in_bytes = size_ * sizeof(T);
+        data_ = static_cast<T*>(
+            mmap(ADDR, ROUND_UP(size_in_bytes), PROTECTION, FLAGS, -1, 0));
+        if (data_ == MAP_FAILED) {
+          LOG(FATAL) << "mmap failed " << filename << " " << strerror(errno)
+                     << "..\n";
+        }
+        FILE* fin = fopen(filename.c_str(), "rb");
+        CHECK_EQ(fread(data_, size_in_bytes, 1, fin), 1);
+      }
+    }
+#endif
   }
 
   void dump(const std::string& filename) {
@@ -170,6 +205,7 @@ class mmap_array {
       if (size == 0) {
         reset();
       } else {
+#ifndef HUGEPAGE
         T* new_data = static_cast<T*>(
             mmap(NULL, size * sizeof(T), PROT_READ | PROT_WRITE,
                  MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0));
@@ -190,6 +226,24 @@ class mmap_array {
           close(fd_);
           fd_ = -1;
         }
+#else
+        LOG(INFO) << "resize with hugepage...";
+        T* new_data = static_cast<T*>(
+            mmap(ADDR, ROUND_UP(size * sizeof(T)), PROTECTION, FLAGS, -1, 0));
+        if (new_data == MAP_FAILED) {
+          LOG(FATAL) << "mmap failed " << strerror(errno) << "..\n";
+        }
+
+        if (data_ != NULL) {
+          size_t copy_size = std::min(size, size_);
+          if (copy_size > 0) {
+            memcpy((void*)new_data, (void*) data_, copy_size * sizeof(T));
+          }
+          munmap(data_, ROUND_UP(size_ * sizeof(T)));
+        }
+        data_ = new_data;
+        size_ = size;
+#endif
       }
     } else {
       if (read_only_) {
@@ -235,6 +289,12 @@ class mmap_array {
 
   bool read_only() const { return read_only_; }
 
+#undef ADDR
+#undef FLAGS
+#undef HUGEPAGE_SIZE
+#undef HUGEPAGE_MASK
+#undef ROUND_UP
+
   void touch(const std::string& filename) {
     if (read_only_) {
       if (filename_ != "") {

diff --git a/flex/utils/property/column.h b/flex/utils/property/column.h
@@ -330,7 +330,18 @@ class TypedColumn<std::string_view> : public ColumnBase {
     } else {
       basic_size_ = basic_buffer_.size();
       extra_size_ = size - basic_size_;
+#ifdef HUGEPAGE
+      if (basic_buffer_.size() != 0) {
+        size_t basic_avg_width =
+            (basic_buffer_.data_size() + basic_buffer_.size() - 1) /
+            basic_buffer_.size();
+        extra_buffer_.resize(extra_size_, extra_size_ * basic_avg_width);
+      } else {
+        extra_buffer_.resize(extra_size_, extra_size_ * width_);
+      }
+#else
       extra_buffer_.resize(extra_size_, extra_size_ * width_);
+#endif
     }
   }