Expose the check for RDMA hugepages support

telescope-proj · Oct 27, 2023 · a964931 · a964931
1 parent 996e43c
commit a964931
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 23 deletions.
diff --git a/tcm/include/tcm_mm.h b/tcm/include/tcm_mm.h
@@ -22,6 +22,16 @@
 #define TCM_DEFAULT_HUGEPAGE_SIZE 2097152
 #endif
 
+/* Check if the RDMA subsystem can use hugepages */
+
+static inline int tcm_can_use_hugepages() {
+    char * e = getenv("RDMAV_HUGEPAGES_SAFE");
+    if (e) {
+        return atoi(e);
+    }
+    return 0;
+}
+
 /* RAII style heap memory class */
 
 class tcm_managed_mem {

diff --git a/tcm/src/tcm_mm.cpp b/tcm/src/tcm_mm.cpp
@@ -71,33 +71,24 @@ void * tcm_mem_align_rdma(size_t size) {
     return NULL;
 #else
     int    ps  = 0;
-    int    hp  = 0;
     int    adv = 0;
-    char * e   = getenv("TCM_PAGE_SIZE");
-    if (e) {
-        ps = atoi(e);
-        if (!ps)
-            ps = tcm_get_page_size();
-        else
-            hp = 1;
-    } else {
-        if (TCM_DEFAULT_HUGEPAGE_SIZE) {
-            ps = TCM_DEFAULT_HUGEPAGE_SIZE;
-            hp = 1;
-        } else {
-            ps = tcm_get_page_size();
-        }
-    }
-    if (hp) {
-        e = getenv("RDMAV_HUGEPAGES_SAFE");
+    if (tcm_can_use_hugepages()) {
+        char * e   = getenv("TCM_PAGE_SIZE");
         if (e) {
-            if (atoi(e) != 1)
-                adv = MADV_NOHUGEPAGE;
-            else
-                adv = MADV_HUGEPAGE;
+            ps = atoi(e);
+            if (!ps)
+                ps = tcm_get_page_size();
         } else {
-            adv = MADV_NOHUGEPAGE;
+            if (TCM_DEFAULT_HUGEPAGE_SIZE) {
+                ps = TCM_DEFAULT_HUGEPAGE_SIZE;
+            } else {
+                ps = tcm_get_page_size();
+            }
         }
+        adv = MADV_HUGEPAGE;
+    } else {
+        ps = tcm_get_page_size();
+        adv = MADV_NOHUGEPAGE;
     }
     void * memptr = tcm_mem_align(size, ps);
     if (!memptr)