[Attributor] Add "free"-based heap2stack deduction

[jdoerfert]

Summary

If there is a unique free of the allocated that has to be reached from
the malloc, we can apply the heap-2-stack transformation even if the
pointer escapes.

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

3647

I need to put this pattern into a helper method in the Iterator class.

[sstefan1]

LGTM

llvm/include/llvm/Analysis/MustExecute.h

struct MustBeExecutedContextExplorer {
   }
 
   /// Return an iterator range to explore the cached context around \p PP.
   llvm::iterator_range<const_iterator> range(const Instruction *PP) const {
     return llvm::make_range(begin(PP), end(PP));
   }
   ///}
 
+  /// Helper to look for \p I in the context of \p PP.
+  ///
+  /// The context is expanded until \p I was found or no more expansion is
+  /// possible.
+  ///
+  /// \returns True, iff \p I was found.
+  bool findInContextOf(const Instruction *I, const Instruction *PP) {
+    auto EIt = begin(PP), EEnd = end(PP);
+    return findInContextOf(I, EIt, EEnd);
+  }
+
+  /// Helper to look for \p I in the context defined by \p EIt and \p EEnd.
+  ///
+  /// The context is expanded until \p I was found or no more expansion is
+  /// possible.
+  ///
+  /// \returns True, iff \p I was found.
+  bool findInContextOf(const Instruction *I, iterator &EIt, iterator &EEnd) {
+    bool Found = EIt.count(I);
+    while (!Found && EIt != EEnd)
+      Found = (++EIt).getCurrentInst() == I;
+    return Found;
+  }
+
   /// Return the next instruction that is guaranteed to be executed after \p PP.
   ///
   /// \param It              The iterator that is used to traverse the must be
   ///                        executed context.
   /// \param PP              The program point for which the next instruction
   ///                        that is guaranteed to execute is determined.
   const Instruction *
   getMustBeExecutedNextInstruction(MustBeExecutedIterator &It,

llvm/lib/Transforms/IPO/Attributor.cpp

ChangeStatus updateImpl(Attributor &A) override {
     if (!CtxI)
       return ChangeStatus::UNCHANGED;
 
     MustBeExecutedContextExplorer &Explorer =
         A.getInfoCache().getMustBeExecutedContextExplorer();
 
     SetVector<const Use *> NextUses;
 
+    auto EIt = Explorer.begin(CtxI), EEnd = Explorer.end(CtxI);
     for (const Use *U : Uses) {
       if (const Instruction *UserI = dyn_cast<Instruction>(U->getUser())) {
-        auto EIt = Explorer.begin(CtxI), EEnd = Explorer.end(CtxI);
-        bool Found = EIt.count(UserI);
-        while (!Found && ++EIt != EEnd)
-          Found = EIt.getCurrentInst() == UserI;
+        bool Found = Explorer.findInContextOf(UserI, EIt, EEnd);
         if (Found && Base::followUse(A, U, UserI))
           for (const Use &Us : UserI->uses())
             NextUses.insert(&Us);
       }
     }
     for (const Use *U : NextUses)
       Uses.insert(U);
 
struct AAHeapToStackImpl : public AAHeapToStack {
 
   ChangeStatus updateImpl(Attributor &A) override;
 };
 
 ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
   const Function *F = getAssociatedFunction();
   const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
 
+  MustBeExecutedContextExplorer &Explorer =
+      A.getInfoCache().getMustBeExecutedContextExplorer();
+
+  auto FreeCheck = [&](Instruction &I) {
+    const auto &Frees = FreesForMalloc.lookup(&I);
+    if (Frees.size() != 1)
+      return false;
+    Instruction *UniqueFree = *Frees.begin();
+    return Explorer.findInContextOf(UniqueFree, I.getNextNode());
+  };
+
   auto UsesCheck = [&](Instruction &I) {
+    bool ValidUsesOnly = true;
+    bool MustUse = true;
+
     SmallPtrSet<const Use *, 8> Visited;
     SmallVector<const Use *, 8> Worklist;
 
     for (Use &U : I.uses())
       Worklist.push_back(&U);
 
     while (!Worklist.empty()) {
       const Use *U = Worklist.pop_back_val();
       if (!Visited.insert(U).second)
         continue;
 
       auto *UserI = U->getUser();
 
       if (isa<LoadInst>(UserI))
         continue;
       if (auto *SI = dyn_cast<StoreInst>(UserI)) {
         if (SI->getValueOperand() == U->get()) {
-          LLVM_DEBUG(dbgs() << "[H2S] escaping store to memory: " << *UserI << "\n");
-          return false;
-        }
+          LLVM_DEBUG(dbgs()
+                     << "[H2S] escaping store to memory: " << *UserI << "\n");
+          ValidUsesOnly = false;
+        } else {
           // A store into the malloc'ed memory is fine.
+        }
         continue;
       }
 
       // NOTE: Right now, if a function that has malloc pointer as an argument
       // frees memory, we assume that the malloc pointer is freed.
 
       // TODO: Add nofree callsite argument attribute to indicate that pointer
       // argument is not freed.
       if (auto *CB = dyn_cast<CallBase>(UserI)) {
         if (!CB->isArgOperand(U))
           continue;
 
         if (CB->isLifetimeStartOrEnd())
           continue;
 
         // Record malloc.
         if (isFreeCall(UserI, TLI)) {
+          if (MustUse) {
             FreesForMalloc[&I].insert(
                 cast<Instruction>(const_cast<User *>(UserI)));
+          } else {
+            LLVM_DEBUG(dbgs() << "[H2S] free potentially on different mallocs: "
+                              << *UserI << "\n");
+            ValidUsesOnly = false;
+          }
           continue;
         }
 
         // If a function does not free memory we are fine
         const auto &NoFreeAA =
             A.getAAFor<AANoFree>(*this, IRPosition::callsite_function(*CB));
 
         unsigned ArgNo = U - CB->arg_begin();
         const auto &NoCaptureAA = A.getAAFor<AANoCapture>(
             *this, IRPosition::callsite_argument(*CB, ArgNo));
 
         if (!NoCaptureAA.isAssumedNoCapture() || !NoFreeAA.isAssumedNoFree()) {
           LLVM_DEBUG(dbgs() << "[H2S] Bad user: " << *UserI << "\n");
-          return false;
+          ValidUsesOnly = false;
         }
         continue;
       }
 
-      if (isa<GetElementPtrInst>(UserI) || isa<BitCastInst>(UserI)) {
+      if (isa<GetElementPtrInst>(UserI) || isa<BitCastInst>(UserI) ||
+          isa<PHINode>(UserI) || isa<SelectInst>(UserI)) {
+        MustUse &= !(isa<PHINode>(UserI) || isa<SelectInst>(UserI));
         for (Use &U : UserI->uses())
           Worklist.push_back(&U);
         continue;
       }
 
-      // Unknown user.
+      // Unknown user for which we can not track uses further (in a way that
+      // makes sense).
       LLVM_DEBUG(dbgs() << "[H2S] Unknown user: " << *UserI << "\n");
-      return false;
+      ValidUsesOnly = false;
     }
-    return true;
+    return ValidUsesOnly;
   };
 
   auto MallocCallocCheck = [&](Instruction &I) {
     if (BadMallocCalls.count(&I))
       return true;
 
     bool IsMalloc = isMallocLikeFn(&I, TLI);
     bool IsCalloc = !IsMalloc && isCallocLikeFn(&I, TLI);
     if (!IsMalloc && !IsCalloc) {
       BadMallocCalls.insert(&I);
       return true;
     }
 
     if (IsMalloc) {
       if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (Size->getValue().sle(MaxHeapToStackSize))
-          if (UsesCheck(I)) {
+          if (UsesCheck(I) || FreeCheck(I)) {
             MallocCalls.insert(&I);
             return true;
           }
     } else if (IsCalloc) {
       bool Overflow = false;
       if (auto *Num = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
           if ((Size->getValue().umul_ov(Num->getValue(), Overflow))
                    .sle(MaxHeapToStackSize))
-            if (!Overflow && UsesCheck(I)) {
+            if (!Overflow && (UsesCheck(I) || FreeCheck(I))) {
               MallocCalls.insert(&I);
               return true;
             }
     }
 
     BadMallocCalls.insert(&I);
     return true;
   };

 }
 
 struct AAHeapToStackFunction final : public AAHeapToStackImpl {
   AAHeapToStackFunction(const IRPosition &IRP) : AAHeapToStackImpl(IRP) {}
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override {
     STATS_DECL(MallocCalls, Function,
-               "Number of MallocCalls converted to allocas");
-    BUILD_STAT_NAME(MallocCalls, Function) += MallocCalls.size();
+               "Number of malloc calls converted to allocas");
+    for (auto *C : MallocCalls)
+      if (!BadMallocCalls.count(C))
+        ++BUILD_STAT_NAME(MallocCalls, Function);
   }
 };
 
 /// -------------------- Memory Behavior Attributes ----------------------------
 /// Includes read-none, read-only, and write-only.
 /// ----------------------------------------------------------------------------
 struct AAMemoryBehaviorImpl : public AAMemoryBehavior {
   AAMemoryBehaviorImpl(const IRPosition &IRP) : AAMemoryBehavior(IRP) {}

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

 ; RUN: opt -passes=attributor --attributor-disable=false -S < %s | FileCheck %s
 
 declare noalias i8* @malloc(i64)
 
 declare void @nocapture_func_frees_pointer(i8* nocapture)
 
 declare void @func_throws(...)
 
 declare void @sync_func(i8* %p)
 
-declare void @sync_will_return(i8* %p) willreturn
+declare void @sync_will_return(i8* %p) willreturn nounwind
 
 declare void @no_sync_func(i8* nocapture %p) nofree nosync willreturn
 
 declare void @nofree_func(i8* nocapture %p) nofree  nosync willreturn
 
 declare void @foo(i32* %p)
 
 declare void @foo_nounw(i32* %p) nounwind nofree
define i32 @test_lifetime() {
   store i32 10, i32* %2
   %3 = load i32, i32* %2
   ; CHECK-NOT: @free(i8* %1)
   tail call void @free(i8* %1)
   ret i32 %3
 }
 
 ; TEST 11 
-; FIXME: should be ok
 
 define void @test11() {
   %1 = tail call noalias i8* @malloc(i64 4)
-  ; CHECK: @malloc(i64 4)
+  ; CHECK: test11
+  ; CHECK-NEXT: alloc
   ; CHECK-NEXT: @sync_will_return(i8* %1)
   tail call void @sync_will_return(i8* %1)
   tail call void @free(i8* %1)
   ret void
 }
 
 ; TEST 12
 define i32 @irreducible_cfg(i32 %0) {
define void @test16a(i8 %v, i8** %P) {
   ret void
 }
 
 define void @test16b(i8 %v, i8** %P) {
   ; CHECK: %1 = tail call noalias i8* @malloc(i64 4)
   %1 = tail call noalias i8* @malloc(i64 4)
   ; CHECK-NEXT: store i8* %1, i8** %P
   store i8* %1, i8** %P
-  ; CHECK-NEXT: @no_sync_func(i8* %1)
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
   tail call void @no_sync_func(i8* %1)
   ; CHECK-NEXT: @free(i8* %1)
   tail call void @free(i8* %1)
   ret void
 }
+
+define void @test16c(i8 %v, i8** %P) {
+  ; CHECK: %1 = alloca
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK-NEXT: store i8* %1, i8** %P
+  store i8* %1, i8** %P
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1) nounwind
+  ; CHECK-NOT: @free
+  tail call void @free(i8* %1)
+  ret void
+}
+
+define void @test16d(i8 %v, i8** %P) {
+  ; CHECK: %1 = tail call noalias i8* @malloc(i64 4)
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK-NEXT: store i8* %1, i8** %P
+  store i8* %1, i8** %P
+  ret void
+}

llvm/test/Transforms/FunctionAttrs/noalias_returned.ll

 declare void @use(i8*)
 define void @test12_1() {
 ; CHECK-LABEL: @test12_1(
 ; CHECK-NEXT:    [[A:%.*]] = alloca i8, align 4
 ; CHECK-NEXT:    [[B:%.*]] = tail call noalias i8* @malloc(i64 4)
 ; CHECK-NEXT:    tail call void @use_nocapture(i8* noalias nonnull align 4 dereferenceable(1) [[A]])
 ; CHECK-NEXT:    tail call void @use_nocapture(i8* noalias nonnull align 4 dereferenceable(1) [[A]])
 ; CHECK-NEXT:    tail call void @use_nocapture(i8* noalias nocapture [[B]])
-; CHECK-NEXT:    tail call void @use_nocapture(i8* noalias [[B]])
+; CHECK-NEXT:    tail call void @use_nocapture(i8* noalias nocapture [[B]])
 ; CHECK-NEXT:    ret void
 ;
   %A = alloca i8, align 4
   %B = tail call noalias i8* @malloc(i64 4)
   tail call void @use_nocapture(i8* %A)
   tail call void @use_nocapture(i8* %A)
   tail call void @use_nocapture(i8* %B)
   tail call void @use_nocapture(i8* %B)
   ret void
 }
 
 define void @test12_2(){
 ; CHECK-LABEL: @test12_2(
 ; CHECK-NEXT:    [[A:%.*]] = tail call noalias i8* @malloc(i64 4)
 ; FIXME: This should be @use_nocapture(i8* noalias [[A]])
 ; CHECK-NEXT:    tail call void @use_nocapture(i8* nocapture [[A]])
-; FIXME: This should be @use_nocapture(i8* noalias [[A]])
-; CHECK-NEXT:    tail call void @use_nocapture(i8* [[A]])
+; FIXME: This should be @use_nocapture(i8* noalias nocapture [[A]])
+; CHECK-NEXT:    tail call void @use_nocapture(i8* nocapture [[A]])
 ; CHECK-NEXT:    tail call void @use(i8* [[A]])
-; CHECK-NEXT:    tail call void @use_nocapture(i8* [[A]])
+; CHECK-NEXT:    tail call void @use_nocapture(i8* nocapture [[A]])
 ; CHECK-NEXT:    ret void
 ;
   %A = tail call noalias i8* @malloc(i64 4)
   tail call void @use_nocapture(i8* %A)
   tail call void @use_nocapture(i8* %A)
   tail call void @use(i8* %A)
   tail call void @use_nocapture(i8* %A)
   ret void
 }
 
 declare void @two_args(i8* nocapture , i8* nocapture)
 define void @test12_3(){
 ; CHECK-LABEL: @test12_3(
   %A = tail call noalias i8* @malloc(i64 4)
-; CHECK: tail call void @two_args(i8* nocapture %A, i8* %A)
+; CHECK: tail call void @two_args(i8* nocapture %A, i8* nocapture %A)
   tail call void @two_args(i8* %A, i8* %A)
   ret void
 }
 
 define void @test12_4(){
 ; CHECK-LABEL: @test12_4(
   %A = tail call noalias i8* @malloc(i64 4)
   %B = tail call noalias i8* @malloc(i64 4)
   %A_0 = getelementptr i8, i8* %A, i64 0
   %A_1 = getelementptr i8, i8* %A, i64 1
   %B_0 = getelementptr i8, i8* %B, i64 0
 
-; CHECK: tail call void @two_args(i8* noalias %A, i8* noalias %B)
+; CHECK: tail call void @two_args(i8* noalias nocapture %A, i8* noalias nocapture %B)
   tail call void @two_args(i8* %A, i8* %B)
 
-; CHECK: tail call void @two_args(i8* %A, i8* nocapture %A_0)
+; CHECK: tail call void @two_args(i8* nocapture %A, i8* nocapture %A_0)
   tail call void @two_args(i8* %A, i8* %A_0)
 
-; CHECK: tail call void @two_args(i8* %A, i8* %A_1)
+; CHECK: tail call void @two_args(i8* nocapture %A, i8* nocapture %A_1)
   tail call void @two_args(i8* %A, i8* %A_1)
 
-; FIXME: This should be @two_args(i8* noalias %A_0, i8* noalias %B_0)
-; CHECK: tail call void @two_args(i8* %A_0, i8* nocapture %B_0)
+; FIXME: This should be @two_args(i8* noalias nocapture %A_0, i8* noalias nocapture %B_0)
+; CHECK: tail call void @two_args(i8* nocapture %A_0, i8* nocapture %B_0)
   tail call void @two_args(i8* %A_0, i8* %B_0)
   ret void
 }