From 9e123a833cc6f56381c46a1656a323f893fa2528 Mon Sep 17 00:00:00 2001
From: Ben Niu <beniu@microsoft.com>
Date: Thu, 21 Dec 2023 20:33:41 -0800
Subject: [PATCH] Leverage new Windows API TlsGetValue2 for performance

---
 include/jemalloc/internal/quantum.h |  2 +-
 include/jemalloc/internal/tsd_win.h | 30 ++++++++++++++++++++++++++---
 src/tsd.c                           |  4 ++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/include/jemalloc/internal/quantum.h b/include/jemalloc/internal/quantum.h
index a97f54ca..b4beb309 100644
--- a/include/jemalloc/internal/quantum.h
+++ b/include/jemalloc/internal/quantum.h
@@ -24,7 +24,7 @@
 #  ifdef __arm__
 #    define LG_QUANTUM		3
 #  endif
-#  ifdef __aarch64__
+#  if defined(__aarch64__) || defined(_M_ARM64)
 #    define LG_QUANTUM		4
 #  endif
 #  ifdef __hppa__
diff --git a/include/jemalloc/internal/tsd_win.h b/include/jemalloc/internal/tsd_win.h
index 8ec7eda7..4b40a8ab 100644
--- a/include/jemalloc/internal/tsd_win.h
+++ b/include/jemalloc/internal/tsd_win.h
@@ -15,6 +15,16 @@ typedef struct {
 extern DWORD tsd_tsd;
 extern tsd_wrapper_t tsd_boot_wrapper;
 extern bool tsd_booted;
+#if defined(_M_ARM64EC)
+#define JEMALLOC_WIN32_TLSGETVALUE2 0
+#else
+#define JEMALLOC_WIN32_TLSGETVALUE2 1
+#endif
+#if JEMALLOC_WIN32_TLSGETVALUE2
+typedef LPVOID (WINAPI *TGV2)(DWORD dwTlsIndex);
+extern TGV2 tls_get_value2;
+extern HMODULE tgv2_mod;
+#endif
 
 /* Initialization/cleanup. */
 JEMALLOC_ALWAYS_INLINE bool
@@ -49,9 +59,17 @@ tsd_wrapper_set(tsd_wrapper_t *wrapper) {
 
 JEMALLOC_ALWAYS_INLINE tsd_wrapper_t *
 tsd_wrapper_get(bool init) {
-	DWORD error = GetLastError();
-	tsd_wrapper_t *wrapper = (tsd_wrapper_t *) TlsGetValue(tsd_tsd);
-	SetLastError(error);
+	tsd_wrapper_t *wrapper;
+#if JEMALLOC_WIN32_TLSGETVALUE2
+	if (tls_get_value2 != NULL) {
+		wrapper = (tsd_wrapper_t *) tls_get_value2(tsd_tsd);
+	} else
+#endif
+	{
+		DWORD error = GetLastError();
+		wrapper = (tsd_wrapper_t *) TlsGetValue(tsd_tsd);
+		SetLastError(error);
+	}
 
 	if (init && unlikely(wrapper == NULL)) {
 		wrapper = (tsd_wrapper_t *)
@@ -78,6 +96,12 @@ tsd_boot0(void) {
 	}
 	_malloc_tsd_cleanup_register(&tsd_cleanup_wrapper);
 	tsd_wrapper_set(&tsd_boot_wrapper);
+#if JEMALLOC_WIN32_TLSGETVALUE2
+	tgv2_mod = LoadLibraryA("api-ms-win-core-processthreads-l1-1-8.dll");
+	if (tgv2_mod != NULL) {
+		tls_get_value2 = (TGV2)GetProcAddress(tgv2_mod, "TlsGetValue2");
+	}
+#endif
 	tsd_booted = true;
 	return false;
 }
diff --git a/src/tsd.c b/src/tsd.c
index a4db8e36..c9ae2d64 100644
--- a/src/tsd.c
+++ b/src/tsd.c
@@ -25,6 +25,10 @@ bool tsd_booted = false;
 DWORD tsd_tsd;
 tsd_wrapper_t tsd_boot_wrapper = {false, TSD_INITIALIZER};
 bool tsd_booted = false;
+#if JEMALLOC_WIN32_TLSGETVALUE2
+TGV2 tls_get_value2 = NULL;
+HMODULE tgv2_mod = NULL;
+#endif
 #else
 
 /*