From e64503ab4318a4a8271bb7da1457d9a091d9e1d9 Mon Sep 17 00:00:00 2001
From: Nick Gasson <nick@nickg.me.uk>
Date: Sat, 5 Nov 2022 14:49:06 +0000
Subject: [PATCH] Use transparent huge pages on Linux where possible

---
 src/rt/model.c  |  4 ++--
 src/rt/mspace.c |  6 +++++-
 src/util.c      | 19 +++++++++++++++++++
 src/util.h      |  1 +
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/rt/model.c b/src/rt/model.c
index c3ef6e26..881617b8 100644
--- a/src/rt/model.c
+++ b/src/rt/model.c
@@ -46,7 +46,7 @@ typedef struct _callback {
 } callback_t;
 
 #define MEMBLOCK_LINE_SZ 64
-#define MEMBLOCK_PAGE_SZ 0x200000
+#define MEMBLOCK_PAGE_SZ 0x800000
 
 typedef struct _memblock {
    memblock_t *chain;
@@ -240,7 +240,7 @@ static void *static_alloc(rt_model_t *m, size_t size)
       mb->pagesz = MAX(MEMBLOCK_PAGE_SZ, nlines * MEMBLOCK_LINE_SZ);
       mb->chain  = m->memblocks;
       mb->free   = mb->pagesz / MEMBLOCK_LINE_SZ;
-      mb->ptr    = nvc_memalign(MEMBLOCK_LINE_SZ, mb->pagesz);
+      mb->ptr    = map_huge_pages(MEMBLOCK_LINE_SZ, mb->pagesz);
 
       m->memblocks = mb;
    }
diff --git a/src/rt/mspace.c b/src/rt/mspace.c
index 5f88b8e4..9c46a4b7 100644
--- a/src/rt/mspace.c
+++ b/src/rt/mspace.c
@@ -90,7 +90,7 @@ mspace_t *mspace_new(size_t size)
 
    DEBUG_ONLY(m->stress = opt_get_int(OPT_GC_STRESS));
 
-   m->space = nvc_memalign(LINE_SIZE, m->maxsize);
+   m->space = map_huge_pages(LINE_SIZE, m->maxsize);
 
    MSPACE_POISON(m->space, m->maxsize);
 
@@ -191,6 +191,10 @@ void *mspace_alloc(mspace_t *m, size_t size)
                (*it)->ptr += asize;
             }
 
+            // Make sure the first fault to the page is a write to
+            // allocate THP on Linux
+            *(volatile char *)base = 0;
+
             return base;
          }
       }
diff --git a/src/util.c b/src/util.c
index a14e6171..7d1bf73d 100644
--- a/src/util.c
+++ b/src/util.c
@@ -105,6 +105,8 @@
 #define MAX_FMT_BUFS    32
 #define MAX_PRINTF_BUFS 8
 
+#define HUGE_PAGE_SIZE  0x200000
+
 typedef void (*print_fn_t)(const char *fmt, ...);
 
 static char *ansi_vasprintf(const char *fmt, va_list ap, bool force_plain);
@@ -1180,6 +1182,23 @@ void *mmap_guarded(size_t sz, guard_fault_fn_t fn, void *ctx)
    return ptr;
 }
 
+void *map_huge_pages(size_t align, size_t sz)
+{
+#ifdef __linux__
+   if (sz >= HUGE_PAGE_SIZE) {
+      const size_t mapsz = ALIGN_UP(sz, HUGE_PAGE_SIZE);
+      void *mem = nvc_memalign(MAX(HUGE_PAGE_SIZE, align), mapsz);
+
+      if (madvise(mem, mapsz, MADV_HUGEPAGE) < 0)
+         warnf("madvise: MADV_HUGEPAGE: %s", last_os_error());
+
+      return mem;
+   }
+#endif
+
+   return nvc_memalign(align, sz);
+}
+
 int checked_sprintf(char *buf, int len, const char *fmt, ...)
 {
    assert(len > 0);
diff --git a/src/util.h b/src/util.h
index 96c345a3..91d882c6 100644
--- a/src/util.h
+++ b/src/util.h
@@ -201,6 +201,7 @@ void *nvc_memalign(size_t align, size_t sz);
 void nvc_munmap(void *ptr, size_t length);
 void nvc_memprotect(void *ptr, size_t length, mem_access_t prot);
 void *mmap_guarded(size_t sz, guard_fault_fn_t fn, void *ctx);
+void *map_huge_pages(size_t align, size_t sz);
 
 void run_program(const char *const *args);
 
-- 
2.39.2