[MSAN] Pass Origin by parameter to __msan_warning functions

[guiand]

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.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

786

This is an unconditional improvement, simply apply it in all cases. Keep __msan_warning_noreturn in the runtime library for backwards compatibility, but never issue a call to it in the instrumentation. Update the tests in test/Instrumentation/MemorySanitizer.

[guiand]

Removed dependency on eager-checks, and updated the tests

[eugenis]

LGTM

compiler-rt/lib/msan/msan.cpp

void __msan_warning_noreturn() {
   (void)sp;
   PrintWarning(pc, bp);
   if (__msan::flags()->print_stats)
     ReportStats();
   Printf("Exiting\n");
   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,
                 common_flags()->fast_unwind_on_fatal);
 }
 
 static void MsanOnDeadlySignal(int signo, void *siginfo, void *context) {
   HandleDeadlySignal(siginfo, context, GetTid(), &OnStackUnwind, nullptr);

compiler-rt/lib/msan/msan_interface_internal.h

 using __sanitizer::uu64;
 using __sanitizer::uu32;
 using __sanitizer::uu16;
 using __sanitizer::u64;
 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
 void __msan_maybe_warning_2(u16 s, u32 o);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_maybe_warning_4(u32 s, u32 o);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_maybe_warning_8(u64 s, u32 o);

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

private:
   Value *OriginTLS;
 
   /// Are the instrumentation callbacks set up?
   bool CallbacksInitialized = false;
 
   /// The run-time callback to print a warning.
   FunctionCallee WarningFn;
 
+  /// True if Origin is passed as a parameter, false if over TLS.
+  bool WarningFnTakesOrigin;
+
   // These arrays are indexed by log2(AccessSize).
   FunctionCallee MaybeWarningFn[kNumberOfAccessSizes];
   FunctionCallee MaybeStoreOriginFn[kNumberOfAccessSizes];
 
   /// Run-time helper that generates a new origin value for a stack
   /// allocation.
   FunctionCallee MsanSetAllocaOrigin4Fn;
 
void MemorySanitizer::createKernelApi(Module &M) {
   ParamOriginTLS = nullptr;
   VAArgTLS = nullptr;
   VAArgOriginTLS = nullptr;
   VAArgOverflowSizeTLS = nullptr;
   // OriginTLS is unused in the kernel.
   OriginTLS = nullptr;
 
   // __msan_warning() in the kernel takes an origin.
+  WarningFnTakesOrigin = true;
   WarningFn = M.getOrInsertFunction("__msan_warning", IRB.getVoidTy(),
                                     IRB.getInt32Ty());
   // Requests the per-task context state (kmsan_context_state*) from the
   // runtime library.
   MsanContextStateTy = StructType::get(
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
return new GlobalVariable(M, Ty, false, GlobalVariable::ExternalLinkage,
                               nullptr, Name, nullptr,
                               GlobalVariable::InitialExecTLSModel);
   });
 }
 
 /// Insert declarations for userspace-specific functions and globals.
 void MemorySanitizer::createUserspaceApi(Module &M) {
   IRBuilder<> IRB(*C);
+
+  // Enable origin as param with eager checks, as a space saving measure.
+  WarningFnTakesOrigin = ClEagerChecks;
+
   // Create the callback.
   // FIXME: this function should have "Cold" calling conv,
   // which is not yet implemented.
-  StringRef WarningFnName = Recover ? "__msan_warning"
-                                    : "__msan_warning_noreturn";
+  std::string WarningFnName = "__msan_warning";
+  if (WarningFnTakesOrigin)
+    WarningFnName += "_with_origin";
+  if (!Recover)
+    WarningFnName += "_noreturn";
+  if (WarningFnTakesOrigin)
+    WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy(), IRB.getInt32Ty());
+  else
     WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy());
 
   // Create the global TLS variables.
   RetvalTLS =
       getOrInsertGlobal(M, "__msan_retval_tls",
                         ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8));
 
   RetvalOriginTLS = getOrInsertGlobal(M, "__msan_retval_origin_tls", OriginTy);
 
for (StoreInst *SI : StoreList) {
                     OriginAlignment, InstrumentWithCalls);
     }
   }
 
   /// Helper function to insert a warning at IRB's current insert point.
   void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     if (!Origin)
       Origin = (Value *)IRB.getInt32(0);
-    if (MS.CompileKernel) {
+    if (MS.WarningFnTakesOrigin) {
+      assert(Origin->getType()->isIntegerTy());
       IRB.CreateCall(MS.WarningFn, Origin);
     } else {
       if (MS.TrackOrigins) {
         IRB.CreateStore(Origin, MS.OriginTLS);
       }
       IRB.CreateCall(MS.WarningFn, {});
     }
     IRB.CreateCall(MS.EmptyAsm, {});