[NFC] factor isStandardLifetime out of HWASan

[fmayer]

Summary

this is so we can use it for aarch64 stack tagging.

[fmayer]

I don't think it's worth to add a .cpp file for those two functions for now, but would be happy to do that if you think that's better.

[fmayer]

add cpp file

[fmayer]

In D118836#3291527, @fmayer wrote:

I don't think it's worth to add a .cpp file for those two functions for now, but would be happy to do that if you think that's better.

Added to Transform/Utils. PTAL.

[eugenis]

LGTM

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
     PtrUse = &I->getOperandUse(OperandNo);
   }
 
   Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
 
   Value *getPtr() { return PtrUse->get(); }
 };
 
-// For an alloca valid between lifetime markers Start and Ends, call the
-// Callback for all possible exits out of the lifetime in the containing
-// function, which can return from the instructions in RetVec.
-//
-// Returns whether Ends covered all possible exits. If they did not,
-// the caller should remove Ends to ensure that work done at the other
-// exits does not happen outside of the lifetime.
-template <typename F>
-bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
-                          const Instruction *Start,
-                          const SmallVectorImpl<IntrinsicInst *> &Ends,
-                          const SmallVectorImpl<Instruction *> &RetVec,
-                          F Callback) {
-  if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
-    Callback(Ends[0]);
-    return true;
-  }
-  SmallVector<Instruction *, 8> ReachableRetVec;
-  unsigned NumCoveredExits = 0;
-  for (auto *RI : RetVec) {
-    if (!isPotentiallyReachable(Start, RI, nullptr, &DT))
-      continue;
-    ReachableRetVec.push_back(RI);
-    // TODO(fmayer): We don't support diamond shapes, where multiple lifetime
-    // ends together dominate the RI, but none of them does by itself.
-    // Check how often this happens and decide whether to support this here.
-    if (std::any_of(Ends.begin(), Ends.end(),
-                    [&](Instruction *End) { return DT.dominates(End, RI); }))
-      ++NumCoveredExits;
-  }
-  // If there's a mix of covered and non-covered exits, just put the untag
-  // on exits, so we avoid the redundancy of untagging twice.
-  if (NumCoveredExits == ReachableRetVec.size()) {
-    for (auto *End : Ends)
-      Callback(End);
-  } else {
-    for (auto *RI : ReachableRetVec)
-      Callback(RI);
-    // We may have inserted untag outside of the lifetime interval.
-    // Signal the caller to remove the lifetime end call for this alloca.
-    return false;
-  }
-  return true;
-}
-
 // Get AddressSanitizer parameters.
 void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
                                bool IsKasan, uint64_t *ShadowBase,
                                int *MappingScale, bool *OrShadowOffset);
 
 } // namespace llvm
 
 #endif

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

-//===--------- Definition of the AddressSanitizer class ---------*- C++ -*-===//
-//
+//===- MemoryTaggingSupport.h - 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 AddressSanitizer and
-// HWAddressSanitizer.
+// This file declares common infrastructure for HWAddressSanitizer and
+// Aarch64StackTagging.
 //
 //===----------------------------------------------------------------------===//
-#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
-#define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
+#ifndef LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
+#define LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
-
-class InterestingMemoryOperand {
-public:
-  Use *PtrUse;
-  bool IsWrite;
-  Type *OpType;
-  uint64_t TypeSize;
-  MaybeAlign Alignment;
-  // The mask Value, if we're looking at a masked load/store.
-  Value *MaybeMask;
-
-  InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
-                           class Type *OpType, MaybeAlign Alignment,
-                           Value *MaybeMask = nullptr)
-      : IsWrite(IsWrite), OpType(OpType), Alignment(Alignment),
-        MaybeMask(MaybeMask) {
-    const DataLayout &DL = I->getModule()->getDataLayout();
-    TypeSize = DL.getTypeStoreSizeInBits(OpType);
-    PtrUse = &I->getOperandUse(OperandNo);
-  }
-
-  Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
-
-  Value *getPtr() { return PtrUse->get(); }
-};
-
 // For an alloca valid between lifetime markers Start and Ends, call the
 // Callback for all possible exits out of the lifetime in the containing
 // function, which can return from the instructions in RetVec.
 //
 // Returns whether Ends covered all possible exits. If they did not,
 // the caller should remove Ends to ensure that work done at the other
 // exits does not happen outside of the lifetime.
 template <typename F>
for (auto *RI : ReachableRetVec)
       Callback(RI);
     // We may have inserted untag outside of the lifetime interval.
     // Signal the caller to remove the lifetime end call for this alloca.
     return false;
   }
   return true;
 }
 
-// Get AddressSanitizer parameters.
-void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
-                               bool IsKasan, uint64_t *ShadowBase,
-                               int *MappingScale, bool *OrShadowOffset);
-
+bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
+                        const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
+                        const DominatorTree &DT, size_t MaxLifetimes);
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 #include <cassert>
 #include <iterator>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-stack-tagging"
 

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

 #include "llvm/PassRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <sstream>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "hwasan"
 
public:
   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);
-  static bool isStandardLifetime(const AllocaInfo &AllocaInfo,
-                                 const DominatorTree &DT);
   bool instrumentStack(
       bool ShouldDetectUseAfterScope,
       MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
       SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
       DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
       SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
       llvm::function_ref<const DominatorTree &()> GetDT,
       llvm::function_ref<const PostDominatorTree &()> GetPDT);
for (auto *LP : LandingPadVec) {
     IRB.CreateCall(
         HWAsanHandleVfork,
         {readRegister(IRB, (TargetTriple.getArch() == Triple::x86_64) ? "rsp"
                                                                       : "sp")});
   }
   return true;
 }
 
-static bool
-maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
-                            const DominatorTree &DT) {
-  // If we have too many lifetime ends, give up, as the algorithm below is N^2.
-  if (Insts.size() > ClMaxLifetimes)
-    return true;
-  for (size_t I = 0; I < Insts.size(); ++I) {
-    for (size_t J = 0; J < Insts.size(); ++J) {
-      if (I == J)
-        continue;
-      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
-        return true;
-    }
-  }
-  return false;
-}
 
-// static
-bool HWAddressSanitizer::isStandardLifetime(const AllocaInfo &AllocaInfo,
-                                            const DominatorTree &DT) {
-  // An alloca that has exactly one start and end in every possible execution.
-  // If it has multiple ends, they have to be unreachable from each other, so
-  // at most one of them is actually used for each execution of the function.
-  return AllocaInfo.LifetimeStart.size() == 1 &&
-         (AllocaInfo.LifetimeEnd.size() == 1 ||
-          (AllocaInfo.LifetimeEnd.size() > 0 &&
-           !maybeReachableFromEachOther(AllocaInfo.LifetimeEnd, DT)));
-}
 
 bool HWAddressSanitizer::instrumentStack(
     bool ShouldDetectUseAfterScope,
     MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
     SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
     DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
     SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
     llvm::function_ref<const DominatorTree &()> GetDT,
for (auto &KV : AllocasToInstrument) {
     size_t Size = getAllocaSizeInBytes(*AI);
     size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
     auto TagEnd = [&](Instruction *Node) {
       IRB.SetInsertPoint(Node);
       Value *UARTag = getUARTag(IRB, StackTag);
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
-        UnrecognizedLifetimes.empty() && isStandardLifetime(Info, GetDT());
+        UnrecognizedLifetimes.empty() &&
+        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, GetDT(),
+                           ClMaxLifetimes);
     if (ShouldDetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
       tagAlloca(IRB, AI, Tag, Size);
       if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd,
                                 RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();

llvm/lib/Transforms/Utils/CMakeLists.txt

add_llvm_component_library(LLVMTransformUtils
   LoopUnrollRuntime.cpp
   LoopUtils.cpp
   LoopVersioning.cpp
   LowerInvoke.cpp
   LowerMemIntrinsics.cpp
   LowerSwitch.cpp
   MatrixUtils.cpp
   MemoryOpRemark.cpp
+  MemoryTaggingSupport.cpp
   Mem2Reg.cpp
   MetaRenamer.cpp
   ModuleUtils.cpp
   NameAnonGlobals.cpp
   PredicateInfo.cpp
   PromoteMemoryToRegister.cpp
   RelLookupTableConverter.cpp
   ScalarEvolutionExpander.cpp

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"
+
+namespace llvm {
+namespace {
+bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
+                                 const DominatorTree &DT, size_t MaxLifetimes) {
+  // If we have too many lifetime ends, give up, as the algorithm below is N^2.
+  if (Insts.size() > MaxLifetimes)
+    return true;
+  for (size_t I = 0; I < Insts.size(); ++I) {
+    for (size_t J = 0; J < Insts.size(); ++J) {
+      if (I == J)
+        continue;
+      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
+        return true;
+    }
+  }
+  return false;
+}
+} // namespace
+
+bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
+                        const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
+                        const DominatorTree &DT, size_t MaxLifetimes) {
+  // An alloca that has exactly one start and end in every possible execution.
+  // If it has multiple ends, they have to be unreachable from each other, so
+  // at most one of them is actually used for each execution of the function.
+  return LifetimeStart.size() == 1 &&
+         (LifetimeEnd.size() == 1 ||
+          (LifetimeEnd.size() > 0 &&
+           !maybeReachableFromEachOther(LifetimeEnd, DT, MaxLifetimes)));
+}
+} // namespace llvm