[MemorySanitizer] Support memcpy.inline and memset.inline

[melver]

Summary

Other sanitizers (ASan, TSan, see added tests) already handle
memcpy.inline and memset.inline by not relying on InstVisitor to turn
the intrinsics into calls. Only MSan instrumentation currently does not
support them due to missing InstVisitor callbacks.

Fix it by actually making InstVisitor handle Mem*InlineInst.

While the mem*.inline intrinsics promise no calls to external functions
as an optimization, for the sanitizers we need to break this guarantee
since access into the runtime is required either way, and performance
can no longer be guaranteed. All other cases, where generating a call is
incorrect, should instead use no_sanitize.

Fixes: https://github.com/llvm/llvm-project/issues/57048

[dvyukov]

llvm/test/Instrumentation/AddressSanitizer/mem-intrinsics.ll

34

llvm is not confused by 2 tail calls in a row?

[vitalybuka]

llvm/include/llvm/IR/InstVisitor.h

210–214

I believe this should be like this

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2583

we don't need this if it's properly fowarded to MemCpyInst

[vitalybuka]

llvm/include/llvm/IR/InstVisitor.h

210–214

because visitor here always delegate to the base class

[vitalybuka]

LGTM if you address the comments above

[vitalybuka]

If it's related to https://github.com/llvm/llvm-project/issues/57048, please include it into commit message

[melver]

• Make InstVisitor delegate Mem*InlineInst to Mem*Inst.

[melver]

llvm/test/Instrumentation/AddressSanitizer/mem-intrinsics.ll

34

No - the test passes, and just follows convention from above. Presumably it was done to make sure instrumentation isn't confused by tail calls? I can't say...

llvm/include/llvm/IR/InstVisitor.h

#include "llvm/IR/Instruction.def"
   // Handle the special intrinsic instruction classes.
   RetTy visitDbgDeclareInst(DbgDeclareInst &I)    { DELEGATE(DbgVariableIntrinsic);}
   RetTy visitDbgValueInst(DbgValueInst &I)        { DELEGATE(DbgVariableIntrinsic);}
   RetTy visitDbgVariableIntrinsic(DbgVariableIntrinsic &I)
                                                   { DELEGATE(DbgInfoIntrinsic);}
   RetTy visitDbgLabelInst(DbgLabelInst &I)        { DELEGATE(DbgInfoIntrinsic);}
   RetTy visitDbgInfoIntrinsic(DbgInfoIntrinsic &I){ DELEGATE(IntrinsicInst); }
   RetTy visitMemSetInst(MemSetInst &I)            { DELEGATE(MemIntrinsic); }
+  RetTy visitMemSetInlineInst(MemSetInlineInst &I){ DELEGATE(MemSetInst); }
   RetTy visitMemCpyInst(MemCpyInst &I)            { DELEGATE(MemTransferInst); }
-  RetTy visitMemCpyInlineInst(MemCpyInlineInst &I) {
-    DELEGATE(MemTransferInst);
-  }
+  RetTy visitMemCpyInlineInst(MemCpyInlineInst &I){ DELEGATE(MemCpyInst); }
   RetTy visitMemMoveInst(MemMoveInst &I)          { DELEGATE(MemTransferInst); }
   RetTy visitMemTransferInst(MemTransferInst &I)  { DELEGATE(MemIntrinsic); }
   RetTy visitMemIntrinsic(MemIntrinsic &I)        { DELEGATE(IntrinsicInst); }
   RetTy visitVAStartInst(VAStartInst &I)          { DELEGATE(IntrinsicInst); }
   RetTy visitVAEndInst(VAEndInst &I)              { DELEGATE(IntrinsicInst); }
   RetTy visitVACopyInst(VACopyInst &I)            { DELEGATE(IntrinsicInst); }
   RetTy visitIntrinsicInst(IntrinsicInst &I)      { DELEGATE(CallInst); }
   RetTy visitCallInst(CallInst &I)                { DELEGATE(CallBase); }
   RetTy visitInvokeInst(InvokeInst &I)            { DELEGATE(CallBase); }
   RetTy visitCallBrInst(CallBrInst &I)            { DELEGATE(CallBase); }
private:
   RetTy delegateCallInst(CallInst &I) {
     if (const Function *F = I.getCalledFunction()) {
       switch (F->getIntrinsicID()) {
       default:                     DELEGATE(IntrinsicInst);
       case Intrinsic::dbg_declare: DELEGATE(DbgDeclareInst);
       case Intrinsic::dbg_value:   DELEGATE(DbgValueInst);
       case Intrinsic::dbg_label:   DELEGATE(DbgLabelInst);
       case Intrinsic::memcpy:      DELEGATE(MemCpyInst);
+      case Intrinsic::memcpy_inline:
+        DELEGATE(MemCpyInlineInst);
       case Intrinsic::memmove:     DELEGATE(MemMoveInst);
       case Intrinsic::memset:      DELEGATE(MemSetInst);
+      case Intrinsic::memset_inline:
+        DELEGATE(MemSetInlineInst);
       case Intrinsic::vastart:     DELEGATE(VAStartInst);
       case Intrinsic::vaend:       DELEGATE(VAEndInst);
       case Intrinsic::vacopy:      DELEGATE(VACopyInst);
       case Intrinsic::not_intrinsic: break;
       }
     }
     DELEGATE(CallInst);
   }

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

void visitMemMoveInst(MemMoveInst &I) {
     IRB.CreateCall(
         MS.MemmoveFn,
         {IRB.CreatePointerCast(I.getArgOperand(0), IRB.getInt8PtrTy()),
          IRB.CreatePointerCast(I.getArgOperand(1), IRB.getInt8PtrTy()),
          IRB.CreateIntCast(I.getArgOperand(2), MS.IntptrTy, false)});
     I.eraseFromParent();
   }
 
-  // Similar to memmove: avoid copying shadow twice.
-  // This is somewhat unfortunate as it may slowdown small constant memcpys.
-  // FIXME: consider doing manual inline for small constant sizes and proper
-  // alignment.
+  /// Instrument memcpy
+  ///
+  /// Similar to memmove: avoid copying shadow twice. This is somewhat
+  /// unfortunate as it may slowdown small constant memcpys.
+  /// FIXME: consider doing manual inline for small constant sizes and proper
+  /// alignment.
+  ///
+  /// Note: This also handles memcpy.inline, which promises no calls to external
+  /// functions as an optimization. However, with instrumentation enabled this
+  /// is difficult to promise; additionally, we know that the MSan runtime
+  /// exists and provides __msan_memcpy(). Therefore, we assume that with
+  /// instrumentation it's safe to turn memcpy.inline into a call to
+  /// __msan_memcpy(). Should this be wrong, such as when implementing memcpy()
+  /// itself, instrumentation should be disabled with the no_sanitize attribute.
   void visitMemCpyInst(MemCpyInst &I) {
     getShadow(I.getArgOperand(1)); // Ensure shadow initialized
     IRBuilder<> IRB(&I);
     IRB.CreateCall(
         MS.MemcpyFn,
         {IRB.CreatePointerCast(I.getArgOperand(0), IRB.getInt8PtrTy()),
          IRB.CreatePointerCast(I.getArgOperand(1), IRB.getInt8PtrTy()),
          IRB.CreateIntCast(I.getArgOperand(2), MS.IntptrTy, false)});
     I.eraseFromParent();
   }
 
   // Same as memcpy.
   void visitMemSetInst(MemSetInst &I) {
     IRBuilder<> IRB(&I);
     IRB.CreateCall(
         MS.MemsetFn,
         {IRB.CreatePointerCast(I.getArgOperand(0), IRB.getInt8PtrTy()),
          IRB.CreateIntCast(I.getArgOperand(1), IRB.getInt32Ty(), false),
          IRB.CreateIntCast(I.getArgOperand(2), MS.IntptrTy, false)});
     I.eraseFromParent();
   }
 
   void visitVAStartInst(VAStartInst &I) {
     VAHelper->visitVAStartInst(I);
   }
 
   void visitVACopyInst(VACopyInst &I) {
     VAHelper->visitVACopyInst(I);
   }
 

llvm/test/Instrumentation/AddressSanitizer/mem-intrinsics.ll

 ; Test memory intrinsics instrumentation
 
 ; RUN: opt < %s -passes='asan-pipeline' -S | FileCheck --check-prefixes=CHECK,CHECK-PREFIX %s
 ; RUN: opt < %s -passes='asan-pipeline' -asan-kernel -S | FileCheck --check-prefixes=CHECK,CHECK-NOPREFIX %s
 ; RUN: opt < %s -passes='asan-pipeline' -asan-kernel -asan-kernel-mem-intrinsic-prefix -S | FileCheck --check-prefixes=CHECK,CHECK-PREFIX %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 target triple = "x86_64-unknown-linux-gnu"
 
 declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i1) nounwind
+declare void @llvm.memset.inline.p0i8.i64(i8* nocapture, i8, i64, i1) nounwind
 declare void @llvm.memmove.p0i8.p0i8.i64(i8* nocapture, i8* nocapture readonly, i64, i1) nounwind
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture readonly, i64, i1) nounwind
+declare void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* nocapture, i8* nocapture readonly, i64, i1) nounwind
 
 define void @memintr_test(i8* %a, i8* %b) nounwind uwtable sanitize_address {
   entry:
   tail call void @llvm.memset.p0i8.i64(i8* %a, i8 0, i64 100, i1 false)
   tail call void @llvm.memmove.p0i8.p0i8.i64(i8* %a, i8* %b, i64 100, i1 false)
   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* %a, i8* %b, i64 100, i1 false)
   ret void
 }
 ; CHECK-LABEL: memintr_test
 ; CHECK-PREFIX: @__asan_memset
 ; CHECK-PREFIX: @__asan_memmove
 ; CHECK-PREFIX: @__asan_memcpy
 ; CHECK-NOPREFIX: @memset
 ; CHECK-NOPREFIX: @memmove
 ; CHECK-NOPREFIX: @memcpy
 ; CHECK: ret void
 
+define void @memintr_inline_test(i8* %a, i8* %b) nounwind uwtable sanitize_address {
+  entry:
+  tail call void @llvm.memset.inline.p0i8.i64(i8* %a, i8 0, i64 100, i1 false)
+  tail call void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* %a, i8* %b, i64 100, i1 false)
+  ret void
+}
+; CHECK-LABEL: memintr_inline_test
+; CHECK-PREFIX: @__asan_memset
+; CHECK-PREFIX: @__asan_memcpy
+; CHECK-NOPREFIX: @memset
+; CHECK-NOPREFIX: @memcpy
+; CHECK: ret void
+
 define void @memintr_test_nosanitize(i8* %a, i8* %b) nounwind uwtable {
   entry:
   tail call void @llvm.memset.p0i8.i64(i8* %a, i8 0, i64 100, i1 false)
   tail call void @llvm.memmove.p0i8.p0i8.i64(i8* %a, i8* %b, i64 100, i1 false)
   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* %a, i8* %b, i64 100, i1 false)
   ret void
 }
 ; CHECK-LABEL: memintr_test_nosanitize

llvm/test/Instrumentation/MemorySanitizer/msan_basic.ll

 }
 
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1) nounwind
 
 ; CHECK-LABEL: @MemCpy
 ; CHECK: call i8* @__msan_memcpy
 ; CHECK: ret void
 
+; memset.inline
+define void @MemSetInline(i8* nocapture %x) nounwind uwtable sanitize_memory {
+entry:
+  call void @llvm.memset.inline.p0i8.i64(i8* %x, i8 42, i64 10, i1 false)
+  ret void
+}
+
+declare void @llvm.memset.inline.p0i8.i64(i8* nocapture, i8, i64, i1) nounwind
+
+; CHECK-LABEL: @MemSetInline
+; CHECK: call i8* @__msan_memset
+; CHECK: ret void
+
+; memcpy.inline
+define void @MemCpyInline(i8* nocapture %x, i8* nocapture %y) nounwind uwtable sanitize_memory {
+entry:
+  call void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* %x, i8* %y, i64 10, i1 false)
+  ret void
+}
+
+declare void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1) nounwind
+
+; CHECK-LABEL: @MemCpyInline
+; CHECK: call i8* @__msan_memcpy
+; CHECK: ret void
 
 ; memmove is lowered to a call
 define void @MemMove(i8* nocapture %x, i8* nocapture %y) nounwind uwtable sanitize_memory {
 entry:
   call void @llvm.memmove.p0i8.p0i8.i64(i8* %x, i8* %y, i64 10, i1 false)
   ret void
 }
 

llvm/test/Instrumentation/ThreadSanitizer/tsan_basic.ll

 ; CHECK-NEXT:   %1 = bitcast i32* %a to i8*
 ; CHECK-NEXT:   call void @__tsan_read4(i8* %1)
 ; CHECK-NEXT:   %tmp1 = load i32, i32* %a, align 4
 ; CHECK-NEXT:   call void @__tsan_func_exit()
 ; CHECK: ret i32
 
 
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1)
+declare void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1)
 declare void @llvm.memmove.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1)
 declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i1)
+declare void @llvm.memset.inline.p0i8.i64(i8* nocapture, i8, i64, i1)
 
 
 ; Check that tsan converts mem intrinsics back to function calls.
 
 define void @MemCpyTest(i8* nocapture %x, i8* nocapture %y) sanitize_thread {
 entry:
     tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %x, i8* align 4 %y, i64 16, i1 false)
     ret void
 ; CHECK: define void @MemCpyTest
 ; CHECK: call i8* @memcpy
 ; CHECK: ret void
 }
 
+define void @MemCpyInlineTest(i8* nocapture %x, i8* nocapture %y) sanitize_thread {
+entry:
+    tail call void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* align 4 %x, i8* align 4 %y, i64 16, i1 false)
+    ret void
+; CHECK: define void @MemCpyInlineTest
+; CHECK: call i8* @memcpy
+; CHECK: ret void
+}
+
 define void @MemMoveTest(i8* nocapture %x, i8* nocapture %y) sanitize_thread {
 entry:
     tail call void @llvm.memmove.p0i8.p0i8.i64(i8* align 4 %x, i8* align 4 %y, i64 16, i1 false)
     ret void
 ; CHECK: define void @MemMoveTest
 ; CHECK: call i8* @memmove
 ; CHECK: ret void
 }
 
 define void @MemSetTest(i8* nocapture %x) sanitize_thread {
 entry:
     tail call void @llvm.memset.p0i8.i64(i8* align 4 %x, i8 77, i64 16, i1 false)
     ret void
 ; CHECK: define void @MemSetTest
 ; CHECK: call i8* @memset
 ; CHECK: ret void
 }
 
+define void @MemSetInlineTest(i8* nocapture %x) sanitize_thread {
+entry:
+    tail call void @llvm.memset.inline.p0i8.i64(i8* align 4 %x, i8 77, i64 16, i1 false)
+    ret void
+; CHECK: define void @MemSetInlineTest
+; CHECK: call i8* @memset
+; CHECK: ret void
+}
+
 ; CHECK-LABEL: @SwiftError
 ; CHECK-NOT: __tsan_read
 ; CHECK-NOT: __tsan_write
 ; CHECK: ret
 define void @SwiftError(i8** swifterror) sanitize_thread {
   %swifterror_ptr_value = load i8*, i8** %0
   store i8* null, i8** %0
   %swifterror_addr = alloca swifterror i8*