[Attributor] Heap-To-Stack Conversion

[sstefan1]

Summary

D53362 gives a prototype heap-to-stack conversion pass. With addition of new attributes in the attributor, this can now be revisted and improved. This will place it in the Attributor to make it easier to use new attributes (eg. nofree, nosync, willreturn, etc.) and other attributor features.

I'm starting with some test cases (more will be added along the way). For now they are in FunctionAttrs, as most other attributor tests, but that can be changed.

[sstefan1]

chenge function names

[xbolva00]

Negative tests:
Malloc in loop body
Irreducible cfg

?

[jdoerfert]

Top-level functions declarations should have attributes to make them as restricted as needed, e.g. sync_func should probably be nofree or renamed to sync_free_func.

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

49

What is the difference between Test 2 and 3?

[sstefan1]

In D65408#1605351, @jdoerfert wrote:

Top-level functions declarations should have attributes to make them as restricted as needed, e.g. sync_func should probably be nofree or renamed to sync_free_func.

Yes. I'll refine them along the way.

In D65408#1605240, @xbolva00 wrote:

Negative tests:
Malloc in loop body
Irreducible cfg

?

Thanks for suggestions. These will also be added. I will probably put some code today, so I can actually test this on something.

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

49

Nothing :). I meant to put call to nosync function here. I'll correct this.

[sstefan1]

Providing an implementation for the conversion. This is few commits behind fromt the latest Attributor, will be rebased.

[sstefan1]

• minor fix

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

2819

Not needed, see my comment further below.

2824

Not needed, we can use the InfoCache to look only at calls.

2830

Let's collect all valid h2s candidate malloc calls and then only "materialize" the ones we deem profitable. Knowing a malloc could be converted to the stack is already a nice piece information.

3047

If you want to start with a small patch, just check the function for willreturn, nounwind, and nofree.

3074

I dislike this exploration. We should not need to traverse all paths all the time. (This does not track visited blocks btw. so it is probably an endless loop if there are loops in the function).

We should (later) follow the uses of the malloc only. The logic I have in mind is as follows (IMHO much easier than this (=the old logic)):

If all (transitive) users of a malloc are either nocapture uses or free uses, we can do heap2stack. In fact, we should probably generalize the capture AA to collect capturing uses instead so we can analyze them later, e.g., here.

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

18

I think this is a bug, see below.

I'll create a patch to *not* annotate free with nocapture.

26

I'd argue, "freeing" a pointer is "capturing" by nature so this combination "nocapture" + "free" cannot exist. If it doesn't, "nocapture" is marked and implies "nofree" on the pointer. Thus, we can actually allocate %1 on the stack here.

[sstefan1]

Taking a different approach using nocapture and nofree.

[sstefan1]

remove llvm_debug

[jdoerfert]

We need tests where stuff is read/written as well.

llvm/lib/Transforms/IPO/Attributor.cpp

2719

Mallocs can become live later so we shouldn't remove them but only ignore them for the moment. In fact, the checkForAllCallLikeInstructions call will already filter "currently assumed dead" ones out so you have to invoke it in every update or collect all calls through the InfoCache explicitly in the initialize method.

2724

This is quadratic in the number of malloc calls.

If you want to exclude malloc calls from the tests below do:
if (MallocCalls.count(&I)) return true;

2729

No need to check, just put it in FreesForMalloc (though I'm unclear why we need it at the moment)

2773

I was hoping we follow and check uses:

UsesCheck = (Instruction &I) {
  SmallPtrSet<const Use *, 8> Visisted;
  SmallVector<const Use *, 8> Worklist;
  Worklist.append(I.use_begin(), I.use_end());
  while (!Worklist.empty()) {
    const Use *U = Worklist.pop_back_val();
    if (!Visited.insert(U).second) continue;
    const auto *UserI = U->getUser();
    if (isa<LoadInst>(UserI) || isa<StoreInst>(UserI))
     continue;
    if (auto *CB = dyn_cast<CallBase>(UserI)) {
      // TODO: Check use once we have nofree on call site args.
      const auto &NoFreeAA = A.getAAFor<AANoFree>(..., IRPosition::call_site(CB));
      if (!NoFreeAA.isAssumedNoFree()) return false;
      // same for nocapture
      // if both are good, continue.
    }
    if (isa<GetElementOperator>(UserI) || isa<BitCastInst>(UserI)) {
      // Add users to worklist and continue
    }
    // unknown user
    return false;
  }
  return true;
}

MalloCheck = (Instruction &I) {
  if (!isaMalloc(I)) return true;
  if (BadMallocs.count(&I)) return true;
  // assumed h2s malloc I
  if (UsesCheck(I))
    Mallocs.insert(&I);
  else
    BadMallocs.insert(&I);
  return true;
}

checkForAllCallLikeInstructions(MallocCheck)

2787

Is this "Number of mallocs converted to allocs"?

2792

I doubt it makes sense to have a H2S call site AA. We should probably only allow the function one.

3260

Still use the Whitelist and group it with the other function AAs please

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

41

We need a test, sync_will_return before free which should (later) be OK.

121

This should be h2s as the free is dead and there are no users that are nofree.

[sstefan1]

addressing comments

[sstefan1]

• remove mistake

[sstefan1]

llvm/lib/Transforms/IPO/Attributor.cpp

2729

The reason for this is if for example a malloc becomes 'bad' then all the frees that use it can be easily disregarded.

[sstefan1]

• Small corrections.

[jdoerfert]

There are some leftover comments but this is almost ready.

llvm/lib/Transforms/IPO/Attributor.cpp

129

I'd enable this by default but for a smaller size, let's say 128?

2652

AI = new BitCastInst(AI, MallocCall->getType(), "malloc_bc", AI->getNextNode());

might do the trick.

2673

remove the const here then you can avoid the const casts

2719

If no-free is set nocaputre is not checked anymore. We need to check for nofree and nocapture.

2724

If this is not the arg operand we should give up. (we may allow it at the callee position but not in the operand bundle)

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

59

outdated comment.

69

maybe elaborate why h2s is still OK.

104

You want branch %6 or you would have a double free (not necessarily a problem for h2s but still).

122

leftover comment.

128

We need this test again with an nounwind @foo which should allow h2s.

[xbolva00]

In D65408#1605240, @xbolva00 wrote:

Negative tests:
Malloc in loop body
Irreducible cfg

?

Still missing.

[xbolva00]

+ support for calloc shouldn’t be very hard to add.

[lebedev.ri]

In D65408#1653669, @jdoerfert wrote:

There are some leftover comments but this is almost ready.

I'd like to see tests with https://llvm.org/docs/LangRef.html#llvm-lifetime-start-intrinsic

[sstefan1]

@xbolva00 I'm sorry this has dragged on a bit.

I will address all the comments, probably tomorrow.

[sstefan1]

addressing comments

[sstefan1]

small fix

[sstefan1]

fix diff

[jdoerfert]

Is this in a state to be reviewed or do you plan to go over it again?

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

252

check lines missing.

[jdoerfert]

Can you rebase and remove the artifacts in the diff? I also provided final comments. I think once they are fixed it is fine. Feel free to commit a fixed version.

llvm/lib/Transforms/IPO/Attributor.cpp

2664

Leftover comment and braces around single stmt.

2745

this has to lead to invalidation. use in a operand bundle could free even if the function is nofree. Maybe move the check to the beginning of the CallBase conditional.

2765

Allow lifetime.start/end as users here as well. That should unbreak one of the test cases.unbreak

2780

You have to put them into the "Bad" instruction set here. Also, move the "is bad" check to the beginning.
There is a test missing for a malloc and a calloc that have to large sizes. (calloc only the multiplication should be too large).
We need to check that the multiplication does not overflow. APInt has a method to do overflow checked multiplication.

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

2

When we run this with the old PM, what will happen? No h2s? We should do it and make sure it doesn't crash etc.

203

This should be fine. Also a lifetime.end user is fine.

[sstefan1]

llvm/lib/Transforms/IPO/Attributor.cpp

2745

It makes sense to me to move the check to the beginning.

use in a operand bundle could free even if the function is nofree

I'm not I understand this, though.

2780

Also, move the "is bad" check to the beginning.

If you mean in initialize, then I'd say this is a nicer approach.

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

2

Yes, no h2s and doesn't crash.

[xbolva00]

I would prefer to use update_test_checks.py instead of hand written checks.

Quite questionable is “stackification” of mallocs in loops, we should avoid it I think. @efriedma, what do you think?

[xbolva00]

IsMallocLikeFn returns true for C++ operator new which does throw on an allocation failure.

So probably you stackify:

int *p = new int;
free(p);

It is a stupid code and buggy, but please add it as test case and the following fix:

For now, just check if (isMallocLike && !isOpNewLike).

I think it is worth to add support for c++’s new - delete pairs in the future..

[lebedev.ri]

In D65408#1670702, @xbolva00 wrote:

Quite questionable is “stackification” of mallocs in loops, we should avoid it I think. @efriedma, what do you think?

Please can you be more specific?

In D65408#1670703, @xbolva00 wrote:

IsMallocLikeFn returns true for C++ operator new which does throw on an allocation failure.

Likewise, please can you be more specific in your messages? This doesn't actually explain what the problem is.

[xbolva00]

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

3229 ↗

(On Diff #220255)

Check for overflow is missing

3263 ↗

(On Diff #220255)

Builder.CreateMemset

[xbolva00]

I think the current implementation needs some important fixes to avoid miscompilation/overflow, can you make it off by default, fix issues, and enable it again?

Thanks.

[xbolva00]

In D65408#1670705, @lebedev.ri wrote:

In D65408#1670702, @xbolva00 wrote:

Quite questionable is “stackification” of mallocs in loops, we should avoid it I think. @efriedma, what do you think?

Please can you be more specific?

Yes :)
Heap to stack idea started some time ago and @efriedma raised some concerns: https://reviews.llvm.org/D47345#1112573.
Imagine

for (int i = 0; i < BIG_NUM; ++i) {

void *p = malloc(128);
....
free(p);

}

->

for (int i = 0; i < BIG_NUM; ++i) {

void *p = alloca(128);
....

}

Nice stack overflow :)

And yes, we shouldn't do this on irreducible cfgs (also noted by Eli).

In D65408#1670703, @xbolva00 wrote:

IsMallocLikeFn returns true for C++ operator new which does throw on an allocation failure.

Likewise, please can you be more specific in your messages? This doesn't actually explain what the problem is.

%1 = tail call dereferenceable(4) i8* @_Znwm(i64 4)
...some code...
tail call void @free(i8* nonnull %1)

is probably converted to alloca. I have no strong opinion what should be done here, ideally warning! :) Anyway, I would like to see some tests with C++'s new and various combinations (there are many variants of 'new', 'delete').

[sstefan1]

In D65408#1670708, @xbolva00 wrote:

I think the current implementation needs some important fixes to avoid miscompilation/overflow, can you make it off by default, fix issues, and enable it again?

Thanks.

Attributor is turned if by dedault right now, so h2s is also turned off. I will revisit this again this week and try to address most of the concerns you raised in a new patch.

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

3229 ↗

(On Diff #220255)

If there is an overflow, this will not be executed.

[xbolva00]

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

3229 ↗

(On Diff #220255)

Ok, great

[sstefan1]

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

3363 ↗

(On Diff #220255)

This was supposed to be if(Overflow). I'll fix this.

[jdoerfert]

In D65408#1670708, @xbolva00 wrote:

I think the current implementation needs some important fixes to avoid miscompilation/overflow, can you make it off by default, fix issues, and enable it again?

The Attributor is off by default (even though we should now pass LLVM-TS + SPEC2006 tests!!).

In D65408#1670703, @xbolva00 wrote:

IsMallocLikeFn returns true for C++ operator new which does throw on an allocation failure.

In D65408#1670713, @xbolva00 wrote:

[...] is probably converted to alloca. I have no strong opinion what should be done here, ideally warning! :) Anyway, I would like to see some tests with C++'s new and various combinations (there are many variants of 'new', 'delete').

I fail to see the problem in converting a possible throwing new to a non-throwing alloca.
If you think there is one could you please try to explain why that should not be valid?

In D65408#1670703, @xbolva00 wrote:

I think it is worth to add support for c++’s new - delete pairs in the future..

We can add this (the delete part) when we actually use free/delete to reason about h2s validity.

In D65408#1670713, @xbolva00 wrote:

In D65408#1670705, @lebedev.ri wrote:

In D65408#1670702, @xbolva00 wrote:

Quite questionable is “stackification” of mallocs in loops, we should avoid it I think. @efriedma, what do you think?

Heap to stack idea started some time ago and @efriedma raised some concerns: https://reviews.llvm.org/D47345#1112573.

From a pure standard point of view, I would argue the current implementation is not wrong in doing h2s for mallocs in loops.
I do however not intent to turn it on and let people deal with the aftermath. As mentioned by @sstefan1, this is not on by default.
My plan was to get the logic in, add the "must-be-freed"-based logic that @hfinkel used in the original version, and then look into heuristics.
The reason is that the "must-be-freed" actually allows different placement of the alloca and we want to see what is out there wrt. sizes etc.
Does that make sense and does my plan sound OK to you?

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

3363 ↗

(On Diff #220255)

Add a test for this please.

[xbolva00]

If you think there is one could you please try to explain why that should not be valid?

Ok, now I think it should be okay.

Does that make sense and does my plan sound OK to you?

Ok.

llvm/include/llvm/Transforms/IPO/Attributor.h

struct Attributor {
   /// \param TLIGetter helper function to get TargetLibraryInfo Analysis result.
   ///
   /// Note that abstract attribute instances are generally created even if the
   /// IR already contains the information they would deduce. The most important
   /// reason for this is the single interface, the one of the abstract attribute
   /// instance, which can be queried without the need to look at the IR in
   /// various places.
   void identifyDefaultAbstractAttributes(
-      Function &F, std::function<TargetLibraryInfo &(Function &)> &TLIGetter,
+      Function &F, std::function<TargetLibraryInfo *(Function &)> &TLIGetter,
       DenseSet<const char *> *Whitelist = nullptr);
 
+  /// Record that \p I is deleted after information was manifested.
+  void deleteAfterManifest(Instruction &I) { ToBeDeletedInsts.insert(&I); }
+
+  /// Record that \p BB is deleted after information was manifested.
+  void deleteAfterManifest(BasicBlock &BB) { ToBeDeletedBlocks.insert(&BB); }
+
+  /// Record that \p F is deleted after information was manifested.
+  void deleteAfterManifest(Function &F) { ToBeDeletedFunctions.insert(&F); }
+
   /// Return true if \p AA (or its context instruction) is assumed dead.
   ///
   /// If \p LivenessAA is not provided it is queried.
   bool isAssumedDead(const AbstractAttribute &AA, const AAIsDead *LivenessAA);
 
   /// Check \p Pred on all function call sites.
   ///
   /// This method will evaluate \p Pred on call sites and return
private:
   ///{
   using QueryMapTy =
       DenseMap<AbstractAttribute *, SetVector<AbstractAttribute *>>;
   QueryMapTy QueryMap;
   ///}
 
   /// The information cache that holds pre-processed (LLVM-IR) information.
   InformationCache &InfoCache;
+
+  /// Functions, blocks, and instructions we delete after manifest is done.
+  ///
+  ///{
+  SmallPtrSet<Function *, 8> ToBeDeletedFunctions;
+  SmallPtrSet<BasicBlock *, 8> ToBeDeletedBlocks;
+  SmallPtrSet<Instruction *, 8> ToBeDeletedInsts;
+  ///}
 };
 
 /// An interface to query the internal state of an abstract attribute.
 ///
 /// The abstract state is a minimal interface that allows the Attributor to
 /// communicate with the abstract attributes about their internal state without
 /// enforcing or exposing implementation details, e.g., the (existence of an)
 /// underlying lattice.

llvm/lib/Transforms/IPO/Attributor.cpp

 
 static cl::opt<bool> VerifyAttributor(
     "attributor-verify", cl::Hidden,
     cl::desc("Verify the Attributor deduction and "
              "manifestation of attributes -- may issue false-positive errors"),
     cl::init(false));
 
 static cl::opt<bool> EnableHeapToStack("enable-heap-to-stack-conversion",
-                                       cl::init(false), cl::Hidden);
+                                       cl::init(true), cl::Hidden);
 
 static cl::opt<int> MaxHeapToStackSize("max-heap-to-stack-size",
-                                       cl::init(1024), cl::Hidden);
+                                       cl::init(128), cl::Hidden);
 
 /// Logic operators for the change status enum class.
 ///
 ///{
 ChangeStatus llvm::operator|(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::CHANGED ? l : r;
 }
 ChangeStatus llvm::operator&(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::UNCHANGED ? l : r;

 
 /// NoCapture attribute deduction for a call site return value.
 using AANoCaptureCallSiteReturned = AANoCaptureFloating;
 
 /// ----------------------- Heap-To-Stack Conversion ---------------------------
 struct AAHeapToStackImpl : public AAHeapToStack {
   AAHeapToStackImpl(const IRPosition &IRP) : AAHeapToStack(IRP) {}
 
-  void initialize(Attributor &A) override {
-    const Function *F = getAssociatedFunction();
-    const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
-
-    auto MallocCheck = [&](Instruction &I) {
-      if (isMallocLikeFn(&I, TLI))
-        if (auto *CSize = dyn_cast<ConstantInt>(I.getOperand(0)))
-          if (CSize->getValue().sle(MaxHeapToStackSize)) {
-            MallocCalls.insert(const_cast<Instruction *>(&I));
-            LLVM_DEBUG(dbgs() << "H2S: Initial malloc call: " << I << "\n");
-          }
-      return true;
-    };
-
-    // Get all malloc calls from this function.
-    A.checkForAllCallLikeInstructions(MallocCheck, *this);
-
-    // No malloc calls, there is no work to be done.
-    if (MallocCalls.empty())
-      indicatePessimisticFixpoint();
+  const std::string getAsStr() const override {
+    return "[H2S] Mallocs: " + std::to_string(MallocCalls.size());
   }
 
   ChangeStatus manifest(Attributor &A) override {
     assert(getState().isValidState() &&
            "Attempted to manifest an invalid state!");
 
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
     Function *F = getAssociatedFunction();
+    const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
 
-    for (const Instruction *Malloc : MallocCalls) {
-      Instruction *MallocCall = const_cast<Instruction *>(Malloc);
+    for (Instruction *MallocCall : MallocCalls) {
       // This malloc cannot be replaced.
       if (BadMallocCalls.count(MallocCall))
         continue;
 
       for (Instruction *FreeCall : FreesForMalloc[MallocCall]) {
         LLVM_DEBUG(dbgs() << "H2S: Removing free call: " << *FreeCall << "\n");
-        FreeCall->eraseFromParent();
+        // FreeCall->eraseFromParent();
+        A.deleteAfterManifest(*FreeCall);
         HasChanged = ChangeStatus::CHANGED;
       }
 
       LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall
                         << "\n");
 
-      unsigned AS = cast<PointerType>(MallocCall->getType())->getAddressSpace();
+      Constant *Size;
+      if (isCallocLikeFn(MallocCall, TLI)) {
+        auto *Num = cast<ConstantInt>(MallocCall->getOperand(0));
+        auto *SizeT = dyn_cast<ConstantInt>(MallocCall->getOperand(1));
+        APInt TotalSize = SizeT->getValue() * Num->getValue();
+        Size =
+            ConstantInt::get(MallocCall->getOperand(0)->getType(), TotalSize);
+      } else {
+        Size = cast<ConstantInt>(MallocCall->getOperand(0));
+      }
+
+        unsigned AS =
+            cast<PointerType>(MallocCall->getType())->getAddressSpace();
         Instruction *AI = new AllocaInst(Type::getInt8Ty(F->getContext()), AS,
-                                       MallocCall->getOperand(0));
-      F->begin()->getInstList().insert(F->begin()->begin(), AI);
+                                         Size, "", MallocCall->getNextNode());
+        // F->begin()->getInstList().insert(F->begin()->begin(), AI);
 
         if (AI->getType() != MallocCall->getType()) {
-        auto *BC = new BitCastInst(AI, MallocCall->getType());
-        BC->insertAfter(AI);
-        AI = BC;
+          AI = new BitCastInst(AI, MallocCall->getType(), "malloc_bc",
+                               AI->getNextNode());
         }
 
       MallocCall->replaceAllUsesWith(AI);
 
       if (auto *II = dyn_cast<InvokeInst>(MallocCall)) {
         auto *NBB = II->getNormalDest();
         BranchInst::Create(NBB, MallocCall->getParent());
-        MallocCall->eraseFromParent();
+        A.deleteAfterManifest(*MallocCall);
       } else {
-        MallocCall->eraseFromParent();
+        A.deleteAfterManifest(*MallocCall);
+      }
+
+      if (isCallocLikeFn(MallocCall, TLI)) {
+        auto *BI = new BitCastInst(AI, MallocCall->getType(), "calloc_bc",
+                                   AI->getNextNode());
+        Value *Ops[] = {
+            BI, ConstantInt::get(F->getContext(), APInt(8, 0, false)), Size,
+            ConstantInt::get(Type::getInt1Ty(F->getContext()), false)};
+
+        Type *Tys[] = {BI->getType(), MallocCall->getOperand(0)->getType()};
+        Module *M = F->getParent();
+        Function *Fn = Intrinsic::getDeclaration(M, Intrinsic::memset, Tys);
+        CallInst::Create(Fn, Ops, "", BI->getNextNode());
       }
       HasChanged = ChangeStatus::CHANGED;
     }
 
     return HasChanged;
   }
 
   /// Collection of all malloc calls in a function.
-  SmallSetVector<const Instruction *, 4> MallocCalls;
+  SmallSetVector<Instruction *, 4> MallocCalls;
 
   /// Collection of malloc calls that cannot be converted.
   DenseSet<const Instruction *> BadMallocCalls;
 
   /// A map for each malloc call to the set of associated free calls.
   DenseMap<Instruction *, SmallPtrSet<Instruction *, 4>> FreesForMalloc;
 
   ChangeStatus updateImpl(Attributor &A) override;
 };
 
 ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
   const Function *F = getAssociatedFunction();
   const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
 
-  auto CallCheck = [&](Instruction &I) {
-    for (const Instruction *MallocCall : MallocCalls) {
-      if (isFreeCall(&I, TLI)) {
-        if (auto *Ptr = dyn_cast<Instruction>(I.getOperand(0)))
-          if (Ptr == MallocCall) {
-            FreesForMalloc[const_cast<Instruction *>(MallocCall)].insert(&I);
-            return true;
-          }
-        continue;
-      }
+  auto UsesCheck = [&](Instruction &I) {
+    SmallPtrSet<const Use *, 8> Visited;
+    SmallVector<const Use *, 8> Worklist;
 
-      CallSite CS(&I);
+    for (Use &U : I.uses())
+      Worklist.push_back(&U);
+      // Worklist.push_back(const_cast<Use *>(&U));
 
-      Function *CalledFunction = CS.getCalledFunction();
-      const IRPosition &IRPos = IRPosition::value(*CalledFunction);
-      auto NoFreeAA = A.getAAFor<AANoFree>(*this, IRPos);
+    while (!Worklist.empty()) {
+      const Use *U = Worklist.pop_back_val();
+      if (!Visited.insert(U).second)
+        continue;
+
+      auto *UserI = U->getUser();
 
-      for (Argument &Arg : CalledFunction->args()) {
-        if (dyn_cast<Instruction>(&Arg) != MallocCall)
-          continue;
+      if (isa<LoadInst>(UserI) || isa<StoreInst>(UserI))
 
-        BadMallocCalls.insert(MallocCall);
-        const IRPosition &IPos = IRPosition::value(Arg);
-        auto NoCaptureAA = A.getAAFor<AANoCapture>(*this, IPos);
+      // NOTE: Right now, if a function that has malloc pointer as an argument
+      // frees memory, we assume that the malloc pointer is freed.
+      if (auto *CB = dyn_cast<CallBase>(UserI)) {
+        // Record malloc.
+        if (isFreeCall(UserI, TLI)) {
+          FreesForMalloc[&I].insert(
+              cast<Instruction>(const_cast<User *>(UserI)));
+          continue;
+        }
 
-        // If a function does not capture the pointer we are fine.
-        if(NoCaptureAA.isAssumedNoCapture())
-          return true;
         // If a function does not free memory we are fine
-        // Note: Right now, if a function that has malloc pointer as an argument
-        // frees memory, we assume that the malloc pointer is freed.
-        //
+        const auto &NoFreeAA =
+            A.getAAFor<AANoFree>(*this, IRPosition::callsite_function(*CB));
+
         // TODO: Add nofree callsite argument attribute to indicate that pointer
         // argument is not freed.
-        if (NoFreeAA.isAssumedNoFree())
-          return true;
+        if (!CB->isArgOperand(U))
+          continue;
+        unsigned ArgNo = U - CB->arg_begin();
+        const auto &NoCaptureAA = A.getAAFor<AANoCapture>(
+            *this, IRPosition::callsite_argument(*CB, ArgNo));
 
-        // Remove all Frees that we thought we could remove.
-        FreesForMalloc.erase(Malloc);
-        BadMallocCalls.insert(MallocCall);
+        if (!NoCaptureAA.isAssumedNoCapture() || !NoFreeAA.isAssumedNoFree()) {
+          LLVM_DEBUG(dbgs() << "[H2S] Bad user: " << *UserI << "\n");
           return false;
         }
+        continue;
+      }
+      if (isa<GetElementPtrInst>(UserI) || isa<BitCastInst>(UserI)) {
+        for (Use &U : UserI->uses())
+          Worklist.push_back(&U);
+        continue;
+      }
+
+      if (isa<PHINode>(UserI))
+        continue;
+
+      // Unknown user.
+      LLVM_DEBUG(dbgs() << "[H2S] Unknown user: " << *UserI << "\n");
+      return false;
     }
+    return true;
+  };
+
+  auto MallocCallocCheck = [&](Instruction &I) {
+    if (isMallocLikeFn(&I, TLI)) {
+      if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(0)))
+        if (!Size->getValue().sle(MaxHeapToStackSize))
+          return true;
+    } else if (isCallocLikeFn(&I, TLI)) {
+      if (auto *Num = dyn_cast<ConstantInt>(I.getOperand(0)))
+        if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
+          if (!(Size->getValue() * Num->getValue()).sle(MaxHeapToStackSize))
+            return true;
+    } else
+      return true;
+    
+
+    if (BadMallocCalls.count(&I))
+      return true;
+
+    if (UsesCheck(I))
+      MallocCalls.insert(&I);
+    else
+      BadMallocCalls.insert(&I);
+    return true;
   };
 
   size_t NumBadMallocs = BadMallocCalls.size();
 
-  A.checkForAllCallLikeInstructions(CallCheck, *this);
+  A.checkForAllCallLikeInstructions(MallocCallocCheck, *this);
 
   if (NumBadMallocs != BadMallocCalls.size())
     return ChangeStatus::CHANGED;
 
   return ChangeStatus::UNCHANGED;
 }
 
+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();
+  }
+};
+
 /// ----------------------------------------------------------------------------
 ///                               Attributor
 /// ----------------------------------------------------------------------------
 
 bool Attributor::isAssumedDead(const AbstractAttribute &AA,
                                const AAIsDead *LivenessAA) {
   const Instruction *CtxI = AA.getIRPosition().getCtxI();
   if (!CtxI)
     return false;
 
   if (!LivenessAA)
     LivenessAA =
         &getAAFor<AAIsDead>(AA, IRPosition::function(*CtxI->getFunction()));
 
   // Don't check liveness for AAIsDead.
   if (&AA == LivenessAA)
     return false;
 
   if (!LivenessAA->isAssumedDead(CtxI))
     return false;
 
   // TODO: Do not track dependences automatically but add it here as only a
do {
     Worklist.insert(AllAbstractAttributes.begin() + NumAAs,
                     AllAbstractAttributes.end());
 
   } while (!Worklist.empty() && ++IterationCounter < MaxFixpointIterations);
 
   size_t NumFinalAAs = AllAbstractAttributes.size();
 
   LLVM_DEBUG(dbgs() << "\n[Attributor] Fixpoint iteration done after: "
                     << IterationCounter << "/" << MaxFixpointIterations
                     << " iterations\n");
 
   bool FinishedAtFixpoint = Worklist.empty();
 
   // Reset abstract arguments not settled in a sound fixpoint by now. This
   // happens when we stopped the fixpoint iteration early. Note that only the
   // ones marked as "changed" *and* the ones transitively depending on them
   // need to be reverted to a pessimistic state. Others might not be in a
if (!State.isAtFixpoint()) {
 
       NumAttributesTimedOut++;
     }
 
     auto &QuerriedAAs = QueryMap[ChangedAA];
     ChangedAAs.append(QuerriedAAs.begin(), QuerriedAAs.end());
   }
 
   LLVM_DEBUG({
     if (!Visited.empty())
       dbgs() << "\n[Attributor] Finalized " << Visited.size()
              << " abstract attributes.\n";
   });
 
   unsigned NumManifested = 0;
   unsigned NumAtFixpoint = 0;
   ChangeStatus ManifestChange = ChangeStatus::UNCHANGED;
ChangeStatus Attributor::run() {
 
   NumAttributesManifested += NumManifested;
   NumAttributesValidFixpoint += NumAtFixpoint;
 
   (void)NumFinalAAs;
   assert(
       NumFinalAAs == AllAbstractAttributes.size() &&
       "Expected the final number of abstract attributes to remain unchanged!");
+
+  // Delete stuff at the end to avoid invalid references and a nice order.
+  LLVM_DEBUG(dbgs() << "\n[Attributor] Delete " << ToBeDeletedFunctions.size()
+                    << " functions and " << ToBeDeletedBlocks.size()
+                    << " blocks and " << ToBeDeletedInsts.size()
+                    << " instructions\n");
+
+  for (Instruction *I : ToBeDeletedInsts) {
+    if (I->hasNUsesOrMore(1))
+      I->replaceAllUsesWith(UndefValue::get(I->getType()));
+    I->eraseFromParent();
+  }
+  for (BasicBlock *BB : ToBeDeletedBlocks) {
+    // TODO: Check if we need to replace users (PHIs, indirect branches?)
+    BB->eraseFromParent();
+  }
+  for (Function *Fn : ToBeDeletedFunctions) {
+    Fn->replaceAllUsesWith(UndefValue::get(Fn->getType()));
+    Fn->eraseFromParent();
+  }
+
   return ManifestChange;
 }
 
 /// Helper function that checks if an abstract attribute of type \p AAType
 /// should be created for IR position \p IRP and if so creates and registers it
 /// with the Attributor \p A.
 ///
 /// This method will look at the provided whitelist. If one is given and the
 /// kind \p AAType::ID is not contained, no abstract attribute is created.
 ///
 /// \returns The created abstract argument, or nullptr if none was created.
 template <typename AAType>
 static const AAType *checkAndRegisterAA(const IRPosition &IRP, Attributor &A,
                                         DenseSet<const char *> *Whitelist) {
   if (Whitelist && !Whitelist->count(&AAType::ID))
     return nullptr;
 
   return &A.registerAA<AAType>(*new AAType(IRP));
 }
 
 void Attributor::identifyDefaultAbstractAttributes(
-    Function &F, DenseSet<const char *> *Whitelist) {
+    Function &F, std::function<TargetLibraryInfo *(Function &)> &TLIGetter,
+    DenseSet<const char *> *Whitelist) {
+
+  if (EnableHeapToStack)
+    InfoCache.FuncTLIMap[&F] = TLIGetter(F);
 
   IRPosition FPos = IRPosition::function(F);
 
   // Check for dead BasicBlocks in every function.
   // We need dead instruction detection because we do not want to deal with
   // broken IR in which SSA rules do not apply.
   checkAndRegisterAA<AAIsDeadFunction>(FPos, *this, /* Whitelist */ nullptr);
 
-  // Every function might be "will-return".
   checkAndRegisterAA<AAWillReturnFunction>(FPos, *this, Whitelist);
+  // Every function might be "will-return".
 
   // Every function can be nounwind.
   checkAndRegisterAA<AANoUnwindFunction>(FPos, *this, Whitelist);
 
   // Every function might be marked "nosync"
   checkAndRegisterAA<AANoSyncFunction>(FPos, *this, Whitelist);
 
   // Every function might be "no-free".
   checkAndRegisterAA<AANoFreeFunction>(FPos, *this, Whitelist);
 
   // Every function might be "no-return".
   checkAndRegisterAA<AANoReturnFunction>(FPos, *this, Whitelist);
 
+  // Every function might be applicable for Heap-To-Stack conversion.
+  if (EnableHeapToStack)
+    checkAndRegisterAA<AAHeapToStackFunction>(FPos, *this, Whitelist);
+
   // Return attributes are only appropriate if the return type is non void.
   Type *ReturnType = F.getReturnType();
   if (!ReturnType->isVoidTy()) {
     // Argument attribute "returned" --- Create only one per function even
     // though it is an argument attribute.
     checkAndRegisterAA<AAReturnedValuesFunction>(FPos, *this, Whitelist);
 
     if (ReturnType->isPointerTy()) {
for (Argument &Arg : F.args()) {
     if (Arg.getType()->isPointerTy()) {
       IRPosition ArgPos = IRPosition::argument(Arg);
       // Every argument with pointer type might be marked nonnull.
       checkAndRegisterAA<AANonNullArgument>(ArgPos, *this, Whitelist);
 
       // Every argument with pointer type might be marked dereferenceable.
       checkAndRegisterAA<AADereferenceableArgument>(ArgPos, *this, Whitelist);
 
-      // Every argument with pointer type might be marked align.
-      checkAndRegisterAA<AAAlignArgument>(ArgPos, *this, Whitelist);
-
       // Every argument with pointer type might be marked nocapture.
       checkAndRegisterAA<AANoCaptureArgument>(ArgPos, *this, Whitelist);
+
+      // Every argument with pointer type might be marked align.
+      checkAndRegisterAA<AAAlignArgument>(ArgPos, *this, Whitelist);
     }
   }
 
   // Walk all instructions to find more attribute opportunities and also
   // interesting instructions that might be queried by abstract attributes
   // during their initialization or update.
   auto &ReadOrWriteInsts = InfoCache.FuncRWInstsMap[&F];
   auto &InstOpcodeMap = InfoCache.FuncInstOpcodeMap[&F];
 
   for (Instruction &I : instructions(&F)) {
     bool IsInterestingOpcode = false;
 
     // To allow easy access to all instructions in a function with a given
     // opcode we store them in the InfoCache. As not all opcodes are interesting
OS << "[P: " << getIRPosition() << "][" << getAsStr() << "][S: " << getState()
      << "]";
 }
 ///}
 
 /// ----------------------------------------------------------------------------
 ///                       Pass (Manager) Boilerplate
 /// ----------------------------------------------------------------------------
 
-static bool runAttributorOnModule(Module &M) {
+static bool runAttributorOnModule(
+    Module &M, std::function<TargetLibraryInfo *(Function &)> &TLIGetter) {
   if (DisableAttributor)
     return false;
 
   LLVM_DEBUG(dbgs() << "[Attributor] Run on module with " << M.size()
                     << " functions.\n");
 
   // Create an Attributor and initially empty information cache that is filled
   // while we identify default attribute opportunities.
for (Function &F : M) {
     if (F.hasFnAttribute(Attribute::Naked) ||
         F.hasFnAttribute(Attribute::OptimizeNone))
       continue;
 
     NumFnWithExactDefinition++;
 
     // Populate the Attributor with abstract attribute opportunities in the
     // function and the information cache with IR information.
-    A.identifyDefaultAbstractAttributes(F);
+    A.identifyDefaultAbstractAttributes(F, TLIGetter);
   }
 
   return A.run() == ChangeStatus::CHANGED;
 }
 
 PreservedAnalyses AttributorPass::run(Module &M, ModuleAnalysisManager &AM) {
-  if (runAttributorOnModule(M)) {
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+
+  std::function<TargetLibraryInfo *(Function &)> TLIGetter =
+      [&](Function &F) -> TargetLibraryInfo * {
+    return &FAM.getResult<TargetLibraryAnalysis>(F);
+  };
+
+  if (runAttributorOnModule(M, TLIGetter)) {
     // FIXME: Think about passes we will preserve and add them here.
     return PreservedAnalyses::none();
   }
   return PreservedAnalyses::all();
 }
 
 namespace {
 
 struct AttributorLegacyPass : public ModulePass {
   static char ID;
 
   AttributorLegacyPass() : ModulePass(ID) {
     initializeAttributorLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
-    return runAttributorOnModule(M);
+    std::function<TargetLibraryInfo *(Function &)> TLIGetter =
+        [&](Function &F) -> TargetLibraryInfo * { return nullptr; };
+
+    return runAttributorOnModule(M, TLIGetter);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     // FIXME: Think about passes we will preserve and add them here.
+    // ModulePass::getAnalysisUsage(AU);
+    AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.setPreservesCFG();
   }
 };
 
 } // end anonymous namespace
 
 Pass *llvm::createAttributorLegacyPass() { return new AttributorLegacyPass(); }
 
 char AttributorLegacyPass::ID = 0;
 
 const char AAReturnedValues::ID = 0;
 const char AANoUnwind::ID = 0;
 const char AANoSync::ID = 0;
 const char AANoFree::ID = 0;
 const char AANonNull::ID = 0;
 const char AANoRecurse::ID = 0;
 const char AAWillReturn::ID = 0;
 const char AANoAlias::ID = 0;
 const char AANoReturn::ID = 0;
 const char AAIsDead::ID = 0;
 const char AADereferenceable::ID = 0;
 const char AAAlign::ID = 0;
 const char AANoCapture::ID = 0;
+const char AAHeapToStack::ID = 0;
 
 // Macro magic to create the static generator function for attributes that
 // follow the naming scheme.
 
 #define SWITCH_PK_INV(CLASS, PK, POS_NAME)                                     \
   case IRPosition::PK:                                                         \
     llvm_unreachable("Cannot create " #CLASS " for a " POS_NAME " position!");
 
 #define SWITCH_PK_CREATE(CLASS, IRP, PK, SUFFIX)                               \
   case IRPosition::PK:                                                         \
     AA = new CLASS##SUFFIX(IRP);                                               \
     break;
 
 #define CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(CLASS)                 \
   CLASS &CLASS::createForPosition(const IRPosition &IRP, Attributor &A) {      \
     CLASS *AA = nullptr;                                                       \
     switch (IRP.getPositionKind()) {                                           \
       SWITCH_PK_INV(CLASS, IRP_INVALID, "invalid")                             \
-      SWITCH_PK_INV(CLASS, IRP_FLOAT, "floating")                              \
       SWITCH_PK_INV(CLASS, IRP_ARGUMENT, "argument")                           \
+      SWITCH_PK_INV(CLASS, IRP_FLOAT, "floating")                              \
       SWITCH_PK_INV(CLASS, IRP_RETURNED, "returned")                           \
       SWITCH_PK_INV(CLASS, IRP_CALL_SITE_RETURNED, "call site returned")       \
       SWITCH_PK_INV(CLASS, IRP_CALL_SITE_ARGUMENT, "call site argument")       \
       SWITCH_PK_CREATE(CLASS, IRP, IRP_FUNCTION, Function)                     \
       SWITCH_PK_CREATE(CLASS, IRP, IRP_CALL_SITE, CallSite)                    \
     }                                                                          \
     AA->initialize(A);                                                         \
     return *AA;                                                                \
   }
 
+#define CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(CLASS)                 \
+  CLASS &CLASS::createForPosition(const IRPosition &IRP, Attributor &A) {      \
+    CLASS *AA = nullptr;                                                       \
+    switch (IRP.getPositionKind()) {                                           \
+      SWITCH_PK_INV(CLASS, IRP_INVALID, "invalid")                             \
+      SWITCH_PK_INV(CLASS, IRP_ARGUMENT, "argument")                           \
+      SWITCH_PK_INV(CLASS, IRP_FLOAT, "floating")                              \
+      SWITCH_PK_INV(CLASS, IRP_RETURNED, "returned")                           \
+      SWITCH_PK_INV(CLASS, IRP_CALL_SITE_RETURNED, "call site returned")       \
+      SWITCH_PK_INV(CLASS, IRP_CALL_SITE_ARGUMENT, "call site argument")       \
+      SWITCH_PK_INV(CLASS, IRP_CALL_SITE, "call site")                       \
+      SWITCH_PK_CREATE(CLASS, IRP, IRP_FUNCTION, Function)                     \
+    }                                                                          \
+    AA->initialize(A);                                                         \
+    return *AA;                                                                \
+  }
+
 #define CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(CLASS)                    \
   CLASS &CLASS::createForPosition(const IRPosition &IRP, Attributor &A) {      \
     CLASS *AA = nullptr;                                                       \
     switch (IRP.getPositionKind()) {                                           \
       SWITCH_PK_INV(CLASS, IRP_INVALID, "invalid")                             \
       SWITCH_PK_INV(CLASS, IRP_FUNCTION, "function")                           \
       SWITCH_PK_INV(CLASS, IRP_CALL_SITE, "call site")                         \
       SWITCH_PK_CREATE(CLASS, IRP, IRP_FLOAT, Floating)                        \

 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoSync)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoFree)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoRecurse)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAWillReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAIsDead)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReturnedValues)
 
+CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAHeapToStack)
+
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANonNull)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoAlias)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AADereferenceable)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAlign)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoCapture)
 
 #undef CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef SWITCH_PK_CREATE
 #undef SWITCH_PK_INV
 
 INITIALIZE_PASS_BEGIN(AttributorLegacyPass, "attributor",
                       "Deduce and propagate attributes", false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(AttributorLegacyPass, "attributor",
                     "Deduce and propagate attributes", false, false)

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

+; RUN: opt -passes=attributor --attributor-disable=false -enable-heap-to-stack-conversion=true -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 @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
+
+declare i32 @no_return_call() noreturn
+
+declare void @free(i8* nocapture)
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) nounwind
+
+; TEST 1 - negative, pointer freed in another function.
+
+define void @test1() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: @malloc(i64 4)
+  ; CHECK-NEXT: @nocapture_func_frees_pointer(i8* nocapture %1)
+  tail call void @nocapture_func_frees_pointer(i8* %1)
+  tail call void (...) @func_throws()
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 2 - negative, call to a sync function.
+
+define void @test2() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: @malloc(i64 4)
+  ; CHECK-NEXT: @sync_func(i8* %1)
+  tail call void @sync_func(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 3 - 1 malloc, 1 free
+
+define void @test3() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = alloca i8, i64 4
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  ; CHECK-NOT: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+declare noalias i8* @calloc(i64, i64)
+
+define void @test0() {
+  %1 = tail call noalias i8* @calloc(i64 2, i64 4)
+  ; CHECK: %1 = alloca i8, i64 8
+  ; CHECK-NEXT: %calloc_bc = bitcast i8* %1 to i8*
+  ; CHECK-NEXT: call void @llvm.memset.p0i8.i64(i8* %calloc_bc, i8 0, i64 8, i1 false)
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  ; CHECK-NOT: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 4 
+define void @test4() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = alloca i8, i64 4
+  ; CHECK-NEXT: @nofree_func(i8* nocapture %1)
+  tail call void @nofree_func(i8* %1)
+  ret void
+}
+
+; TEST 5 - not all exit paths have a call to free, but all uses of malloc
+; are in nofree functions and are not captured
+
+define void @test5(i32) {
+  %2 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %2 = alloca i8, i64 4
+  ; CHECK-NEXT: icmp eq i32 %0, 0
+  %3 = icmp eq i32 %0, 0
+  br i1 %3, label %5, label %4
+
+4:                                                ; preds = %1
+  tail call void @nofree_func(i8* %2)
+  br label %6
+
+5:                                                ; preds = %1
+  tail call void @free(i8* %2)
+  ; CHECK-NOT: @free(i8* %2)
+  br label %6
+
+6:                                                ; preds = %5, %4
+  ret void
+}
+
+; TEST 6 - all exit paths have a call to free
+
+define void @test6(i32) {
+  %2 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %2 = alloca i8, i64 4
+  ; CHECK-NEXT: icmp eq i32 %0, 0
+  %3 = icmp eq i32 %0, 0
+  br i1 %3, label %5, label %4
+
+4:                                                ; preds = %1
+  tail call void @nofree_func(i8* %2)
+  tail call void @free(i8* %2)
+  ; CHECK-NOT: @free(i8* %2)
+  br label %6
+
+5:                                                ; preds = %1
+  tail call void @free(i8* %2)
+  ; CHECK-NOT: @free(i8* %2)
+  br label %6
+
+6:                                                ; preds = %5, %4
+  ret void
+}
+
+; TEST 7 - free is dead.
+
+define void @test7() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: alloca i8, i64 4
+  ; CHECK-NEXT: tail call i32 @no_return_call()
+  tail call i32 @no_return_call()
+  ; CHECK-NOT: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 8 - Negative: bitcast pointer used in capture function
+
+define void @test8() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  %2 = bitcast i8* %1 to i32*
+  store i32 10, i32* %2
+  %3 = load i32, i32* %2
+  tail call void @foo(i32* %2)
+  ; CHECK: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 9 - FIXME: malloc should be converted.
+define void @test9() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  %2 = bitcast i8* %1 to i32*
+  store i32 10, i32* %2
+  %3 = load i32, i32* %2
+  tail call void @foo_nounw(i32* %2)
+  ; CHECK: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; TEST 10 - 1 malloc, 1 free
+
+define i32 @test10() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = alloca i8, i64 4
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  %2 = bitcast i8* %1 to i32*
+  store i32 10, i32* %2
+  %3 = load i32, i32* %2
+  ; CHECK-NOT: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret i32 %3
+}
+
+define i32 @test_lifetime() {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK: %1 = tail call noalias i8* @malloc(i64 4)
+  ; CHECK-NEXT: @no_sync_func(i8* nocapture %1)
+  tail call void @no_sync_func(i8* %1)
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %1)
+  %2 = bitcast i8* %1 to i32*
+  store i32 10, i32* %2
+  %3 = load i32, i32* %2
+  ; CHECK: @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-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) {
+  ; CHECK: alloca i8, i64 4
+  ; CHECK-NEXT: %3 = bitcast
+  %2 = tail call noalias i8* @malloc(i64 4)
+  %3 = bitcast i8* %2 to i32*
+  store i32 10, i32* %3, align 4
+  %4 = icmp eq i32 %0, 1
+  br i1 %4, label %8, label %5
+
+5:                                                ; preds = %1
+  %6 = getelementptr inbounds i8, i8* %2, i64 -4
+  %7 = bitcast i8* %6 to i32*
+  br label %13
+
+8:                                                ; preds = %1, %13
+  %9 = phi i32 [ %15, %13 ], [ 10, %1 ]
+  %10 = phi i32* [ %14, %13 ], [ %3, %1 ]
+  %11 = getelementptr inbounds i32, i32* %10, i64 -1
+  %12 = icmp eq i32 %9, 0
+  br i1 %12, label %16, label %13
+
+13:                                               ; preds = %5, %8
+  %14 = phi i32* [ %7, %5 ], [ %11, %8 ]
+  %15 = load i32, i32* %14, align 4
+  br label %8
+
+16:                                               ; preds = %8
+  %17 = bitcast i32* %11 to i8*
+  ; CHECK-NOT: @free
+  tail call void @free(i8* %17)
+  %18 = load i32, i32* %11, align 4
+  ret i32 %18
+}
+
+define i32 @malloc_in_loop(i32 %0) #0 {
+  %2 = alloca i32, align 4
+  %3 = alloca i32*, align 8
+  store i32 %0, i32* %2, align 4
+  br label %4
+
+4:                                                ; preds = %8, %1
+  %5 = load i32, i32* %2, align 4
+  %6 = add nsw i32 %5, -1
+  store i32 %6, i32* %2, align 4
+  %7 = icmp sgt i32 %6, 0
+  br i1 %7, label %8, label %11
+
+8:                                                ; preds = %4
+  %9 = call noalias i8* @malloc(i64 4) #2
+  %10 = bitcast i8* %9 to i32*
+  store i32* %10, i32** %3, align 8
+  br label %4
+
+11:                                               ; preds = %4
+  %12 = load i32*, i32** %3, align 8
+  %13 = bitcast i32* %12 to i8*
+  call void @free(i8* %13) #2
+  ret i32 5
+}