[MSAN] Pass Origin by parameter to __msan_warning functions

Summary: Normally, the Origin is passed over TLS, which seems like it introduces unnecessary overhead. It's in the (extremely) cold path though, so the only overhead is in code size. But with eager-checks, calls to __msan_warning functions are extremely common, so this becomes a useful optimization. This can save ~5% code size. Reviewers: eugenis, vitalybuka Reviewed By: eugenis, vitalybuka Subscribers: hiraditya, #sanitizers, llvm-commits Tags: #sanitizers, #llvm Differential Revision: https://reviews.llvm.org/D81700
llvm · Jun 16, 2020 · b0ffa8b · b0ffa8b
1 parent 576fa5a
commit b0ffa8b
Show file tree

Hide file tree

Showing 14 changed files with 92 additions and 73 deletions.
diff --git a/compiler-rt/lib/msan/msan.cpp b/compiler-rt/lib/msan/msan.cpp
@@ -380,6 +380,28 @@ void __msan_warning_noreturn() {
   Die();
 }
 
+void __msan_warning_with_origin(u32 origin) {
+  GET_CALLER_PC_BP_SP;
+  (void)sp;
+  PrintWarningWithOrigin(pc, bp, origin);
+  if (__msan::flags()->halt_on_error) {
+    if (__msan::flags()->print_stats)
+      ReportStats();
+    Printf("Exiting\n");
+    Die();
+  }
+}
+
+void __msan_warning_with_origin_noreturn(u32 origin) {
+  GET_CALLER_PC_BP_SP;
+  (void)sp;
+  PrintWarningWithOrigin(pc, bp, origin);
+  if (__msan::flags()->print_stats)
+    ReportStats();
+  Printf("Exiting\n");
+  Die();
+}
+
 static void OnStackUnwind(const SignalContext &sig, const void *,
                           BufferedStackTrace *stack) {
   stack->Unwind(StackTrace::GetNextInstructionPc(sig.pc), sig.bp, sig.context,

diff --git a/compiler-rt/lib/msan/msan_interface_internal.h b/compiler-rt/lib/msan/msan_interface_internal.h
@@ -46,6 +46,12 @@ using __sanitizer::u32;
 using __sanitizer::u16;
 using __sanitizer::u8;
 
+// Versions of the above which take Origin as a parameter
+SANITIZER_INTERFACE_ATTRIBUTE
+void __msan_warning_with_origin(u32 origin);
+SANITIZER_INTERFACE_ATTRIBUTE __attribute__((noreturn)) void
+__msan_warning_with_origin_noreturn(u32 origin);
+
 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_maybe_warning_1(u8 s, u32 o);
 SANITIZER_INTERFACE_ATTRIBUTE

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -543,10 +543,6 @@ class MemorySanitizer {
   /// (x86_64-specific).
   Value *VAArgOverflowSizeTLS;
 
-  /// Thread-local space used to pass origin value to the UMR reporting
-  /// function.
-  Value *OriginTLS;
-
   /// Are the instrumentation callbacks set up?
   bool CallbacksInitialized = false;
 
@@ -715,10 +711,7 @@ void MemorySanitizer::createKernelApi(Module &M) {
   VAArgTLS = nullptr;
   VAArgOriginTLS = nullptr;
   VAArgOverflowSizeTLS = nullptr;
-  // OriginTLS is unused in the kernel.
-  OriginTLS = nullptr;
 
-  // __msan_warning() in the kernel takes an origin.
   WarningFn = M.getOrInsertFunction("__msan_warning", IRB.getVoidTy(),
                                     IRB.getInt32Ty());
   // Requests the per-task context state (kmsan_context_state*) from the
@@ -773,12 +766,14 @@ static Constant *getOrInsertGlobal(Module &M, StringRef Name, Type *Ty) {
 /// Insert declarations for userspace-specific functions and globals.
 void MemorySanitizer::createUserspaceApi(Module &M) {
   IRBuilder<> IRB(*C);
+
   // Create the callback.
   // FIXME: this function should have "Cold" calling conv,
   // which is not yet implemented.
-  StringRef WarningFnName = Recover ? "__msan_warning"
-                                    : "__msan_warning_noreturn";
-  WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy());
+  StringRef WarningFnName = Recover ? "__msan_warning_with_origin"
+                                    : "__msan_warning_with_origin_noreturn";
+  WarningFn =
+      M.getOrInsertFunction(WarningFnName, IRB.getVoidTy(), IRB.getInt32Ty());
 
   // Create the global TLS variables.
   RetvalTLS =
@@ -805,7 +800,6 @@ void MemorySanitizer::createUserspaceApi(Module &M) {
 
   VAArgOverflowSizeTLS =
       getOrInsertGlobal(M, "__msan_va_arg_overflow_size_tls", IRB.getInt64Ty());
-  OriginTLS = getOrInsertGlobal(M, "__msan_origin_tls", IRB.getInt32Ty());
 
   for (size_t AccessSizeIndex = 0; AccessSizeIndex < kNumberOfAccessSizes;
        AccessSizeIndex++) {
@@ -1216,14 +1210,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     if (!Origin)
       Origin = (Value *)IRB.getInt32(0);
-    if (MS.CompileKernel) {
-      IRB.CreateCall(MS.WarningFn, Origin);
-    } else {
-      if (MS.TrackOrigins) {
-        IRB.CreateStore(Origin, MS.OriginTLS);
-      }
-      IRB.CreateCall(MS.WarningFn, {});
-    }
+    assert(Origin->getType()->isIntegerTy());
+    IRB.CreateCall(MS.WarningFn, Origin);
     IRB.CreateCall(MS.EmptyAsm, {});
     // FIXME: Insert UnreachableInst if !MS.Recover?
     // This may invalidate some of the following checks and needs to be done

diff --git a/llvm/test/Instrumentation/MemorySanitizer/atomics.ll b/llvm/test/Instrumentation/MemorySanitizer/atomics.ll
@@ -54,7 +54,7 @@ entry:
 ; CHECK: store { i32, i1 } zeroinitializer,
 ; CHECK: icmp
 ; CHECK: br
-; CHECK: @__msan_warning
+; CHECK: @__msan_warning_with_origin
 ; CHECK: cmpxchg {{.*}} seq_cst seq_cst
 ; CHECK: store i32 0, {{.*}} @__msan_retval_tls
 ; CHECK: ret i32
@@ -73,7 +73,7 @@ entry:
 ; CHECK: store { i32, i1 } zeroinitializer,
 ; CHECK: icmp
 ; CHECK: br
-; CHECK: @__msan_warning
+; CHECK: @__msan_warning_with_origin
 ; CHECK: cmpxchg {{.*}} release monotonic
 ; CHECK: store i32 0, {{.*}} @__msan_retval_tls
 ; CHECK: ret i32

diff --git a/llvm/test/Instrumentation/MemorySanitizer/check-constant-shadow.ll b/llvm/test/Instrumentation/MemorySanitizer/check-constant-shadow.ll
@@ -14,7 +14,7 @@ entry:
 }
 
 ; CHECK-LABEL: @main
-; CHECK: call void @__msan_warning_noreturn
+; CHECK: call void @__msan_warning_with_origin_noreturn
 ; CHECK: ret i32 undef
 
 

diff --git a/llvm/test/Instrumentation/MemorySanitizer/check_access_address.ll b/llvm/test/Instrumentation/MemorySanitizer/check_access_address.ll
@@ -45,7 +45,7 @@ entry:
 ; CHECK: icmp
 ; CHECK: br i1
 ; CHECK: {{^[0-9]+}}:
-; CHECK: call void @__msan_warning_noreturn
+; CHECK: call void @__msan_warning_with_origin_noreturn
 ; CHECK: {{^[0-9]+}}:
 ; CHECK: xor
 ; CHECK: store

diff --git a/llvm/test/Instrumentation/MemorySanitizer/csr.ll b/llvm/test/Instrumentation/MemorySanitizer/csr.ll
@@ -28,7 +28,7 @@ entry:
 ; ADDR: %[[A:.*]] = load i64, i64* getelementptr inbounds {{.*}} @__msan_param_tls, i32 0, i32 0), align 8
 ; ADDR: %[[B:.*]] = icmp ne i64 %[[A]], 0
 ; ADDR: br i1 %[[B]], label {{.*}}, label
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 ; ADDR: call void @llvm.x86.sse.stmxcsr(
 ; ADDR: ret void
 
@@ -44,14 +44,14 @@ entry:
 ; CHECK: %[[A:.*]] = load i32, i32* %{{.*}}, align 1
 ; CHECK: %[[B:.*]] = icmp ne i32 %[[A]], 0
 ; CHECK: br i1 %[[B]], label {{.*}}, label
-; CHECK: call void @__msan_warning_noreturn()
+; CHECK: call void @__msan_warning_with_origin_noreturn(i32 0)
 ; CHECK: call void @llvm.x86.sse.ldmxcsr(
 ; CHECK: ret void
 
 ; ADDR-LABEL: @setcsr(
 ; ADDR: %[[A:.*]] = load i64, i64* getelementptr inbounds {{.*}} @__msan_param_tls, i32 0, i32 0), align 8
 ; ADDR: %[[B:.*]] = icmp ne i64 %[[A]], 0
 ; ADDR: br i1 %[[B]], label {{.*}}, label
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 ; ADDR: call void @llvm.x86.sse.ldmxcsr(
 ; ADDR: ret void
diff --git a/llvm/test/Instrumentation/MemorySanitizer/masked-store-load.ll b/llvm/test/Instrumentation/MemorySanitizer/masked-store-load.ll
@@ -54,12 +54,12 @@ entry:
 
 ; ADDR: %[[ADDRBAD:.*]] = icmp ne i64 %[[ADDRSHADOW]], 0
 ; ADDR: br i1 %[[ADDRBAD]], label {{.*}}, label {{.*}}
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 
 ; ADDR: %[[MASKSHADOWFLAT:.*]] = bitcast <4 x i1> %[[MASKSHADOW]] to i4
 ; ADDR: %[[MASKBAD:.*]] = icmp ne i4 %[[MASKSHADOWFLAT]], 0
 ; ADDR: br i1 %[[MASKBAD]], label {{.*}}, label {{.*}}
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 
 ; ADDR: tail call void @llvm.masked.store.v4i64.p0v4i64(<4 x i64> %v, <4 x i64>* %p, i32 1, <4 x i1> %mask)
 ; ADDR: ret void
@@ -94,12 +94,12 @@ entry:
 
 ; ADDR: %[[ADDRBAD:.*]] = icmp ne i64 %[[ADDRSHADOW]], 0
 ; ADDR: br i1 %[[ADDRBAD]], label {{.*}}, label {{.*}}
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 
 ; ADDR: %[[MASKSHADOWFLAT:.*]] = bitcast <4 x i1> %[[MASKSHADOW]] to i4
 ; ADDR: %[[MASKBAD:.*]] = icmp ne i4 %[[MASKSHADOWFLAT]], 0
 ; ADDR: br i1 %[[MASKBAD]], label {{.*}}, label {{.*}}
-; ADDR: call void @__msan_warning_noreturn()
+; ADDR: call void @__msan_warning_with_origin_noreturn(i32 0)
 
 ; ADDR: = call <4 x double> @llvm.masked.load.v4f64.p0v4f64(<4 x double>* %p, i32 1, <4 x i1> %mask, <4 x double> %v)
 ; ADDR: ret <4 x double>