[MTE] [HWASan] unify isInterestingAlloca

[eugenis]

LGTM

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 #define LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Support/Alignment.h"
 
 namespace llvm {
 class DominatorTree;
 class DbgVariableIntrinsic;
 class IntrinsicInst;
 class PostDominatorTree;
 class AllocaInst;
struct StackInfo {
   MapVector<AllocaInst *, AllocaInfo> AllocasToInstrument;
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
   SmallVector<Instruction *, 8> RetVec;
   bool CallsReturnTwice = false;
 };
 
 class StackInfoBuilder {
 public:
-  StackInfoBuilder(std::function<bool(const AllocaInst &)> IsInterestingAlloca)
-      : IsInterestingAlloca(IsInterestingAlloca) {}
+  StackInfoBuilder(const StackSafetyGlobalInfo *SSI) : SSI(SSI) {}
 
   void visit(Instruction &Inst);
+  bool isInterestingAlloca(const AllocaInst &AI);
   StackInfo &get() { return Info; };
 
 private:
   StackInfo Info;
-  std::function<bool(const AllocaInst &)> IsInterestingAlloca;
+  const StackSafetyGlobalInfo *SSI;
 };
 
 uint64_t getAllocaSizeInBytes(const AllocaInst &AI);
 void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Align);
 
 } // namespace memtag
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

public:
   AArch64StackTagging(bool IsOptNone = false)
       : FunctionPass(ID),
         MergeInit(ClMergeInit.getNumOccurrences() ? ClMergeInit : !IsOptNone),
         UseStackSafety(ClUseStackSafety.getNumOccurrences() ? ClUseStackSafety
                                                             : !IsOptNone) {
     initializeAArch64StackTaggingPass(*PassRegistry::getPassRegistry());
   }
 
-  bool isInterestingAlloca(const AllocaInst &AI);
-
   void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr,
                  uint64_t Size);
   void untagAlloca(AllocaInst *AI, Instruction *InsertBefore, uint64_t Size);
 
   Instruction *collectInitializers(Instruction *StartInst, Value *StartPtr,
                                    uint64_t Size, InitializerBuilder &IB);
 
   Instruction *insertBaseTaggedPointer(
if (StoreInst *NextStore = dyn_cast<StoreInst>(BI)) {
       if (!IB.addMemSet(*Offset, MSI))
         break;
       LastInst = MSI;
     }
   }
   return LastInst;
 }
 
-bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) {
-  // FIXME: support dynamic allocas
-  bool IsInteresting =
-      AI.getAllocatedType()->isSized() && AI.isStaticAlloca() &&
-      // alloca() may be called with 0 size, ignore it.
-      *AI.getAllocationSizeInBits(*DL) > 0 &&
-      // inalloca allocas are not treated as static, and we don't want
-      // dynamic alloca instrumentation for them as well.
-      !AI.isUsedWithInAlloca() &&
-      // swifterror allocas are register promoted by ISel
-      !AI.isSwiftError() &&
-      // safe allocas are not interesting
-      !(SSI && SSI->isSafe(AI));
-  return IsInteresting;
-}
-
 void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore,
                                     Value *Ptr, uint64_t Size) {
   auto SetTagZeroFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag_zero);
   auto StgpFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_stgp);
 
   InitializerBuilder IB(Size, DL, Ptr, SetTagFunc, SetTagZeroFunc, StgpFunc);
bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
   if (UseStackSafety)
     SSI = &getAnalysis<StackSafetyGlobalInfoWrapperPass>().getResult();
   F = &Fn;
   DL = &Fn.getParent()->getDataLayout();
   if (MergeInit)
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
-  memtag::StackInfoBuilder SIB(
-      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+  memtag::StackInfoBuilder SIB(SSI);
   for (Instruction &I : instructions(F))
     SIB.visit(I);
   memtag::StackInfo &SInfo = SIB.get();
 
   if (SInfo.AllocasToInstrument.empty())
     return false;
 
   std::unique_ptr<DominatorTree> DeleteDT;
bool AArch64StackTagging::runOnFunction(Function &Fn) {
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
   Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
 
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
-    assert(Info.AI && isInterestingAlloca(*Info.AI));
+    assert(Info.AI && SIB.isInterestingAlloca(*Info.AI));
     TrackingVH<Instruction> OldAI = Info.AI;
     memtag::alignAndPadAlloca(Info, kTagGranuleSize);
     AllocaInst *AI = Info.AI;
     int Tag = NextTag;
     NextTag = (NextTag + 1) % 16;
     // Replace alloca with tagp(alloca).
     IRBuilder<> IRB(Info.AI->getNextNode());
     Function *TagP = Intrinsic::getDeclaration(

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

void instrumentMemAccessInline(Value *Ptr, bool IsWrite,
                                  Instruction *InsertBefore);
   bool ignoreMemIntrinsic(MemIntrinsic *MI);
   void instrumentMemIntrinsic(MemIntrinsic *MI);
   bool instrumentMemAccess(InterestingMemoryOperand &O);
   bool ignoreAccess(Instruction *Inst, Value *Ptr);
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
-  bool isInterestingAlloca(const AllocaInst &AI);
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
   bool instrumentStack(memtag::StackInfo &Info, Value *StackTag,
                        const DominatorTree &DT, const PostDominatorTree &PDT,
                        const LoopInfo &LI);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
for (auto &KV : SInfo.AllocasToInstrument) {
     }
     memtag::alignAndPadAlloca(Info, Align(Mapping.getObjectAlignment()));
   }
   for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
   return true;
 }
 
-bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
-  return (AI.getAllocatedType()->isSized() &&
-          // FIXME: instrument dynamic allocas, too
-          AI.isStaticAlloca() &&
-          // alloca() may be called with 0 size, ignore it.
-          memtag::getAllocaSizeInBytes(AI) > 0 &&
-          // We are only interested in allocas not promotable to registers.
-          // Promotable allocas are common under -O0.
-          !isAllocaPromotable(&AI) &&
-          // inalloca allocas are not treated as static, and we don't want
-          // dynamic alloca instrumentation for them as well.
-          !AI.isUsedWithInAlloca() &&
-          // swifterror allocas are register promoted by ISel
-          !AI.isSwiftError()) &&
-         // safe allocas are not interesting
-         !(SSI && SSI->isSafe(AI));
-}
-
 bool HWAddressSanitizer::sanitizeFunction(Function &F,
                                           FunctionAnalysisManager &FAM) {
   if (&F == HwasanCtorFunction)
     return false;
 
   if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
     return false;
 
   LLVM_DEBUG(dbgs() << "Function: " << F.getName() << "\n");
 
   SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
   SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
   SmallVector<Instruction *, 8> LandingPadVec;
 
-  memtag::StackInfoBuilder SIB(
-      [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+  memtag::StackInfoBuilder SIB(SSI);
   for (auto &Inst : instructions(F)) {
     if (InstrumentStack) {
       SIB.visit(Inst);
     }
 
     if (InstrumentLandingPads && isa<LandingPadInst>(Inst))
       LandingPadVec.push_back(&Inst);
 

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

 //== MemoryTaggingSupport.cpp - helpers for memory tagging implementations ===//
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares common infrastructure for HWAddressSanitizer and
 // Aarch64StackTagging.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 
 namespace llvm {
 namespace memtag {
 namespace {
 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
                                  const DominatorTree *DT, const LoopInfo *LI,
                                  size_t MaxLifetimes) {
   // If we have too many lifetime ends, give up, as the algorithm below is N^2.

 
 void StackInfoBuilder::visit(Instruction &Inst) {
   if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
     if (CI->canReturnTwice()) {
       Info.CallsReturnTwice = true;
     }
   }
   if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
-    if (IsInterestingAlloca(*AI)) {
+    if (isInterestingAlloca(*AI)) {
       Info.AllocasToInstrument[AI].AI = AI;
     }
     return;
   }
   auto *II = dyn_cast<IntrinsicInst>(&Inst);
   if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
              II->getIntrinsicID() == Intrinsic::lifetime_end)) {
     AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
     if (!AI) {
       Info.UnrecognizedLifetimes.push_back(&Inst);
       return;
     }
-    if (!IsInterestingAlloca(*AI))
+    if (!isInterestingAlloca(*AI))
       return;
     if (II->getIntrinsicID() == Intrinsic::lifetime_start)
       Info.AllocasToInstrument[AI].LifetimeStart.push_back(II);
     else
       Info.AllocasToInstrument[AI].LifetimeEnd.push_back(II);
     return;
   }
   if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
     for (Value *V : DVI->location_ops()) {
       if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
-        if (!IsInterestingAlloca(*AI))
+        if (!isInterestingAlloca(*AI))
           continue;
         AllocaInfo &AInfo = Info.AllocasToInstrument[AI];
         auto &DVIVec = AInfo.DbgVariableIntrinsics;
         if (DVIVec.empty() || DVIVec.back() != DVI)
           DVIVec.push_back(DVI);
       }
     }
   }
   Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
   if (ExitUntag)
     Info.RetVec.push_back(ExitUntag);
 }
 
+bool StackInfoBuilder::isInterestingAlloca(const AllocaInst &AI) {
+  return (AI.getAllocatedType()->isSized() &&
+          // FIXME: instrument dynamic allocas, too
+          AI.isStaticAlloca() &&
+          // alloca() may be called with 0 size, ignore it.
+          memtag::getAllocaSizeInBytes(AI) > 0 &&
+          // We are only interested in allocas not promotable to registers.
+          // Promotable allocas are common under -O0.
+          !isAllocaPromotable(&AI) &&
+          // inalloca allocas are not treated as static, and we don't want
+          // dynamic alloca instrumentation for them as well.
+          !AI.isUsedWithInAlloca() &&
+          // swifterror allocas are register promoted by ISel
+          !AI.isSwiftError()) &&
+         // safe allocas are not interesting
+         !(SSI && SSI->isSafe(AI));
+}
+
 uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
   auto DL = AI.getModule()->getDataLayout();
   return *AI.getAllocationSizeInBits(DL) / 8;
 }
 
 void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) {
   const Align NewAlignment = std::max(Info.AI->getAlign(), Alignment);
   Info.AI->setAlignment(NewAlignment);

llvm/test/CodeGen/AArch64/stack-tagging-split-lifetime.ll

 ; RUN: opt -S -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 %s -o - | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-arm-unknown-eabi"
 
+declare void @use8(i8*)
+
 define  void @f(i1 %cond) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
 next0:
 ; CHECK-LABEL: next0:
 ; CHECK: call void @llvm.aarch64.settag
   call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a)
   br label %exit0
; CHECK-NOT: call void @llvm.aarch64.settag
   ret void
 }
 
 define  void @diamond(i1 %cond) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
 next0:
 ; CHECK-LABEL: next0:
 ; CHECK: call void @llvm.aarch64.settag
   call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a)
   br label %exit1
; CHECK-NOT: call void @llvm.aarch64.settag
   ret void
 }
 
 define  void @diamond_nocover(i1 %cond) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %next1
 
 next0:
 ; CHECK-LABEL: next0:
 ; CHECK-NOT: llvm.lifetime.end
   call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a)
   br label %exit1
; CHECK: call void @llvm.aarch64.settag
   ret void
 }
 
 define  void @diamond3(i1 %cond, i1 %cond1) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %start1
 
 start1:
   br i1 %cond1, label %next1, label %next2
 
 next0:
 ; CHECK-LABEL: next0:
; CHECK-NOT: call void @llvm.aarch64.settag
   ret void
 }
 
 define  void @diamond3_nocover(i1 %cond, i1 %cond1) local_unnamed_addr sanitize_memtag {
 start:
 ; CHECK-LABEL: start:
   %a = alloca i8, i32 48, align 8
   call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+  call void @use8(i8* %a)
 ; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
   br i1 %cond, label %next0, label %start1
 
 start1:
   br i1 %cond1, label %next1, label %next2
 
 next0:
 ; CHECK-LABEL: next0:

llvm/test/CodeGen/AArch64/stack-tagging.ll

 ; RUN: opt < %s -aarch64-stack-tagging -S -o - | FileCheck %s --check-prefixes=CHECK,SSI
-; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -o - | FileCheck %s --check-prefixes=CHECK,NOSSI
+; RUN: opt < %s -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 -S -o - | FileCheck %s --check-prefixes=CHECK
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-android"
 
 declare void @use8(i8*)
 declare void @use32(i32*)
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)

 ; CHECK:  call { i32, [12 x i8] }* @llvm.aarch64.tagp.{{.*}}({ i32, [12 x i8] }* {{.*}}, i64 0)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  alloca { i8, [15 x i8] }, align 16
 ; CHECK:  call { i8, [15 x i8] }* @llvm.aarch64.tagp.{{.*}}({ i8, [15 x i8] }* {{.*}}, i64 1)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  alloca { [11 x i32], [4 x i8] }, align 16
 ; CHECK:  call { [11 x i32], [4 x i8] }* @llvm.aarch64.tagp.{{.*}}({ [11 x i32], [4 x i8] }* {{.*}}, i64 2)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)
-; SSI:    alloca i32, align 4
-; NOSSI:  alloca { i32, [12 x i8] }, align 16
-; NOSSI: @llvm.aarch64.tagp.
-; NOSSI: call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
+; CHECK:  alloca i32, align 4
 ; SSI-NOT: @llvm.aarch64.tagp
 ; SSI-NOT: @llvm.aarch64.settag
 
 ; CHECK:  call void @use32(
 ; CHECK:  call void @use8(
 ; CHECK:  call void @use32(
 
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)
-; NOSSI:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK-NEXT:  ret void
 
 
 define void @Scope(i32 %b) sanitize_memtag {
 entry:
   %x = alloca i32, align 4
   %tobool = icmp eq i32 %b, 0
   br i1 %tobool, label %if.end, label %if.then