[llvm] [msan] Add 32-bit platforms support (PR #109284)

via llvm-commits llvm-commits at lists.llvm.org
Sat Sep 21 12:53:38 PDT 2024


https://github.com/k-kashapov updated https://github.com/llvm/llvm-project/pull/109284

>From eed180e19264a95f47ee4a7a62e5f3906fbfb44b Mon Sep 17 00:00:00 2001
From: Kamil Kashapov <kashapov at ispras.ru>
Date: Thu, 19 Sep 2024 16:39:29 +0300
Subject: [PATCH 1/2] [msan] Add 32-bit platforms support

---
 .../Instrumentation/MemorySanitizer.cpp       | 485 +++++++++++++-----
 1 file changed, 366 insertions(+), 119 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 07d667434e0710..e18236769ff372 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -411,6 +411,30 @@ static const MemoryMapParams Linux_X86_64_MemoryMapParams = {
     0x100000000000, // OriginBase
 };
 
+// riscv32 Linux
+static const MemoryMapParams Linux_RISCV32_MemoryMapParams = {
+    0x000080000000, // AndMask
+    0,              // XorMask (not used)
+    0,              // ShadowBase (not used)
+    0x000040000000, // OriginBase
+};
+
+// riscv64 Linux
+static const MemoryMapParams Linux_RISCV64_MemoryMapParams = {
+    0,              // AndMask (not used)
+    0x008000000000, // XorMask
+    0,              // ShadowBase (not used)
+    0x002000000000, // OriginBase
+};
+
+// mips32 Linux
+static const MemoryMapParams Linux_MIPS32_MemoryMapParams = {
+    0x000080000000, // AndMask
+    0,              // XorMask (not used)
+    0,              // ShadowBase (not used)
+    0x000040000000, // OriginBase
+};
+
 // mips64 Linux
 static const MemoryMapParams Linux_MIPS64_MemoryMapParams = {
     0,              // AndMask (not used)
@@ -419,6 +443,14 @@ static const MemoryMapParams Linux_MIPS64_MemoryMapParams = {
     0x002000000000, // OriginBase
 };
 
+// ppc32 Linux
+static const MemoryMapParams Linux_PowerPC32_MemoryMapParams = {
+    0x000080000000, // AndMask
+    0,              // XorMask (not used)
+    0,              // ShadowBase (not used)
+    0x000040000000, // OriginBase
+};
+
 // ppc64 Linux
 static const MemoryMapParams Linux_PowerPC64_MemoryMapParams = {
     0xE00000000000, // AndMask
@@ -435,6 +467,14 @@ static const MemoryMapParams Linux_S390X_MemoryMapParams = {
     0x1C0000000000, // OriginBase
 };
 
+// ARM32 Linux
+static const MemoryMapParams Linux_ARM32_MemoryMapParams = {
+    0x000080000000, // AndMask
+    0,              // XorMask (not used)
+    0,              // ShadowBase (not used)
+    0x000040000000, // OriginBase
+};
+
 // aarch64 Linux
 static const MemoryMapParams Linux_AArch64_MemoryMapParams = {
     0,               // AndMask (not used)
@@ -488,13 +528,18 @@ static const PlatformMemoryMapParams Linux_X86_MemoryMapParams = {
     &Linux_X86_64_MemoryMapParams,
 };
 
+static const PlatformMemoryMapParams Linux_RISCV_MemoryMapParams = {
+    &Linux_RISCV32_MemoryMapParams,
+    &Linux_RISCV64_MemoryMapParams,
+};
+
 static const PlatformMemoryMapParams Linux_MIPS_MemoryMapParams = {
-    nullptr,
+    &Linux_MIPS32_MemoryMapParams,
     &Linux_MIPS64_MemoryMapParams,
 };
 
 static const PlatformMemoryMapParams Linux_PowerPC_MemoryMapParams = {
-    nullptr,
+    &Linux_PowerPC32_MemoryMapParams,
     &Linux_PowerPC64_MemoryMapParams,
 };
 
@@ -504,7 +549,7 @@ static const PlatformMemoryMapParams Linux_S390_MemoryMapParams = {
 };
 
 static const PlatformMemoryMapParams Linux_ARM_MemoryMapParams = {
-    nullptr,
+    &Linux_ARM32_MemoryMapParams,
     &Linux_AArch64_MemoryMapParams,
 };
 
@@ -556,10 +601,11 @@ class MemorySanitizer {
   friend struct MemorySanitizerVisitor;
   friend struct VarArgHelperBase;
   friend struct VarArgAMD64Helper;
-  friend struct VarArgMIPS64Helper;
   friend struct VarArgAArch64Helper;
-  friend struct VarArgPowerPC64Helper;
+  friend struct VarArgPowerPCHelper;
   friend struct VarArgSystemZHelper;
+  friend struct VarArgI386Helper;
+  friend struct VarArgGenericHelper;
 
   void initializeModule(Module &M);
   void initializeCallbacks(Module &M, const TargetLibraryInfo &TLI);
@@ -1007,16 +1053,30 @@ void MemorySanitizer::initializeModule(Module &M) {
       break;
     case Triple::Linux:
       switch (TargetTriple.getArch()) {
+      case Triple::x86:
+        MapParams = Linux_X86_MemoryMapParams.bits32;
+        break;
       case Triple::x86_64:
         MapParams = Linux_X86_MemoryMapParams.bits64;
         break;
-      case Triple::x86:
-        MapParams = Linux_X86_MemoryMapParams.bits32;
+      case Triple::riscv32:
+        MapParams = Linux_RISCV_MemoryMapParams.bits32;
+        break;
+      case Triple::riscv64:
+        MapParams = Linux_RISCV_MemoryMapParams.bits64;
+        break;
+      case Triple::mips:
+      case Triple::mipsel:
+        MapParams = Linux_MIPS_MemoryMapParams.bits32;
         break;
       case Triple::mips64:
       case Triple::mips64el:
         MapParams = Linux_MIPS_MemoryMapParams.bits64;
         break;
+      case Triple::ppc:
+      case Triple::ppcle:
+        MapParams = Linux_PowerPC_MemoryMapParams.bits32;
+        break;
       case Triple::ppc64:
       case Triple::ppc64le:
         MapParams = Linux_PowerPC_MemoryMapParams.bits64;
@@ -1024,6 +1084,10 @@ void MemorySanitizer::initializeModule(Module &M) {
       case Triple::systemz:
         MapParams = Linux_S390_MemoryMapParams.bits64;
         break;
+      case Triple::arm:
+      case Triple::armeb:
+        MapParams = Linux_ARM_MemoryMapParams.bits32;
+        break;
       case Triple::aarch64:
       case Triple::aarch64_be:
         MapParams = Linux_ARM_MemoryMapParams.bits64;
@@ -4487,6 +4551,12 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       if (EagerCheck) {
         insertShadowCheck(A, &CB);
         Size = DL.getTypeAllocSize(A->getType());
+        if (ArgOffset + Size > kParamTLSSize)
+          break;
+
+        Value *ArgShadow = getShadow(A);
+        Value *ArgShadowBase = getShadowPtrForArgument(IRB, ArgOffset);
+        IRB.CreateAlignedStore(ArgShadow, ArgShadowBase, kShadowTLSAlignment);
       } else {
         Value *Store = nullptr;
         // Compute the Shadow for arg even if it is ByVal, because
@@ -5037,13 +5107,13 @@ struct VarArgHelperBase : public VarArgHelper {
     Value *TailSize =
         ConstantInt::getSigned(IRB.getInt32Ty(), kParamTLSSize - BaseOffset);
     IRB.CreateMemSet(ShadowBase, ConstantInt::getNullValue(IRB.getInt8Ty()),
-                     TailSize, Align(8));
+                     TailSize, Align(4));
   }
 
   void unpoisonVAListTagForInst(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *VAListTag = I.getArgOperand(0);
-    const Align Alignment = Align(8);
+    const Align Alignment = Align(4);
     auto [ShadowPtr, OriginPtr] = MSV.getShadowOriginPtr(
         VAListTag, IRB, IRB.getInt8Ty(), Alignment, /*isStore*/ true);
     // Unpoison the whole __va_list_tag.
@@ -5290,86 +5360,6 @@ struct VarArgAMD64Helper : public VarArgHelperBase {
   }
 };
 
-/// MIPS64-specific implementation of VarArgHelper.
-/// NOTE: This is also used for LoongArch64.
-struct VarArgMIPS64Helper : public VarArgHelperBase {
-  AllocaInst *VAArgTLSCopy = nullptr;
-  Value *VAArgSize = nullptr;
-
-  VarArgMIPS64Helper(Function &F, MemorySanitizer &MS,
-                     MemorySanitizerVisitor &MSV)
-      : VarArgHelperBase(F, MS, MSV, /*VAListTagSize=*/8) {}
-
-  void visitCallBase(CallBase &CB, IRBuilder<> &IRB) override {
-    unsigned VAArgOffset = 0;
-    const DataLayout &DL = F.getDataLayout();
-    for (Value *A :
-         llvm::drop_begin(CB.args(), CB.getFunctionType()->getNumParams())) {
-      Triple TargetTriple(F.getParent()->getTargetTriple());
-      Value *Base;
-      uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
-      if (TargetTriple.getArch() == Triple::mips64) {
-        // Adjusting the shadow for argument with size < 8 to match the
-        // placement of bits in big endian system
-        if (ArgSize < 8)
-          VAArgOffset += (8 - ArgSize);
-      }
-      Base = getShadowPtrForVAArgument(IRB, VAArgOffset, ArgSize);
-      VAArgOffset += ArgSize;
-      VAArgOffset = alignTo(VAArgOffset, 8);
-      if (!Base)
-        continue;
-      IRB.CreateAlignedStore(MSV.getShadow(A), Base, kShadowTLSAlignment);
-    }
-
-    Constant *TotalVAArgSize = ConstantInt::get(IRB.getInt64Ty(), VAArgOffset);
-    // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
-    // a new class member i.e. it is the total size of all VarArgs.
-    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
-  }
-
-  void finalizeInstrumentation() override {
-    assert(!VAArgSize && !VAArgTLSCopy &&
-           "finalizeInstrumentation called twice");
-    IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
-    Value *CopySize =
-        IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, 0), VAArgSize);
-
-    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.
-      VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
-      VAArgTLSCopy->setAlignment(kShadowTLSAlignment);
-      IRB.CreateMemSet(VAArgTLSCopy, Constant::getNullValue(IRB.getInt8Ty()),
-                       CopySize, kShadowTLSAlignment, false);
-
-      Value *SrcSize = IRB.CreateBinaryIntrinsic(
-          Intrinsic::umin, CopySize,
-          ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
-                       kShadowTLSAlignment, SrcSize);
-    }
-
-    // Instrument va_start.
-    // Copy va_list shadow from the backup copy of the TLS contents.
-    for (CallInst *OrigInst : VAStartInstrumentationList) {
-      NextNodeIRBuilder IRB(OrigInst);
-      Value *VAListTag = OrigInst->getArgOperand(0);
-      Value *RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
-          IRB.CreatePtrToInt(VAListTag, MS.IntptrTy), MS.PtrTy);
-      Value *RegSaveAreaPtr = IRB.CreateLoad(MS.PtrTy, RegSaveAreaPtrPtr);
-      Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
-      const Align Alignment = Align(8);
-      std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
-          MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
-                                 Alignment, /*isStore*/ true);
-      IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
-                       CopySize);
-    }
-  }
-};
-
 /// AArch64-specific implementation of VarArgHelper.
 struct VarArgAArch64Helper : public VarArgHelperBase {
   static const unsigned kAArch64GrArgSize = 64;
@@ -5613,14 +5603,16 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
   }
 };
 
-/// PowerPC64-specific implementation of VarArgHelper.
-struct VarArgPowerPC64Helper : public VarArgHelperBase {
+/// PowerPC-specific implementation of VarArgHelper.
+struct VarArgPowerPCHelper : public VarArgHelperBase {
   AllocaInst *VAArgTLSCopy = nullptr;
   Value *VAArgSize = nullptr;
+  const DataLayout &DL = F.getDataLayout();
+  unsigned IntptrSize = DL.getTypeStoreSize(MS.IntptrTy);
 
-  VarArgPowerPC64Helper(Function &F, MemorySanitizer &MS,
-                        MemorySanitizerVisitor &MSV)
-      : VarArgHelperBase(F, MS, MSV, /*VAListTagSize=*/8) {}
+  VarArgPowerPCHelper(Function &F, MemorySanitizer &MS,
+                        MemorySanitizerVisitor &MSV, unsigned VAListTagSize)
+      : VarArgHelperBase(F, MS, MSV, VAListTagSize) {}
 
   void visitCallBase(CallBase &CB, IRBuilder<> &IRB) override {
     // For PowerPC, we need to deal with alignment of stack arguments -
@@ -5634,12 +5626,17 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
     // Parameter save area starts at 48 bytes from frame pointer for ABIv1,
     // and 32 bytes for ABIv2.  This is usually determined by target
     // endianness, but in theory could be overridden by function attribute.
-    if (TargetTriple.getArch() == Triple::ppc64)
-      VAArgBase = 48;
-    else
-      VAArgBase = 32;
+    if (TargetTriple.isPPC64()) {
+      if (TargetTriple.isPPC64ELFv2ABI()) {
+        VAArgBase = 32;
+      } else {
+        VAArgBase = 48;
+      }
+    } else {
+      // Parameter save area is 8 bytes from frame pointer in PPC32
+      VAArgBase = 8;
+    }
     unsigned VAArgOffset = VAArgBase;
-    const DataLayout &DL = F.getDataLayout();
     for (const auto &[ArgNo, A] : llvm::enumerate(CB.args())) {
       bool IsFixed = ArgNo < CB.getFunctionType()->getNumParams();
       bool IsByVal = CB.paramHasAttr(ArgNo, Attribute::ByVal);
@@ -5697,12 +5694,13 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
         VAArgOffset += ArgSize;
         VAArgOffset = alignTo(VAArgOffset, Align(8));
       }
-      if (IsFixed)
+      if (IsFixed) {
         VAArgBase = VAArgOffset;
+      }
     }
 
     Constant *TotalVAArgSize =
-        ConstantInt::get(IRB.getInt64Ty(), VAArgOffset - VAArgBase);
+        ConstantInt::get(MS.IntptrTy, VAArgOffset - VAArgBase);
     // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
     // a new class member i.e. it is the total size of all VarArgs.
     IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
@@ -5712,7 +5710,7 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
     assert(!VAArgSize && !VAArgTLSCopy &&
            "finalizeInstrumentation called twice");
     IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
     Value *CopySize =
         IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, 0), VAArgSize);
 
@@ -5737,11 +5735,24 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
     for (CallInst *OrigInst : VAStartInstrumentationList) {
       NextNodeIRBuilder IRB(OrigInst);
       Value *VAListTag = OrigInst->getArgOperand(0);
-      Value *RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
+      Value *RegSaveAreaPtrPtr;
+
+      Triple TargetTriple(F.getParent()->getTargetTriple());
+
+      // In PPC32 va_list_tag is a struct, whereas in PPC64 it's a pointer
+      if (TargetTriple.isPPC64()) {
+        RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
           IRB.CreatePtrToInt(VAListTag, MS.IntptrTy), MS.PtrTy);
+      } else {
+        RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
+            IRB.CreateAdd(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
+                          ConstantInt::get(MS.IntptrTy, 8)), MS.PtrTy);
+      }
+
+       
       Value *RegSaveAreaPtr = IRB.CreateLoad(MS.PtrTy, RegSaveAreaPtrPtr);
       Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
-      const Align Alignment = Align(8);
+      const Align Alignment = Align(IntptrSize);
       std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
           MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
                                  Alignment, /*isStore*/ true);
@@ -6032,9 +6043,209 @@ struct VarArgSystemZHelper : public VarArgHelperBase {
   }
 };
 
-// Loongarch64 is not a MIPS, but the current vargs calling convention matches
-// the MIPS.
-using VarArgLoongArch64Helper = VarArgMIPS64Helper;
+/// i386-specific implementation of VarArgHelper.
+struct VarArgI386Helper : public VarArgHelperBase {
+  AllocaInst *VAArgTLSCopy = nullptr;
+  Value *VAArgSize = nullptr;
+  const DataLayout &DL = F.getDataLayout();
+  unsigned IntptrSize = DL.getTypeStoreSize(MS.IntptrTy);
+
+  VarArgI386Helper(Function &F, MemorySanitizer &MS,
+                        MemorySanitizerVisitor &MSV)
+      : VarArgHelperBase(F, MS, MSV, /*VAListTagSize=*/4) {}
+
+  void visitCallBase(CallBase &CB, IRBuilder<> &IRB) override {
+    unsigned VAArgOffset = 0;
+    for (const auto &[ArgNo, A] : llvm::enumerate(CB.args())) {
+      bool IsFixed = ArgNo < CB.getFunctionType()->getNumParams();
+      bool IsByVal = CB.paramHasAttr(ArgNo, Attribute::ByVal);
+      if (IsByVal) {
+        assert(A->getType()->isPointerTy());
+        Type *RealTy = CB.getParamByValType(ArgNo);
+        uint64_t ArgSize = DL.getTypeAllocSize(RealTy);
+        Align ArgAlign = CB.getParamAlign(ArgNo).value_or(Align(IntptrSize));
+        if (ArgAlign < IntptrSize)
+          ArgAlign = Align(IntptrSize);
+        VAArgOffset = alignTo(VAArgOffset, ArgAlign);
+        if (!IsFixed) {
+          Value *Base = getShadowPtrForVAArgument(IRB, VAArgOffset, ArgSize);
+          if (Base) {
+            Value *AShadowPtr, *AOriginPtr;
+            std::tie(AShadowPtr, AOriginPtr) =
+                MSV.getShadowOriginPtr(A, IRB, IRB.getInt8Ty(),
+                                       kShadowTLSAlignment, /*isStore*/ false);
+
+            IRB.CreateMemCpy(Base, kShadowTLSAlignment, AShadowPtr,
+                             kShadowTLSAlignment, ArgSize);
+          }
+          VAArgOffset += alignTo(ArgSize, Align(IntptrSize));
+        }
+      } else {
+        Value *Base;
+        uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
+        Align ArgAlign = Align(IntptrSize);
+        VAArgOffset = alignTo(VAArgOffset, ArgAlign);
+        if (DL.isBigEndian()) {
+          // Adjusting the shadow for argument with size < IntptrSize to match the
+          // placement of bits in big endian system
+          if (ArgSize < IntptrSize)
+            VAArgOffset += (IntptrSize - ArgSize);
+        }
+        if (!IsFixed) {
+          Base = getShadowPtrForVAArgument(IRB, VAArgOffset, ArgSize);
+          if (Base)
+            IRB.CreateStore(MSV.getShadow(A), Base);
+          VAArgOffset += ArgSize;
+          VAArgOffset = alignTo(VAArgOffset, Align(IntptrSize));
+        }
+      }
+    }
+
+    Constant *TotalVAArgSize = ConstantInt::get(MS.IntptrTy, VAArgOffset);
+    // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
+    // a new class member i.e. it is the total size of all VarArgs.
+    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+  }
+
+  void finalizeInstrumentation() override {
+    assert(!VAArgSize && !VAArgTLSCopy &&
+           "finalizeInstrumentation called twice");
+
+    IRBuilder<> IRB(MSV.FnPrologueEnd);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
+    Value *CopySize =
+        IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, 0), VAArgSize);
+
+    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.
+      VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+      VAArgTLSCopy->setAlignment(kShadowTLSAlignment);
+      IRB.CreateMemSet(VAArgTLSCopy, Constant::getNullValue(IRB.getInt8Ty()),
+                       CopySize, kShadowTLSAlignment, false);
+
+      Value *SrcSize = IRB.CreateBinaryIntrinsic(
+          Intrinsic::umin, CopySize,
+          ConstantInt::get(MS.IntptrTy, kParamTLSSize));
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+                       kShadowTLSAlignment, SrcSize);
+    }
+
+    // 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];
+      NextNodeIRBuilder IRB(OrigInst);
+      Value *VAListTag = OrigInst->getArgOperand(0);
+      Type *RegSaveAreaPtrTy = PointerType::getUnqual(*MS.C);
+      Value *RegSaveAreaPtrPtr =
+          IRB.CreateIntToPtr(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
+                             PointerType::get(RegSaveAreaPtrTy, 0));
+      Value *RegSaveAreaPtr =
+          IRB.CreateLoad(RegSaveAreaPtrTy, RegSaveAreaPtrPtr);
+      Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
+      const Align Alignment = Align(IntptrSize);
+      std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
+          MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
+                                 Alignment, /*isStore*/ true);
+      IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
+                       CopySize);
+    }
+  }
+};
+
+/// Implementation of VarArgHelper that is used for ARM32, MIPS, RISCV, LoongArch64.
+struct VarArgGenericHelper : public VarArgHelperBase {
+  AllocaInst *VAArgTLSCopy = nullptr;
+  Value *VAArgSize = nullptr;
+  const DataLayout &DL = F.getDataLayout();
+  unsigned IntptrSize = DL.getTypeStoreSize(MS.IntptrTy);
+
+  VarArgGenericHelper(Function &F, MemorySanitizer &MS,
+                     MemorySanitizerVisitor &MSV, const unsigned VAListTagSize)
+      : VarArgHelperBase(F, MS, MSV, VAListTagSize) {}
+
+  void visitCallBase(CallBase &CB, IRBuilder<> &IRB) override {
+    unsigned VAArgOffset = 0;
+    for (Value *A :
+         llvm::drop_begin(CB.args(), CB.getFunctionType()->getNumParams())) {
+      Triple TargetTriple(F.getParent()->getTargetTriple());
+      Value *Base;
+      uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
+      if (DL.isBigEndian()) {
+        // Adjusting the shadow for argument with size < IntptrSize to match the
+        // placement of bits in big endian system
+        if (ArgSize < IntptrSize)
+          VAArgOffset += (IntptrSize - ArgSize);
+      }
+      Base = getShadowPtrForVAArgument(IRB, VAArgOffset, ArgSize);
+      VAArgOffset += ArgSize;
+      VAArgOffset = alignTo(VAArgOffset, IntptrSize);
+      if (!Base) {
+        continue;
+      }
+      Value *Shadow = MSV.getShadow(A);
+      IRB.CreateStore(Shadow, Base);
+    }
+
+    Constant *TotalVAArgSize = ConstantInt::get(MS.IntptrTy, VAArgOffset);
+    // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
+    // a new class member i.e. it is the total size of all VarArgs.
+    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+  }
+
+  void finalizeInstrumentation() override {
+    assert(!VAArgSize && !VAArgTLSCopy &&
+           "finalizeInstrumentation called twice");
+
+    IRBuilder<> IRB(MSV.FnPrologueEnd);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
+    Value *CopySize =
+        IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, 0), VAArgSize);
+
+    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.
+      VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+      VAArgTLSCopy->setAlignment(kShadowTLSAlignment);
+      IRB.CreateMemSet(VAArgTLSCopy, Constant::getNullValue(IRB.getInt8Ty()),
+                       CopySize, kShadowTLSAlignment, false);
+
+      Value *SrcSize = IRB.CreateBinaryIntrinsic(
+          Intrinsic::umin, CopySize,
+          ConstantInt::get(MS.IntptrTy, kParamTLSSize));
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+                       kShadowTLSAlignment, SrcSize);
+    }
+
+    // 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];
+      NextNodeIRBuilder IRB(OrigInst);
+      Value *VAListTag = OrigInst->getArgOperand(0);
+      Type *RegSaveAreaPtrTy = PointerType::getUnqual(*MS.C);
+      Value *RegSaveAreaPtrPtr =
+          IRB.CreateIntToPtr(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
+                             PointerType::get(RegSaveAreaPtrTy, 0));
+      Value *RegSaveAreaPtr =
+          IRB.CreateLoad(RegSaveAreaPtrTy, RegSaveAreaPtrPtr);
+      Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
+      const Align Alignment = Align(IntptrSize);
+      std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
+          MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
+                                 Alignment, /*isStore*/ true);
+      IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
+                       CopySize);
+    }
+  }
+};
+
+// ARM32, Loongarch64, MIPS and RISCV share the same calling conventions regarding VAArgs.
+using VarArgARM32Helper = VarArgGenericHelper;
+using VarArgRISCVHelper = VarArgGenericHelper;
+using VarArgMIPSHelper = VarArgGenericHelper;
+using VarArgLoongArch64Helper = VarArgGenericHelper;
 
 /// A no-op implementation of VarArgHelper.
 struct VarArgNoOpHelper : public VarArgHelper {
@@ -6057,21 +6268,57 @@ static VarArgHelper *CreateVarArgHelper(Function &Func, MemorySanitizer &Msan,
   // VarArg handling is only implemented on AMD64. False positives are possible
   // on other platforms.
   Triple TargetTriple(Func.getParent()->getTargetTriple());
-  if (TargetTriple.getArch() == Triple::x86_64)
+
+  if (TargetTriple.getArch() == Triple::x86) {
+    return new VarArgI386Helper(Func, Msan, Visitor);
+  }
+
+  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)
+  }
+
+  if (TargetTriple.isARM()) {
+    return new VarArgARM32Helper(Func, Msan, Visitor, /*VAListTagSize=*/4);
+  }
+
+  if (TargetTriple.isAArch64()) {
     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)
+  }
+
+  if (TargetTriple.isSystemZ()) {
     return new VarArgSystemZHelper(Func, Msan, Visitor);
-  else if (TargetTriple.isLoongArch64())
-    return new VarArgLoongArch64Helper(Func, Msan, Visitor);
-  else
-    return new VarArgNoOpHelper(Func, Msan, Visitor);
+  }
+
+  if (TargetTriple.isPPC32()) {
+    // On PowerPC32 VAListTag is a struct {char, char, i16 padding, char *, char *}
+    return new VarArgPowerPCHelper(Func, Msan, Visitor, /*VAListTagSize=*/12);
+  }
+
+  if (TargetTriple.isPPC64()) {
+    return new VarArgPowerPCHelper(Func, Msan, Visitor, /*VAListTagSize=*/8);
+  }
+
+  if (TargetTriple.isRISCV32()) {
+    return new VarArgRISCVHelper(Func, Msan, Visitor, /*VAListTagSize=*/4);
+  }
+
+  if (TargetTriple.isRISCV64()) {
+    return new VarArgRISCVHelper(Func, Msan, Visitor, /*VAListTagSize=*/8);
+  }
+
+  if (TargetTriple.isMIPS32()) {
+    return new VarArgMIPSHelper(Func, Msan, Visitor, /*VAListTagSize=*/4);
+  }
+
+  if (TargetTriple.isMIPS64()) {
+    return new VarArgMIPSHelper(Func, Msan, Visitor, /*VAListTagSize=*/8);
+  }
+
+  if (TargetTriple.isLoongArch64()) {
+    return new VarArgLoongArch64Helper(Func, Msan, Visitor, /*VAListTagSize=*/8);
+  }
+
+  return new VarArgNoOpHelper(Func, Msan, Visitor);
 }
 
 bool MemorySanitizer::sanitizeFunction(Function &F, TargetLibraryInfo &TLI) {

>From 1f46f528e5ccba9a59edb4c345ccf7a606636b41 Mon Sep 17 00:00:00 2001
From: Kamil Kashapov <kashapov at ispras.ru>
Date: Sat, 21 Sep 2024 22:52:38 +0300
Subject: [PATCH 2/2] removed MemoryMapParams, removed EagerCheck shadow copy,
 f fixed alignment

---
 .../Instrumentation/MemorySanitizer.cpp       | 81 ++-----------------
 1 file changed, 7 insertions(+), 74 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index e18236769ff372..9efea41d9d8319 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -411,30 +411,6 @@ static const MemoryMapParams Linux_X86_64_MemoryMapParams = {
     0x100000000000, // OriginBase
 };
 
-// riscv32 Linux
-static const MemoryMapParams Linux_RISCV32_MemoryMapParams = {
-    0x000080000000, // AndMask
-    0,              // XorMask (not used)
-    0,              // ShadowBase (not used)
-    0x000040000000, // OriginBase
-};
-
-// riscv64 Linux
-static const MemoryMapParams Linux_RISCV64_MemoryMapParams = {
-    0,              // AndMask (not used)
-    0x008000000000, // XorMask
-    0,              // ShadowBase (not used)
-    0x002000000000, // OriginBase
-};
-
-// mips32 Linux
-static const MemoryMapParams Linux_MIPS32_MemoryMapParams = {
-    0x000080000000, // AndMask
-    0,              // XorMask (not used)
-    0,              // ShadowBase (not used)
-    0x000040000000, // OriginBase
-};
-
 // mips64 Linux
 static const MemoryMapParams Linux_MIPS64_MemoryMapParams = {
     0,              // AndMask (not used)
@@ -443,14 +419,6 @@ static const MemoryMapParams Linux_MIPS64_MemoryMapParams = {
     0x002000000000, // OriginBase
 };
 
-// ppc32 Linux
-static const MemoryMapParams Linux_PowerPC32_MemoryMapParams = {
-    0x000080000000, // AndMask
-    0,              // XorMask (not used)
-    0,              // ShadowBase (not used)
-    0x000040000000, // OriginBase
-};
-
 // ppc64 Linux
 static const MemoryMapParams Linux_PowerPC64_MemoryMapParams = {
     0xE00000000000, // AndMask
@@ -467,14 +435,6 @@ static const MemoryMapParams Linux_S390X_MemoryMapParams = {
     0x1C0000000000, // OriginBase
 };
 
-// ARM32 Linux
-static const MemoryMapParams Linux_ARM32_MemoryMapParams = {
-    0x000080000000, // AndMask
-    0,              // XorMask (not used)
-    0,              // ShadowBase (not used)
-    0x000040000000, // OriginBase
-};
-
 // aarch64 Linux
 static const MemoryMapParams Linux_AArch64_MemoryMapParams = {
     0,               // AndMask (not used)
@@ -528,18 +488,13 @@ static const PlatformMemoryMapParams Linux_X86_MemoryMapParams = {
     &Linux_X86_64_MemoryMapParams,
 };
 
-static const PlatformMemoryMapParams Linux_RISCV_MemoryMapParams = {
-    &Linux_RISCV32_MemoryMapParams,
-    &Linux_RISCV64_MemoryMapParams,
-};
-
 static const PlatformMemoryMapParams Linux_MIPS_MemoryMapParams = {
-    &Linux_MIPS32_MemoryMapParams,
+    nullptr,
     &Linux_MIPS64_MemoryMapParams,
 };
 
 static const PlatformMemoryMapParams Linux_PowerPC_MemoryMapParams = {
-    &Linux_PowerPC32_MemoryMapParams,
+    nullptr,
     &Linux_PowerPC64_MemoryMapParams,
 };
 
@@ -549,7 +504,7 @@ static const PlatformMemoryMapParams Linux_S390_MemoryMapParams = {
 };
 
 static const PlatformMemoryMapParams Linux_ARM_MemoryMapParams = {
-    &Linux_ARM32_MemoryMapParams,
+    nullptr,
     &Linux_AArch64_MemoryMapParams,
 };
 
@@ -1059,24 +1014,10 @@ void MemorySanitizer::initializeModule(Module &M) {
       case Triple::x86_64:
         MapParams = Linux_X86_MemoryMapParams.bits64;
         break;
-      case Triple::riscv32:
-        MapParams = Linux_RISCV_MemoryMapParams.bits32;
-        break;
-      case Triple::riscv64:
-        MapParams = Linux_RISCV_MemoryMapParams.bits64;
-        break;
-      case Triple::mips:
-      case Triple::mipsel:
-        MapParams = Linux_MIPS_MemoryMapParams.bits32;
-        break;
       case Triple::mips64:
       case Triple::mips64el:
         MapParams = Linux_MIPS_MemoryMapParams.bits64;
         break;
-      case Triple::ppc:
-      case Triple::ppcle:
-        MapParams = Linux_PowerPC_MemoryMapParams.bits32;
-        break;
       case Triple::ppc64:
       case Triple::ppc64le:
         MapParams = Linux_PowerPC_MemoryMapParams.bits64;
@@ -1084,10 +1025,6 @@ void MemorySanitizer::initializeModule(Module &M) {
       case Triple::systemz:
         MapParams = Linux_S390_MemoryMapParams.bits64;
         break;
-      case Triple::arm:
-      case Triple::armeb:
-        MapParams = Linux_ARM_MemoryMapParams.bits32;
-        break;
       case Triple::aarch64:
       case Triple::aarch64_be:
         MapParams = Linux_ARM_MemoryMapParams.bits64;
@@ -4553,10 +4490,6 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
         Size = DL.getTypeAllocSize(A->getType());
         if (ArgOffset + Size > kParamTLSSize)
           break;
-
-        Value *ArgShadow = getShadow(A);
-        Value *ArgShadowBase = getShadowPtrForArgument(IRB, ArgOffset);
-        IRB.CreateAlignedStore(ArgShadow, ArgShadowBase, kShadowTLSAlignment);
       } else {
         Value *Store = nullptr;
         // Compute the Shadow for arg even if it is ByVal, because
@@ -5107,13 +5040,13 @@ struct VarArgHelperBase : public VarArgHelper {
     Value *TailSize =
         ConstantInt::getSigned(IRB.getInt32Ty(), kParamTLSSize - BaseOffset);
     IRB.CreateMemSet(ShadowBase, ConstantInt::getNullValue(IRB.getInt8Ty()),
-                     TailSize, Align(4));
+                     TailSize, Align(8));
   }
 
   void unpoisonVAListTagForInst(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *VAListTag = I.getArgOperand(0);
-    const Align Alignment = Align(4);
+    const Align Alignment = Align(8);
     auto [ShadowPtr, OriginPtr] = MSV.getShadowOriginPtr(
         VAListTag, IRB, IRB.getInt8Ty(), Alignment, /*isStore*/ true);
     // Unpoison the whole __va_list_tag.
@@ -6094,7 +6027,7 @@ struct VarArgI386Helper : public VarArgHelperBase {
         if (!IsFixed) {
           Base = getShadowPtrForVAArgument(IRB, VAArgOffset, ArgSize);
           if (Base)
-            IRB.CreateStore(MSV.getShadow(A), Base);
+            IRB.CreateAlignedStore(MSV.getShadow(A), Base, kShadowTLSAlignment);
           VAArgOffset += ArgSize;
           VAArgOffset = alignTo(VAArgOffset, Align(IntptrSize));
         }
@@ -6185,7 +6118,7 @@ struct VarArgGenericHelper : public VarArgHelperBase {
         continue;
       }
       Value *Shadow = MSV.getShadow(A);
-      IRB.CreateStore(Shadow, Base);
+      IRB.CreateAlignedStore(Shadow, Base, kShadowTLSAlignment);
     }
 
     Constant *TotalVAArgSize = ConstantInt::get(MS.IntptrTy, VAArgOffset);



More information about the llvm-commits mailing list