[MSan] Add instrumentation for SystemZ

[iii]

Summary

This patch establishes memory layout and adds instrumentation. It does
not add runtime support and does not enable MSan, which will be done
separately.

Memory layout is based on PPC64, with the exception that XorMask
is not used - low and high memory addresses are chosen in a way that
applying AndMask to low and high memory produces non-overlapping
results.

VarArgHelper is based on AMD64. It might be tempting to share some
code between the two implementations, but we need to keep in mind that
all the ABI similarities are coincidental, and therefore any such
sharing might backfire.

copyRegSaveArea() indiscriminately copies the entire register save area
shadow, however, fragments thereof not filled by the corresponding
visitCallSite() invocation contain irrelevant data. Whether or not this
can lead to practical problems is unclear, hence a simple TODO comment.
Note that the behavior of the related copyOverflowArea() is correct: it
copies only the vararg-related fragment of the overflow area shadow.

VarArgHelper test is based on the AArch64 one.

s390x ABI requires that arguments are zero-extended to 64 bits. This is
particularly important for msan_maybe_warning_*() and
msan_maybe_store_origin_*() shadow and origin arguments, since non
zeroed upper parts thereof confuse these functions. Therefore, add ZExt
attribute to the corresponding parameters.

Add ZExt attribute checks to msan-basic.ll. Since with
-msan-instrumentation-with-call-threshold=0 instrumentation looks quite
different, introduce the new CHECK-CALLS check prefix.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

815

Add a test for zext attribute somewhere, ex. in msan_basic.ll

llvm/test/Instrumentation/MemorySanitizer/SystemZ/vararg.ll

97

CHECK-LABEL: @many_args

[vitalybuka]

LGTM

[iii]

Added a test for zext attribute to msan_basic.ll, added missing CHECK-LABEL: @many_args.

[iii]

Fixed a typo in vararg.ll test.

[uweigand]

Had a look at the vararg handling. Reviewed only the ABI-relevant parts.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

4637

typo?

4662

Is argument type coercion also taken care of? This is about the implicit promotion to 64 bits.

For example, a 32-bit integer is passed as zero- or sign-extended 64-bit value, so all 8 bytes of the slot in either the GPR save area or overflow area are filled by the caller and may be validly accessed by the callee.

In clang's SystemZABIInfo::EmitVAArg this is done by this call:

if (AI.getCoerceToType())
  ArgTy = AI.getCoerceToType();

4664

We recently added -msoft-float support in particular to support kernel builds.

I guess this also needs to be supported here ... (in soft-float mode, floating-point arguments are passed in GPRs).

4669

What about vector arguments? The variable arguments of vector type will always be in memory, but the fixed portion of the argument list can still include vector type arguments that would be passed in vector registers (and not via memory, so your memory overflow offset would be off).

4696

Might be best to add an ArgAllocSize <= ArgSize assertion here.

4716

Actually, 32-bit floats are passed in the high part of the register, so there should be no offset here. (This applies only for the in-memory FPR copies. For floating-point arguments passed in the overflow area, the offset is there.)

[iii]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

4662

No, we get the original (non-coerced) type here. The offset calculation is still correct in presence of type coercion, because the slot size is always 8, but the code currently lacks instrumentation that would fill the shadow for the extensions bits. I'll add it.

4664

I couldn't find a way to query floating point ABI, so I'll use the indirect way to detect this and assume that userspace is always hardfloat, and kernel is always softfloat.

[iii]

Implemented Ulrich's suggestions.

[iii]

Resolved shadow zext/sext TODO using CreateShadowCast.
Resolved softfloat TODO using getFnAttribute("use-soft-float").
Added shadow propagation checks to tests.

[uweigand]

The Z ABI implementation now looks correct to me, except for one corner case: an LLVM-level argument type of f128 is passed via implicit reference. Now, in most cases this is already handled at the clang level, i.e. when you have a "long double" at C source level, you'll already see a pointer type in the LLVM IR instead. However, there are still a few cases where there is a f128 at the LLVM IR level, e.g. as arguments to some compiler builtin / libgcc routines. Those are only transformed to implicit reference in the LLVM back-end, so I believe for completeness this case should also be handled here.

[iii]

Added fp128 handling and a test for it.

[iii]

Added i128 handling and a test for it.

[uweigand]

Just a minor cosmetic suggestion inline. Otherwise the Z ABI parts now all LGTM.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

4667

The size check now seems redundant with the i128 check. OTOH maybe it would be preferable to move the fp128 and i128 checks to here and return a new ArgKind::Indirect -- that way all the type classification checks would be in one place.

[iii]

Introduced ArgKind::Indirect.

[uweigand]

Thanks!

[iii]

@eugenis now that VarArgSystemZHelper part seems to be in order, could you please have another look?

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

397

wrong indent here and below

[iii]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

397

I tried changing this, but the result doesn't pass arcanist linter. Apparently it thinks that the existing indentation is wrong.

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

397

Ha, you (and arcanist) are correct, this is actually the right indentation.
I don't want to reformat the entire file at once because it would mess with the history.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 // ppc64 Linux
 static const MemoryMapParams Linux_PowerPC64_MemoryMapParams = {
   0xE00000000000,  // AndMask
   0x100000000000,  // XorMask
   0x080000000000,  // ShadowBase
   0x1C0000000000,  // OriginBase
 };
 
+// s390x Linux
+static const MemoryMapParams Linux_S390X_MemoryMapParams = {
+    0xC00000000000, // AndMask
+    0,              // XorMask (not used)
+    0x080000000000, // ShadowBase
+    0x1C0000000000, // OriginBase
+};
+
 // aarch64 Linux
 static const MemoryMapParams Linux_AArch64_MemoryMapParams = {
   0,               // AndMask (not used)
   0x06000000000,   // XorMask
   0,               // ShadowBase (not used)
   0x01000000000,   // OriginBase
 };
 
static const PlatformMemoryMapParams Linux_MIPS_MemoryMapParams = {
   &Linux_MIPS64_MemoryMapParams,
 };
 
 static const PlatformMemoryMapParams Linux_PowerPC_MemoryMapParams = {
   nullptr,
   &Linux_PowerPC64_MemoryMapParams,
 };
 
+static const PlatformMemoryMapParams Linux_S390_MemoryMapParams = {
+    nullptr,
+    &Linux_S390X_MemoryMapParams,
+};
+
 static const PlatformMemoryMapParams Linux_ARM_MemoryMapParams = {
   nullptr,
   &Linux_AArch64_MemoryMapParams,
 };
 
 static const PlatformMemoryMapParams FreeBSD_X86_MemoryMapParams = {
   &FreeBSD_I386_MemoryMapParams,
   &FreeBSD_X86_64_MemoryMapParams,
public:
   bool sanitizeFunction(Function &F, TargetLibraryInfo &TLI);
 
 private:
   friend struct MemorySanitizerVisitor;
   friend struct VarArgAMD64Helper;
   friend struct VarArgMIPS64Helper;
   friend struct VarArgAArch64Helper;
   friend struct VarArgPowerPC64Helper;
+  friend struct VarArgSystemZHelper;
 
   void initializeModule(Module &M);
   void initializeCallbacks(Module &M);
   void createKernelApi(Module &M);
   void createUserspaceApi(Module &M);
 
   /// True if we're compiling the Linux kernel.
   bool CompileKernel;
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++) {
     unsigned AccessSize = 1 << AccessSizeIndex;
     std::string FunctionName = "__msan_maybe_warning_" + itostr(AccessSize);
+    SmallVector<std::pair<unsigned, Attribute>, 2> MaybeWarningFnAttrs;
+    MaybeWarningFnAttrs.push_back(std::make_pair(
+        AttributeList::FirstArgIndex, Attribute::get(*C, Attribute::ZExt)));
+    MaybeWarningFnAttrs.push_back(std::make_pair(
+        AttributeList::FirstArgIndex + 1, Attribute::get(*C, Attribute::ZExt)));
     MaybeWarningFn[AccessSizeIndex] = M.getOrInsertFunction(
-        FunctionName, IRB.getVoidTy(), IRB.getIntNTy(AccessSize * 8),
-        IRB.getInt32Ty());
+        FunctionName, AttributeList::get(*C, MaybeWarningFnAttrs),
+        IRB.getVoidTy(), IRB.getIntNTy(AccessSize * 8), IRB.getInt32Ty());
 
     FunctionName = "__msan_maybe_store_origin_" + itostr(AccessSize);
+    SmallVector<std::pair<unsigned, Attribute>, 2> MaybeStoreOriginFnAttrs;
+    MaybeStoreOriginFnAttrs.push_back(std::make_pair(
+        AttributeList::FirstArgIndex, Attribute::get(*C, Attribute::ZExt)));
+    MaybeStoreOriginFnAttrs.push_back(std::make_pair(
+        AttributeList::FirstArgIndex + 2, Attribute::get(*C, Attribute::ZExt)));
     MaybeStoreOriginFn[AccessSizeIndex] = M.getOrInsertFunction(
-        FunctionName, IRB.getVoidTy(), IRB.getIntNTy(AccessSize * 8),
-        IRB.getInt8PtrTy(), IRB.getInt32Ty());
+        FunctionName, AttributeList::get(*C, MaybeStoreOriginFnAttrs),
+        IRB.getVoidTy(), IRB.getIntNTy(AccessSize * 8), IRB.getInt8PtrTy(),
+        IRB.getInt32Ty());
   }
 
   MsanSetAllocaOrigin4Fn = M.getOrInsertFunction(
     "__msan_set_alloca_origin4", IRB.getVoidTy(), IRB.getInt8PtrTy(), IntptrTy,
     IRB.getInt8PtrTy(), IntptrTy);
   MsanPoisonStackFn =
       M.getOrInsertFunction("__msan_poison_stack", IRB.getVoidTy(),
                             IRB.getInt8PtrTy(), IntptrTy);
switch (TargetTriple.getOS()) {
           case Triple::mips64:
           case Triple::mips64el:
             MapParams = Linux_MIPS_MemoryMapParams.bits64;
             break;
           case Triple::ppc64:
           case Triple::ppc64le:
             MapParams = Linux_PowerPC_MemoryMapParams.bits64;
             break;
+          case Triple::systemz:
+            MapParams = Linux_S390_MemoryMapParams.bits64;
+            break;
           case Triple::aarch64:
           case Triple::aarch64_be:
             MapParams = Linux_ARM_MemoryMapParams.bits64;
             break;
           default:
             report_fatal_error("unsupported architecture");
         }
         break;
for (size_t i = 0, n = VAStartInstrumentationList.size(); i < n; i++) {
           MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
                                  Alignment, /*isStore*/ true);
       IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
                        CopySize);
     }
   }
 };
 
+/// SystemZ-specific implementation of VarArgHelper.
+struct VarArgSystemZHelper : public VarArgHelper {
+  static const unsigned SystemZGpOffset = 16;
+  static const unsigned SystemZGpEndOffset = 56;
+  static const unsigned SystemZFpOffset = 128;
+  static const unsigned SystemZFpEndOffset = 160;
+  static const unsigned SysteZRegSaveAreaSize = 160;
+  static const unsigned SystemZOverflowOffset = 160;
+  static const unsigned SystemZVAListTagSize = 32;
+  static const unsigned SystemZOverflowArgAreaPtrOffset = 16;
+  static const unsigned SystemZRegSaveAreaPtrOffset = 24;
+
+  Function &F;
+  MemorySanitizer &MS;
+  MemorySanitizerVisitor &MSV;
+  Value *VAArgTLSCopy = nullptr;
+  Value *VAArgTLSOriginCopy = nullptr;
+  Value *VAArgOverflowSize = nullptr;
+
+  SmallVector<CallInst *, 16> VAStartInstrumentationList;
+
+  enum ArgKind { AK_GeneralPurpose, AK_FloatingPoint, AK_Memory };
+
+  VarArgSystemZHelper(Function &F, MemorySanitizer &MS,
+                      MemorySanitizerVisitor &MSV)
+      : F(F), MS(MS), MSV(MSV) {}
+
+  ArgKind classifyArgument(Value *arg) {
+    Type *T = arg->getType();
+    // T is a SystemZABIInfo output, and there are only a few possibilities
+    // of what it can be. In particular, single element structs and large
+    // types have already been taken care of.
+    if (T->isFloatingPointTy())
+      return AK_FloatingPoint;
+    if (T->isIntegerTy() && T->getPrimitiveSizeInBits() <= 64)
+      return AK_GeneralPurpose;
+    if (T->isPointerTy())
+      return AK_GeneralPurpose;
+    return AK_Memory;
+  }
+
+  void visitCallSite(CallSite &CS, IRBuilder<> &IRB) override {
+    unsigned GpOffset = SystemZGpOffset;
+    unsigned FpOffset = SystemZFpOffset;
+    unsigned OverflowOffset = SystemZOverflowOffset;
+    const DataLayout &DL = F.getParent()->getDataLayout();
+    for (CallSite::arg_iterator ArgIt = CS.arg_begin(), End = CS.arg_end();
+         ArgIt != End; ++ArgIt) {
+      Value *A = *ArgIt;
+      unsigned ArgNo = CS.getArgumentNo(ArgIt);
+      bool IsFixed = ArgNo < CS.getFunctionType()->getNumParams();
+      // SystemZABIInfo does not produce ByVal parameters.
+      assert(!CS.paramHasAttr(ArgNo, Attribute::ByVal));
+      ArgKind AK = classifyArgument(A);
+      if (AK == AK_GeneralPurpose && GpOffset >= SystemZGpEndOffset)
+        AK = AK_Memory;
+      if (AK == AK_FloatingPoint && FpOffset >= SystemZFpEndOffset)
+        AK = AK_Memory;
+      Value *ShadowBase = nullptr, *OriginBase = nullptr;
+      switch (AK) {
+      case AK_GeneralPurpose: {
+        // Always keep track of GpOffset, but store shadow only for varargs.
+        uint64_t ArgSize = 8;
+        if (GpOffset + ArgSize <= kParamTLSSize) {
+          if (!IsFixed) {
+            uint64_t ArgAllocSize = DL.getTypeAllocSize(A->getType());
+            ShadowBase = getShadowPtrForVAArgument(
+                A->getType(), IRB, GpOffset + (ArgSize - ArgAllocSize));
+            if (MS.TrackOrigins)
+              OriginBase = getOriginPtrForVAArgument(
+                  A->getType(), IRB, GpOffset + (ArgSize - ArgAllocSize));
+          }
+          GpOffset += ArgSize;
+        } else {
+          GpOffset = kParamTLSSize;
+        }
+        break;
+      }
+      case AK_FloatingPoint: {
+        // Always keep track of FpOffset, but store shadow only for varargs.
+        uint64_t ArgSize = 8;
+        if (FpOffset + ArgSize <= kParamTLSSize) {
+          if (!IsFixed) {
+            uint64_t ArgAllocSize = DL.getTypeAllocSize(A->getType());
+            ShadowBase = getShadowPtrForVAArgument(
+                A->getType(), IRB, FpOffset + (ArgSize - ArgAllocSize));
+            if (MS.TrackOrigins)
+              OriginBase = getOriginPtrForVAArgument(
+                  A->getType(), IRB, FpOffset + (ArgSize - ArgAllocSize));
+          }
+          FpOffset += ArgSize;
+        } else {
+          FpOffset = kParamTLSSize;
+        }
+        break;
+      }
+      case AK_Memory: {
+        // Keep track of OverflowOffset and store shadow only for varargs.
+        uint64_t ArgAllocSize = DL.getTypeAllocSize(A->getType());
+        uint64_t ArgSize = alignTo(ArgAllocSize, 8);
+        if (!IsFixed) {
+          if (OverflowOffset + ArgSize <= kParamTLSSize) {
+            ShadowBase = getShadowPtrForVAArgument(
+                A->getType(), IRB, OverflowOffset + (ArgSize - ArgAllocSize));
+            if (MS.TrackOrigins)
+              OriginBase = getOriginPtrForVAArgument(
+                  A->getType(), IRB, OverflowOffset + (ArgSize - ArgAllocSize));
+            OverflowOffset += ArgSize;
+          } else {
+            OverflowOffset = kParamTLSSize;
+          }
+        }
+        break;
+      }
+      }
+      if (!ShadowBase)
+        continue;
+      Value *Shadow = MSV.getShadow(A);
+      IRB.CreateAlignedStore(Shadow, ShadowBase, kShadowTLSAlignment);
+      if (MS.TrackOrigins) {
+        Value *Origin = MSV.getOrigin(A);
+        unsigned StoreSize = DL.getTypeStoreSize(Shadow->getType());
+        MSV.paintOrigin(IRB, Origin, OriginBase, StoreSize,
+                        std::max(kShadowTLSAlignment, kMinOriginAlignment));
+      }
+    }
+    Constant *OverflowSize = ConstantInt::get(
+        IRB.getInt64Ty(), OverflowOffset - SystemZOverflowOffset);
+    IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
+  }
+
+  Value *getShadowPtrForVAArgument(Type *Ty, IRBuilder<> &IRB,
+                                   unsigned ArgOffset) {
+    Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
+    Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
+    return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
+                              "_msarg_va_s");
+  }
+
+  Value *getOriginPtrForVAArgument(Type *Ty, IRBuilder<> &IRB, int ArgOffset) {
+    Value *Base = IRB.CreatePointerCast(MS.VAArgOriginTLS, MS.IntptrTy);
+    Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
+    return IRB.CreateIntToPtr(Base, PointerType::get(MS.OriginTy, 0),
+                              "_msarg_va_o");
+  }
+
+  void unpoisonVAListTagForInst(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Value *VAListTag = I.getArgOperand(0);
+    Value *ShadowPtr, *OriginPtr;
+    const Align Alignment = Align(8);
+    std::tie(ShadowPtr, OriginPtr) =
+        MSV.getShadowOriginPtr(VAListTag, IRB, IRB.getInt8Ty(), Alignment,
+                               /*isStore*/ true);
+    IRB.CreateMemSet(ShadowPtr, Constant::getNullValue(IRB.getInt8Ty()),
+                     SystemZVAListTagSize, Alignment, false);
+  }
+
+  void visitVAStartInst(VAStartInst &I) override {
+    VAStartInstrumentationList.push_back(&I);
+    unpoisonVAListTagForInst(I);
+  }
+
+  void visitVACopyInst(VACopyInst &I) override { unpoisonVAListTagForInst(I); }
+
+  void copyRegSaveArea(IRBuilder<> &IRB, Value *VAListTag) {
+    Type *RegSaveAreaPtrTy = Type::getInt64PtrTy(*MS.C);
+    Value *RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
+        IRB.CreateAdd(
+            IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
+            ConstantInt::get(MS.IntptrTy, SystemZRegSaveAreaPtrOffset)),
+        PointerType::get(RegSaveAreaPtrTy, 0));
+    Value *RegSaveAreaPtr = IRB.CreateLoad(RegSaveAreaPtrTy, RegSaveAreaPtrPtr);
+    Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
+    const Align Alignment = Align(8);
+    std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
+        MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(), Alignment,
+                               /*isStore*/ true);
+    // TODO(iii): copy only fragments filled by visitCallSite()
+    IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
+                     SysteZRegSaveAreaSize);
+    if (MS.TrackOrigins)
+      IRB.CreateMemCpy(RegSaveAreaOriginPtr, Alignment, VAArgTLSOriginCopy,
+                       Alignment, SysteZRegSaveAreaSize);
+  }
+
+  void copyOverflowArea(IRBuilder<> &IRB, Value *VAListTag) {
+    Type *OverflowArgAreaPtrTy = Type::getInt64PtrTy(*MS.C);
+    Value *OverflowArgAreaPtrPtr = IRB.CreateIntToPtr(
+        IRB.CreateAdd(
+            IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
+            ConstantInt::get(MS.IntptrTy, SystemZOverflowArgAreaPtrOffset)),
+        PointerType::get(OverflowArgAreaPtrTy, 0));
+    Value *OverflowArgAreaPtr =
+        IRB.CreateLoad(OverflowArgAreaPtrTy, OverflowArgAreaPtrPtr);
+    Value *OverflowArgAreaShadowPtr, *OverflowArgAreaOriginPtr;
+    const Align Alignment = Align(8);
+    std::tie(OverflowArgAreaShadowPtr, OverflowArgAreaOriginPtr) =
+        MSV.getShadowOriginPtr(OverflowArgAreaPtr, IRB, IRB.getInt8Ty(),
+                               Alignment, /*isStore*/ true);
+    Value *SrcPtr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), VAArgTLSCopy,
+                                           SystemZOverflowOffset);
+    IRB.CreateMemCpy(OverflowArgAreaShadowPtr, Alignment, SrcPtr, Alignment,
+                     VAArgOverflowSize);
+    if (MS.TrackOrigins) {
+      SrcPtr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), VAArgTLSOriginCopy,
+                                      SystemZOverflowOffset);
+      IRB.CreateMemCpy(OverflowArgAreaOriginPtr, Alignment, SrcPtr, Alignment,
+                       VAArgOverflowSize);
+    }
+  }
+
+  void finalizeInstrumentation() override {
+    assert(!VAArgOverflowSize && !VAArgTLSCopy &&
+           "finalizeInstrumentation called twice");
+    if (!VAStartInstrumentationList.empty()) {
+      // If there is a va_start in this function, make a backup copy of
+      // va_arg_tls somewhere in the function entry block.
+      IRBuilder<> IRB(MSV.ActualFnStart->getFirstNonPHI());
+      VAArgOverflowSize =
+          IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+      Value *CopySize =
+          IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, SystemZOverflowOffset),
+                        VAArgOverflowSize);
+      VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+      IRB.CreateMemCpy(VAArgTLSCopy, Align(8), MS.VAArgTLS, Align(8), CopySize);
+      if (MS.TrackOrigins) {
+        VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+        IRB.CreateMemCpy(VAArgTLSOriginCopy, Align(8), MS.VAArgOriginTLS,
+                         Align(8), CopySize);
+      }
+    }
+
+    // Instrument va_start.
+    // Copy va_list shadow from the backup copy of the TLS contents.
+    for (size_t i = 0, n = VAStartInstrumentationList.size(); i < n; i++) {
+      CallInst *OrigInst = VAStartInstrumentationList[i];
+      IRBuilder<> IRB(OrigInst->getNextNode());
+      Value *VAListTag = OrigInst->getArgOperand(0);
+      copyRegSaveArea(IRB, VAListTag);
+      copyOverflowArea(IRB, VAListTag);
+    }
+  }
+};
+
 /// A no-op implementation of VarArgHelper.
 struct VarArgNoOpHelper : public VarArgHelper {
   VarArgNoOpHelper(Function &F, MemorySanitizer &MS,
                    MemorySanitizerVisitor &MSV) {}
 
   void visitCallSite(CallSite &CS, IRBuilder<> &IRB) override {}
 
   void visitVAStartInst(VAStartInst &I) override {}
if (TargetTriple.getArch() == Triple::x86_64)
     return new VarArgAMD64Helper(Func, Msan, Visitor);
   else if (TargetTriple.isMIPS64())
     return new VarArgMIPS64Helper(Func, Msan, Visitor);
   else if (TargetTriple.getArch() == Triple::aarch64)
     return new VarArgAArch64Helper(Func, Msan, Visitor);
   else if (TargetTriple.getArch() == Triple::ppc64 ||
            TargetTriple.getArch() == Triple::ppc64le)
     return new VarArgPowerPC64Helper(Func, Msan, Visitor);
+  else if (TargetTriple.getArch() == Triple::systemz)
+    return new VarArgSystemZHelper(Func, Msan, Visitor);
   else
     return new VarArgNoOpHelper(Func, Msan, Visitor);
 }
 
 bool MemorySanitizer::sanitizeFunction(Function &F, TargetLibraryInfo &TLI) {
   if (!CompileKernel && F.getName() == kMsanModuleCtorName)
     return false;
 

llvm/test/Instrumentation/MemorySanitizer/SystemZ/vararg.ll

+; RUN: opt < %s -S -passes=msan 2>&1 | FileCheck %s
+; RUN: opt < %s -msan -S | FileCheck %s
+
+target triple = "s390x-unknown-linux-gnu"
+
+%struct.__va_list = type { i64, i64, i8*, i8* }
+
+define i64 @foo(i64 %guard, ...) {
+  %vl = alloca %struct.__va_list, align 8
+  %1 = bitcast %struct.__va_list* %vl to i8*
+  call void @llvm.lifetime.start.p0i8(i64 32, i8* %1)
+  call void @llvm.va_start(i8* %1)
+  call void @llvm.va_end(i8* %1)
+  call void @llvm.lifetime.end.p0i8(i64 32, i8* %1)
+  ret i64 0
+}
+
+; First check if the variadic shadow values are saved in stack with correct
+; size (which is 160 - size of the register save area).
+
+; CHECK-LABEL: @foo
+; CHECK: [[A:%.*]] = load {{.*}} @__msan_va_arg_overflow_size_tls
+; CHECK: [[B:%.*]] = add i64 160, [[A]]
+; CHECK: alloca {{.*}} [[B]]
+
+; We expect two memcpy operations: one for the register save area, and one for
+; the overflow arg area.
+
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 {{%.*}}, i8* align 8 {{%.*}}, i64 160, i1 false)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 {{%.*}}, i8* align 8 {{%.*}}, i64 [[A]], i1 false)
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+declare void @llvm.va_start(i8*) #2
+declare void @llvm.va_end(i8*) #2
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+define i64 @bar() {
+  %1 = call i64 (i64, ...) @foo(i64 0, i64 1, i64 2, double 3.000000e+00,
+                                double 4.000000e+00, i64 5, i64 6,
+                                double 7.000000e+00, i64 8, i64 9, i64 10, i64 11,
+                                double 12.000000e+00, double 13.000000e+00)
+  ret i64 %1
+}
+
+; Save the incoming shadow values from the varargs in the __msan_va_arg_tls
+; array at the offsets equal to those defined by the ABI for the corresponding
+; registers in the register save area, and for the corresponding overflow args
+; in the overflow arg area:
+; - r2@16 is skipped, because it's fixed
+; - r3@24 == 1
+; - r4@32 == 2
+; - f0@128 == 3.0
+; - f2@136 == 4.0
+; - r5@40 == 5
+; - r6@48 == 6
+; - f4@144 == 7.0
+; - overflow@[160::8] == [8, 9, 10, 11]
+; - f6@152 == 12.0
+; - overflow@192 = 13.0
+; Overflow arg area size is therefore 40.
+
+; CHECK-LABEL: @bar
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 24
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 32
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 128
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 136
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 40
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 48
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 144
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 160
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 168
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 176
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 184
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 152
+; CHECK: store {{.*}} @__msan_va_arg_tls {{.*}} 192
+; CHECK: store {{.*}} 40, {{.*}} @__msan_va_arg_overflow_size_tls
+
+; Test that MSan doesn't generate code overflowing __msan_va_arg_tls when too many arguments are
+; passed to a variadic function.
+
+define dso_local i64 @many_args() {
+entry:
+  %ret = call i64 (i64, ...) @sum(i64 120,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1,
+    i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1
+  )
+  ret i64 %ret
+}
+
+; If the size of __msan_va_arg_tls changes the second argument of `add` must also be changed.
+; CHECK-LABEL: @many_args
+; CHECK: i64 add (i64 ptrtoint ([100 x i64]* @__msan_va_arg_tls to i64), i64 792)
+; CHECK-NOT: i64 add (i64 ptrtoint ([100 x i64]* @__msan_va_arg_tls to i64), i64 800)
+
+declare i64 @sum(i64 %n, ...)

llvm/test/Instrumentation/MemorySanitizer/msan_basic.ll

 ; RUN: opt < %s -msan-check-access-address=0 -S -passes='module(msan-module),function(msan)' 2>&1 | FileCheck -allow-deprecated-dag-overlap %s
 ; RUN: opt < %s --passes='module(msan-module),function(msan)' -msan-check-access-address=0 -S | FileCheck -allow-deprecated-dag-overlap %s
 ; RUN: opt < %s -msan-check-access-address=0 -msan-track-origins=1 -S -passes='module(msan-module),function(msan)' 2>&1 | \
 ; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,CHECK-ORIGINS %s
 ; RUN: opt < %s -passes='module(msan-module),function(msan)' -msan-check-access-address=0 -msan-track-origins=1 -S | \
 ; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,CHECK-ORIGINS %s
+; RUN: opt < %s -passes='module(msan-module),function(msan)' -msan-instrumentation-with-call-threshold=0 -msan-track-origins=1 -S | \
+; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK-CALLS %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-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; CHECK: @llvm.global_ctors {{.*}} { i32 0, void ()* @msan.module_ctor, i8* null }
 
 ; Check the presence and the linkage type of __msan_track_origins and
 ; other interface symbols.

 ; CHECK: musttail call i32* @MismatchingMustTailCall
 ; No instrumentation between call and ret.
 ; CHECK-NEXT: bitcast i32* {{.*}} to i8*
 ; CHECK-NEXT: ret i8*
 
 
 ; CHECK-LABEL: define internal void @msan.module_ctor() {
 ; CHECK: call void @__msan_init()
+
+; CHECK-CALLS: declare void @__msan_maybe_warning_1(i8 zeroext, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_store_origin_1(i8 zeroext, i8*, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_warning_2(i16 zeroext, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_store_origin_2(i16 zeroext, i8*, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_warning_4(i32 zeroext, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_store_origin_4(i32 zeroext, i8*, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_warning_8(i64 zeroext, i32 zeroext)
+; CHECK-CALLS: declare void @__msan_maybe_store_origin_8(i64 zeroext, i8*, i32 zeroext)