From 945982479255a95e7b8fbfd54c1a832513a84a50 Mon Sep 17 00:00:00 2001
From: Qingqing Li <liqingqing3@huawei.com>
Date: Tue, 10 Sep 2024 15:15:23 +0800
Subject: [PATCH] Use THP for dynamic shared librarys

When add env "LD_HUGEPAGE_LIB=2" for a exec file,
the load will use THP for it's shared librarys.

---
 elf/dl-load.c                  |  7 ++--
 elf/dl-map-segments-hugepage.h | 62 +++++++++++++++++++++++++---------
 elf/rtld.c                     |  1 +
 sysdeps/generic/ldsodefs.h     |  2 ++
 4 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/elf/dl-load.c b/elf/dl-load.c
index 1c3ccdfe..063548ac 100644
--- a/elf/dl-load.c
+++ b/elf/dl-load.c
@@ -1187,9 +1187,11 @@ _dl_map_object_from_fd (const char *name, const char *origname, int fd,
 
 #ifdef HUGEPAGE_SHARED_LIB
           if (GLRO(dl_hugepage_mask) & DL_HUGEPAGE_LIB_LARGE_IN_FLAG) {
-              hp_bitmap[nloadcmds] = 1;
+              hp_bitmap[nloadcmds] = DL_HUGEPAGE_LARGE_IN;
+          } else if (GLRO(dl_hugepage_mask) & DL_HUGEPAGE_LIB_THP_IN_FLAG) {
+              hp_bitmap[nloadcmds] = DL_HUGEPAGE_THP_IN;
           } else if (ph->p_flags & PF_HUGEPAGE) {
-              hp_bitmap[nloadcmds] = 1;
+              hp_bitmap[nloadcmds] = DL_HUGEPAGE_LARGE_IN;
               use_hugepage = true;
           } else
               hp_bitmap[nloadcmds] = 0;
@@ -1305,6 +1307,7 @@ _dl_map_object_from_fd (const char *name, const char *origname, int fd,
     int hp_errcode = 0;
     char hp_buf[ERRSTRING_BUF_LEN];
     if ((GLRO(dl_hugepage_mask) & DL_HUGEPAGE_LIB_LARGE_IN_FLAG) ||
+	(GLRO(dl_hugepage_mask) & DL_HUGEPAGE_LIB_THP_IN_FLAG) ||
             ((GLRO(dl_hugepage_mask) & DL_HUGEPAGE_PROBE_FLAG) && use_hugepage))
       {
         _llvm_ld_relro_hp_fallback(phdr, relro_ph, l, hp_bitmap);
diff --git a/elf/dl-map-segments-hugepage.h b/elf/dl-map-segments-hugepage.h
index 37993fa7..31b0c6a0 100644
--- a/elf/dl-map-segments-hugepage.h
+++ b/elf/dl-map-segments-hugepage.h
@@ -180,7 +181,7 @@ _zero_tail_page(const struct loadcmd *c, ElfW(Addr) zero, ElfW(Addr) zeropage,
 }
 
 static __always_inline const char *
-_mmap_remain_zero_page(ElfW(Addr) zeropage, ElfW(Addr) zeroend, int prot)
+_mmap_remain_zero_page(ElfW(Addr) zeropage, ElfW(Addr) zeroend, int prot, char hp_flag)
 {
   ElfW(Addr) hp_start = ALIGN_UP(zeropage, SIZE_2MB);
   size_t len = 0, mod = 0;
@@ -213,11 +214,24 @@ _mmap_remain_zero_page(ElfW(Addr) zeropage, ElfW(Addr) zeroend, int prot)
 
   if (__glibc_unlikely (GLRO(dl_debug_mask) & DL_DEBUG_FILES))
     _dl_debug_printf("\t\t\t=> %s mmap anonymous hugepage: [%lx-%lx)\n", __func__, hp_start, hp_start + len - mod);
-  mapat = __mmap((caddr_t) hp_start, len - mod, prot,
-                  MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB|(SHFIT_2MB << MAP_HUGE_SHIFT),
-                  -1, 0);
-  if (__glibc_unlikely (mapat == MAP_FAILED))
-    return DL_MAP_SEGMENTS_ERROR_MAP_ZERO_FILL;
+
+  if (hp_flag == DL_HUGEPAGE_LARGE_IN)
+  {
+    mapat = __mmap((caddr_t) hp_start, len - mod, prot,
+		    MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB|(SHFIT_2MB << MAP_HUGE_SHIFT),
+                    -1, 0);
+
+    if (__glibc_unlikely (mapat == MAP_FAILED))
+      return DL_MAP_SEGMENTS_ERROR_MAP_ZERO_FILL;
+  } else {
+    mapat = __mmap((caddr_t) hp_start, len - mod, prot,
+                    MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+
+    if (__glibc_unlikely (mapat == MAP_FAILED))
+      return DL_MAP_SEGMENTS_ERROR_MAP_ZERO_FILL;
+
+    (void) madvise((void *)hp_start, len - mod, MADV_HUGEPAGE);
+  }
 
   if (mod > 0)
     {
@@ -238,7 +252,7 @@ _mmap_remain_zero_page(ElfW(Addr) zeropage, ElfW(Addr) zeroend, int prot)
  */
 static __always_inline const char *
 _mmap_segment_memsz(struct link_map *l, const struct loadcmd * c,
-                ElfW(Addr) mapstart, size_t extra_len, size_t *memsz_len)
+                ElfW(Addr) mapstart, size_t extra_len, size_t *memsz_len, char hp_flag)
 {
   const char * errstring = NULL;
 
@@ -305,7 +319,7 @@ _mmap_segment_memsz(struct link_map *l, const struct loadcmd * c,
       *memsz_len = ALIGN_UP(zeroend, GLRO(dl_pagesize)) - zeropage;
       if (__glibc_unlikely (GLRO(dl_debug_mask) & DL_DEBUG_FILES))
         _dl_debug_printf("\t\tmmap remain anonymous page [%lx-%lx)\n", zeropage, ALIGN_UP(zeroend, GLRO(dl_pagesize)));
-      errstring = _mmap_remain_zero_page(zeropage, zeroend, c->prot);
+      errstring = _mmap_remain_zero_page(zeropage, zeroend, c->prot, hp_flag);
     }
   return errstring;
 }
@@ -317,7 +331,7 @@ _mmap_segment_memsz(struct link_map *l, const struct loadcmd * c,
  */
 static __always_inline const char *
 _mmap_segment_filesz(struct link_map *l, const struct loadcmd *c, ElfW(Addr) mapstart,
-              size_t extra_len, int fd)
+              size_t extra_len, int fd, char hp_flag)
 {
   void *map_addr = 0;
 
@@ -382,11 +396,24 @@ _mmap_segment_filesz(struct link_map *l, const struct loadcmd *c, ElfW(Addr) map
   size_t mod = len % SIZE_2MB;
   if (__glibc_unlikely (GLRO(dl_debug_mask) & DL_DEBUG_FILES))
     _dl_debug_printf("\t\tmmap hugepage: [%lx-%lx)\n", mapstart, mapstart + len - mod);
-  mapstart = (ElfW(Addr))__mmap((void *)mapstart, len - mod, c->prot,
-                      MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB|(SHFIT_2MB << MAP_HUGE_SHIFT),
-                      -1, 0);
-  if (__glibc_unlikely ((void *)mapstart == MAP_FAILED))
-    return DL_MAP_SEGMENTS_ERROR_MAP_SEGMENT;
+
+  if (hp_flag == DL_HUGEPAGE_LARGE_IN)
+  {
+    mapstart = (ElfW(Addr))__mmap((void *)mapstart, len - mod, c->prot,
+                      		  MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB|(SHFIT_2MB << MAP_HUGE_SHIFT),
+	                          -1, 0);
+
+    if (__glibc_unlikely ((void *)mapstart == MAP_FAILED))
+      return DL_MAP_SEGMENTS_ERROR_MAP_SEGMENT;
+  } else {
+    mapstart = (ElfW(Addr))__mmap((void *)mapstart, len - mod, c->prot,
+                                  MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+
+    if (__glibc_unlikely ((void *)mapstart == MAP_FAILED))
+      return DL_MAP_SEGMENTS_ERROR_MAP_SEGMENT;
+
+    (void) madvise((void *)mapstart, len - mod, MADV_HUGEPAGE);
+  }
 
   if ((c->prot & PROT_WRITE) == 0 && __mprotect((void *)mapstart, len - mod, c->prot | PROT_WRITE) < 0)
     {
@@ -561,11 +588,13 @@ _mmap_segment(struct link_map *l, const struct loadcmd loadcmds[], size_t nloadc
 {
   const char * errstring = NULL;
   size_t extra_len, memsz_len = 0;
+  char hp_flag;
 
   if (__glibc_unlikely (GLRO(dl_debug_mask) & DL_DEBUG_FILES))
     _dl_debug_printf("\t%s mmap from %lx: \n\t{\n", __func__, mapstart);
 
-  if (!hp_bitmap[((void *)c - (void *)loadcmds) / sizeof(struct loadcmd)])
+  hp_flag = hp_bitmap[((void *)c - (void *)loadcmds) / sizeof(struct loadcmd)];
+  if (!hp_flag)
     {
       errstring = __mmap_segment_normalsz(c, mapstart, fd, mapseglen);
       if (__glibc_unlikely (GLRO(dl_debug_mask) & DL_DEBUG_FILES))
@@ -574,10 +603,10 @@ _mmap_segment(struct link_map *l, const struct loadcmd loadcmds[], size_t nloadc
     }
 
   extra_len = _extra_mmap(l, loadcmds, nloadcmds, c, mapstart);
-  errstring = _mmap_segment_filesz(l, c, mapstart, extra_len, fd);
+  errstring = _mmap_segment_filesz(l, c, mapstart, extra_len, fd, hp_flag);
   if (__glibc_unlikely (errstring != NULL))
     return errstring;
-  errstring = _mmap_segment_memsz(l, c, mapstart, extra_len, &memsz_len);
+  errstring = _mmap_segment_memsz(l, c, mapstart, extra_len, &memsz_len, hp_flag);
   if (__glibc_unlikely (errstring != NULL))
     return errstring;
 
diff --git a/elf/rtld.c b/elf/rtld.c
index 85596ee4..7b15e203 100644
--- a/elf/rtld.c
+++ b/elf/rtld.c
@@ -2556,6 +2556,7 @@ process_dl_hugepage (const char *dl_hugepage)
   } hpopts[] =
     {
       {DL_HUGEPAGE_LARGE_IN, DL_HUGEPAGE_LIB_LARGE_IN_FLAG},
+      {DL_HUGEPAGE_THP_IN, DL_HUGEPAGE_LIB_THP_IN_FLAG}
     };
 #define nhpopts (sizeof (hpopts) / sizeof (hpopts[0]))
 
diff --git a/sysdeps/generic/ldsodefs.h b/sysdeps/generic/ldsodefs.h
index a9fffd66..cbe9b48d 100644
--- a/sysdeps/generic/ldsodefs.h
+++ b/sysdeps/generic/ldsodefs.h
@@ -553,6 +553,8 @@ struct rtld_global_ro
 #define DL_HUGEPAGE_PROBE_FLAG  (1 << 31)
 #define DL_HUGEPAGE_LIB_LARGE_IN_FLAG  (1 << 30)
 #define DL_HUGEPAGE_LARGE_IN 1
+#define DL_HUGEPAGE_LIB_THP_IN_FLAG  (1 << 29)
+#define DL_HUGEPAGE_THP_IN 2
 #endif
   /* OS version.  */
   EXTERN unsigned int _dl_osversion;
-- 
2.27.0

