From 0181aaa495bc6ef3dcd570ea5d37cb7b72375614 Mon Sep 17 00:00:00 2001
From: Guangli Dai <gdai@meta.com>
Date: Fri, 13 Sep 2024 15:52:22 -0700
Subject: [PATCH] Optimize edata_cmp_summary_compare when __uint128_t is
 available

---
 configure.ac                                     | 15 +++++++++++++++
 include/jemalloc/internal/edata.h                | 16 ++++++++++++++++
 .../internal/jemalloc_internal_defs.h.in         |  3 +++
 3 files changed, 34 insertions(+)

diff --git a/configure.ac b/configure.ac
index dcf357e4..2e7f610d 100644
--- a/configure.ac
+++ b/configure.ac
@@ -562,6 +562,21 @@ AC_CACHE_CHECK([asm volatile support],
 if test "x${je_cv_asm_volatile}" = "xyes"; then
   AC_DEFINE([JEMALLOC_HAVE_ASM_VOLATILE], [ ], [ ])
 fi
+AC_CACHE_CHECK([__int128 support],
+               [je_cv_int128],
+               AC_RUN_IFELSE([AC_LANG_PROGRAM(
+[[
+]],
+[[
+      __int128 temp = 0;
+      return temp;
+]])],
+[je_cv_int128=yes],
+[je_cv_int128=no],
+[je_cv_int128=no]))
+if test "x${je_cv_int128}" = "xyes"; then
+  AC_DEFINE([JEMALLOC_HAVE_INT128], [ ], [ ])
+fi
 
 LD_PRELOAD_VAR="LD_PRELOAD"
 so="so"
diff --git a/include/jemalloc/internal/edata.h b/include/jemalloc/internal/edata.h
index 17befd92..2381ccbc 100644
--- a/include/jemalloc/internal/edata.h
+++ b/include/jemalloc/internal/edata.h
@@ -666,6 +666,21 @@ edata_cmp_summary_get(const edata_t *edata) {
 	return result;
 }
 
+#ifdef JEMALLOC_HAVE_INT128
+JEMALLOC_ALWAYS_INLINE unsigned __int128
+edata_cmp_summary_encode(edata_cmp_summary_t src) {
+	return ((unsigned __int128)src.sn << 64) | src.addr;
+}
+
+static inline int
+edata_cmp_summary_comp(edata_cmp_summary_t a, edata_cmp_summary_t b) {
+    unsigned __int128 a_encoded = edata_cmp_summary_encode(a);
+    unsigned __int128 b_encoded = edata_cmp_summary_encode(b);
+    if (a_encoded < b_encoded) return -1;
+    if (a_encoded == b_encoded) return 0;
+    return 1;
+}
+#else
 static inline int
 edata_cmp_summary_comp(edata_cmp_summary_t a, edata_cmp_summary_t b) {
 	/*
@@ -683,6 +698,7 @@ edata_cmp_summary_comp(edata_cmp_summary_t a, edata_cmp_summary_t b) {
 	return (2 * ((a.sn > b.sn) - (a.sn < b.sn))) +
 	       ((a.addr > b.addr) - (a.addr < b.addr));
 }
+#endif
 
 static inline int
 edata_snad_comp(const edata_t *a, const edata_t *b) {
diff --git a/include/jemalloc/internal/jemalloc_internal_defs.h.in b/include/jemalloc/internal/jemalloc_internal_defs.h.in
index 7498bc48..7f369873 100644
--- a/include/jemalloc/internal/jemalloc_internal_defs.h.in
+++ b/include/jemalloc/internal/jemalloc_internal_defs.h.in
@@ -454,6 +454,9 @@
  */
 #undef JEMALLOC_HAVE_RDTSCP
 
+/* If defined, use __int128 for optimization. */
+#undef JEMALLOC_HAVE_INT128
+
 #include "jemalloc/internal/jemalloc_internal_overrides.h"
 
 #endif /* JEMALLOC_INTERNAL_DEFS_H_ */