From f11bfc6ccad3e07fde4e40493635bd832d108477 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <vchuravy@users.noreply.github.com>
Date: Thu, 27 Apr 2023 16:29:25 -0400
Subject: [PATCH] Use NewPM for ASAN/MSAN (#49530)

Co-authored-by: Gabriel Baraldi <baraldigabriel@gmail.com>
Co-authored-by: Prem Chintalapudi <prem.chintalapudi@gmail.com>
---
 src/aotcompile.cpp |  4 ++++
 src/cgmemmgr.cpp   |  4 ++--
 src/jitlayers.h    | 13 +++++++++++--
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp
index 2a14e2a4fa0ab..b89cdf550171f 100644
--- a/src/aotcompile.cpp
+++ b/src/aotcompile.cpp
@@ -1775,6 +1775,7 @@ void addOptimizationPasses(legacy::PassManagerBase *PM, int opt_level,
                 PM->add(createCFGSimplificationPass(basicSimplifyCFGOptions));
             }
         }
+#if JL_LLVM_VERSION < 150000
 #if defined(_COMPILER_ASAN_ENABLED_)
         PM->add(createAddressSanitizerFunctionPass());
 #endif
@@ -1783,6 +1784,7 @@ void addOptimizationPasses(legacy::PassManagerBase *PM, int opt_level,
 #endif
 #if defined(_COMPILER_TSAN_ENABLED_)
         PM->add(createThreadSanitizerLegacyPassPass());
+#endif
 #endif
         return;
     }
@@ -1934,6 +1936,7 @@ void addOptimizationPasses(legacy::PassManagerBase *PM, int opt_level,
     }
     PM->add(createCombineMulAddPass());
     PM->add(createDivRemPairsPass());
+#if JL_LLVM_VERSION < 150000
 #if defined(_COMPILER_ASAN_ENABLED_)
     PM->add(createAddressSanitizerFunctionPass());
 #endif
@@ -1943,6 +1946,7 @@ void addOptimizationPasses(legacy::PassManagerBase *PM, int opt_level,
 #if defined(_COMPILER_TSAN_ENABLED_)
     PM->add(createThreadSanitizerLegacyPassPass());
 #endif
+#endif
 }
 
 // An LLVM module pass that just runs all julia passes in order. Useful for
diff --git a/src/cgmemmgr.cpp b/src/cgmemmgr.cpp
index 9f4d69137c0fd..15d28ff270c55 100644
--- a/src/cgmemmgr.cpp
+++ b/src/cgmemmgr.cpp
@@ -860,8 +860,8 @@ uint8_t *RTDyldMemoryManagerJL::allocateCodeSection(uintptr_t Size,
                                                     StringRef SectionName)
 {
     // allocating more than one code section can confuse libunwind.
-#if !defined(_COMPILER_MSAN_ENABLED_)
-    // TODO: Figure out why msan needs this.
+#if !defined(_COMPILER_MSAN_ENABLED_) && !defined(_COMPILER_ASAN_ENABLED_)
+    // TODO: Figure out why msan and now asan too need this.
     assert(!code_allocated);
     code_allocated = true;
 #endif
diff --git a/src/jitlayers.h b/src/jitlayers.h
index 7f07034586c80..f63f3a42842f1 100644
--- a/src/jitlayers.h
+++ b/src/jitlayers.h
@@ -42,7 +42,14 @@
 // and feature support (e.g. Windows, JITEventListeners for various profilers,
 // etc.). Thus, we currently only use JITLink where absolutely required, that is,
 // for Mac/aarch64.
-#if defined(_OS_DARWIN_) && defined(_CPU_AARCH64_) || defined(_COMPILER_ASAN_ENABLED_) || defined(JL_FORCE_JITLINK)
+// #define JL_FORCE_JITLINK
+
+#if defined(_COMPILER_ASAN_ENABLED_) || defined(_COMPILER_MSAN_ENABLED_) || defined(_COMPILER_TSAN_ENABLED_)
+# define HAS_SANITIZER
+#endif
+// The sanitizers don't play well with our memory manager
+
+#if defined(_OS_DARWIN_) && defined(_CPU_AARCH64_) || defined(JL_FORCE_JITLINK) || JL_LLVM_VERSION >= 150000 && defined(HAS_SANITIZER)
 # if JL_LLVM_VERSION < 130000
 #  pragma message("On aarch64-darwin, LLVM version >= 13 is required for JITLink; fallback suffers from occasional segfaults")
 # endif
@@ -93,7 +100,9 @@ struct OptimizationOptions {
 // for middle-end IR optimizations. However, we have not qualified the new
 // pass manager on our optimization pipeline yet, so this remains an optional
 // define
-// #define JL_USE_NEW_PM
+#if defined(HAS_SANITIZER) && JL_LLVM_VERSION >= 150000
+#define JL_USE_NEW_PM
+#endif
 
 struct NewPM {
     std::unique_ptr<TargetMachine> TM;
