[HotColdSplit] Disable splitting for sanitized functions

Splitting can make sanitizer errors harder to understand, as the
trapping instruction may not be in the function where the bug was
detected.

rdar://48142697

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@354931 91177308-0d34-0410-b5e6-96231b3b80d8

Author: vedantk

lib/Transforms/IPO/HotColdSplitting.cpp

@@ -111,10 +111,11 @@ bool unlikelyExecuted(BasicBlock &BB) {
   if (BB.isEHPad() || isa<ResumeInst>(BB.getTerminator()))
     return true;
 
-  // The block is cold if it calls/invokes a cold function.
+  // The block is cold if it calls/invokes a cold function. However, do not
+  // mark sanitizer traps as cold.
   for (Instruction &I : BB)
     if (auto CS = CallSite(&I))
-      if (CS.hasFnAttr(Attribute::Cold))
+      if (CS.hasFnAttr(Attribute::Cold) && !CS->getMetadata("nosanitize"))
         return true;
 
   // The block is cold if it has an unreachable terminator, unless it's
@@ -235,6 +236,12 @@ bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
   if (F.hasFnAttribute(Attribute::NoInline))
     return false;
 
+  if (F.hasFnAttribute(Attribute::SanitizeAddress) ||
+      F.hasFnAttribute(Attribute::SanitizeHWAddress) ||
+      F.hasFnAttribute(Attribute::SanitizeThread) ||
+      F.hasFnAttribute(Attribute::SanitizeMemory))
+    return false;
+
   return true;
 }
 

test/Transforms/HotColdSplit/X86/do-not-split.ll

@@ -110,6 +110,78 @@ if.end:                                           ; preds = %entry
   ret void
 }
 
+; CHECK-LABEL: @sanitize_address
+; CHECK-NOT: sanitize_address.cold.1
+define void @sanitize_address() sanitize_address {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_hwaddress
+; CHECK-NOT: sanitize_hwaddress.cold.1
+define void @sanitize_hwaddress() sanitize_hwaddress {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_thread
+; CHECK-NOT: sanitize_thread.cold.1
+define void @sanitize_thread() sanitize_thread {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_memory
+; CHECK-NOT: sanitize_memory.cold.1
+define void @sanitize_memory() sanitize_memory {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+declare void @llvm.trap() cold noreturn
+
+; CHECK-LABEL: @nosanitize_call
+; CHECK-NOT: nosanitize_call.cold.1
+define void @nosanitize_call() sanitize_memory {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @llvm.trap(), !nosanitize !2
+  unreachable
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 declare void @sink() cold