[Prototype] Heap-To-Stack Conversion Pass

[hfinkel]

Summary

This is a prototype heap-to-stack conversion pass. It's not really ready to be reviewed yet. Posting this now because I'm going to talk about it tomorrow at the dev meeting (lightning talk).

[xbolva00]

Nice, this looks much better than my D47345!

Does this patch also solves all raised concerns from D47345?

[hfinkel]

Add the pointer-capturing check when potential exits are found.

[hfinkel]

Fix variable initialization so that we don't infinte-loop when we refine.

[xbolva00]

In D53362#1267885, @hfinkel wrote:

Add the pointer-capturing check when potential exits are found.

Not sure if enough, since @efriedma noted in D47345:

"Your "capture" function captures the pointer, but it's possible to write a function which takes the pointer as a parameter but doesn't capture it. For example, given the function void f(void *p) { free(p); }, LLVM will mark the parameter "p" nocapture."

[hfinkel]

In D53362#1267929, @xbolva00 wrote:

In D53362#1267885, @hfinkel wrote:

Add the pointer-capturing check when potential exits are found.

Not sure if enough, since @efriedma noted in D47345:

"Your "capture" function captures the pointer, but it's possible to write a function which takes the pointer as a parameter but doesn't capture it. For example, given the function void f(void *p) { free(p); }, LLVM will mark the parameter "p" nocapture."

free doesn't capture, but that's okay. I'm not depending on the capture check for that because, for free, the path reachability check takes care of that. For everything else, that's exactly correct, and in fact, I want to be able to handle cases where functions take the pointers and don't capture them (the tricky part there is actually that, if the pointer is captured, we also need to prove eventual termination).

I need to update the description to describe the algorithm and why I think that it works. ;) -- I'll do that soon.

[efriedma]

I think you have to be careful about something like the following:

void f() {
  void *p = malloc(4);
  nocapture_func_frees_pointer(p);
  func_throws();
  free(p);
}

I don't think it's possible to trigger the issue at the moment, though, because we don't have a mustreturn attribute (so isGuaranteedToTransferExecutionToSuccessor() will return false for nocapture_func_frees_pointer).

lib/Transforms/Scalar/HeapToStack.cpp

139

isGuaranteedToTransferExecutionToSuccessor is true for atomics.

[hfinkel]

Add some optimization remarks. Now we'll get things like his:

In file included from /src/llvm/lib/Target/AArch64/AArch64A57FPLoadBalancing.cpp:31:
In file included from /src/llvm/lib/Target/AArch64/AArch64.h:19:
In file included from /src/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h:23:
In file included from /src/llvm/include/llvm/ADT/STLExtras.h:21:
In file included from /src/llvm/include/llvm/ADT/SmallVector.h:21:
/src/llvm/include/llvm/Support/MemAlloc.h:27:18: remark: removed memory allocation (placed on stack) [-Rpass=heap-to-stack]
  void *Result = std::malloc(Sz);
                 ^
In file included from /src/llvm/lib/Target/AArch64/AArch64A57FPLoadBalancing.cpp:32:
In file included from /src/llvm/lib/Target/AArch64/AArch64InstrInfo.h:20:
In file included from /src/llvm/include/llvm/CodeGen/TargetInstrInfo.h:21:
In file included from /src/llvm/include/llvm/CodeGen/LiveRegUnits.h:18:
/src/llvm/include/llvm/ADT/BitVector.h:164:18: remark: removed memory deallocation (placed on stack) [-Rpass=heap-to-stack]
  ~BitVector() { std::free(Bits.data()); }
                 ^

During a check-clang self host, we currently remove 79 allocations (using the "very aggressive mode, at least").

[xbolva00]

Thanks for stats, interesting! Maybe try with some big C benchmarks/codebases? SPECs? :)

[hfinkel]

In D53362#1267978, @efriedma wrote:

I think you have to be careful about something like the following:

void f() {
  void *p = malloc(4);
  nocapture_func_frees_pointer(p);
  func_throws();
  free(p);
}

I don't think it's possible to trigger the issue at the moment, though, because we don't have a mustreturn attribute (so isGuaranteedToTransferExecutionToSuccessor() will return false for nocapture_func_frees_pointer).

I think you're right. Thanks! -- I also rebased my nofree-attribute patch the other day. I'll post that update too. I think that can be used to solve this problem because in cases where we can infer nocapture we should also be able to infer nofree (unless, of course, as in this case, where the function actually does free the data).

lib/Transforms/Scalar/HeapToStack.cpp

139

Indeed (at least some of them)...

bool llvm::isGuaranteedToTransferExecutionToSuccessor(const Instruction *I) {
  // A memory operation returns normally if it isn't volatile. A volatile
  // operation is allowed to trap.
  //
  // An atomic operation isn't guaranteed to return in a reasonable amount of
  // time because it's possible for another thread to interfere with it for an
  // arbitrary length of time, but programs aren't allowed to rely on that.
  if (const LoadInst *LI = dyn_cast<LoadInst>(I))
    return !LI->isVolatile();
  if (const StoreInst *SI = dyn_cast<StoreInst>(I))
    return !SI->isVolatile();
  if (const AtomicCmpXchgInst *CXI = dyn_cast<AtomicCmpXchgInst>(I))
    return !CXI->isVolatile();
  if (const AtomicRMWInst *RMWI = dyn_cast<AtomicRMWInst>(I))
    return !RMWI->isVolatile();
  if (const MemIntrinsic *MII = dyn_cast<MemIntrinsic>(I))
    return !MII->isVolatile();

[sstefan1]

Hi Hal! Are you planning on working on this one? If not maybe I can take over? Now that we have nofree & nosync added and deduced in the Attributor, maybe this can be revisited.

[hfinkel]

In D53362#1592224, @sstefan1 wrote:

Hi Hal! Are you planning on working on this one? If not maybe I can take over? Now that we have nofree & nosync added and deduced in the Attributor, maybe this can be revisited.

I would certainly like to see progress on this, and it's great that we can now get back to this. If you have time in the short term to make progress here, I'm happy for you to do so (and, of course, I'll help review).

include/llvm/InitializePasses.h

 void initializeGVNLegacyPassPass(PassRegistry&);
 void initializeGVNSinkLegacyPassPass(PassRegistry&);
 void initializeGlobalDCELegacyPassPass(PassRegistry&);
 void initializeGlobalMergePass(PassRegistry&);
 void initializeGlobalOptLegacyPassPass(PassRegistry&);
 void initializeGlobalSplitPass(PassRegistry&);
 void initializeGlobalsAAWrapperPassPass(PassRegistry&);
 void initializeGuardWideningLegacyPassPass(PassRegistry&);
+void initializeHeapToStackLegacyPassPass(PassRegistry&);
 void initializeHotColdSplittingLegacyPassPass(PassRegistry&);
 void initializeHWAddressSanitizerPass(PassRegistry&);
 void initializeIPCPPass(PassRegistry&);
 void initializeIPSCCPLegacyPassPass(PassRegistry&);
 void initializeIRCELegacyPassPass(PassRegistry&);
 void initializeIRTranslatorPass(PassRegistry&);
 void initializeIVUsersWrapperPassPass(PassRegistry&);
 void initializeIfConverterPass(PassRegistry&);

include/llvm/LinkAllPasses.h

ForcePassLinking() {
       (void) llvm::createSeparateConstOffsetFromGEPPass();
       (void) llvm::createSpeculativeExecutionPass();
       (void) llvm::createSpeculativeExecutionIfHasBranchDivergencePass();
       (void) llvm::createRewriteSymbolsPass();
       (void) llvm::createStraightLineStrengthReducePass();
       (void) llvm::createMemDerefPrinter();
       (void) llvm::createMustExecutePrinter();
       (void) llvm::createFloat2IntPass();
+      (void) llvm::createHeapToStackPass();
       (void) llvm::createEliminateAvailableExternallyPass();
       (void) llvm::createScalarizeMaskedMemIntrinPass();
 
       (void)new llvm::IntervalPartition();
       (void)new llvm::ScalarEvolutionWrapperPass();
       llvm::Function::Create(nullptr, llvm::GlobalValue::ExternalLinkage)->viewCFGOnly();
       llvm::RGPassManager RGM;
       llvm::TargetLibraryInfoImpl TLII;

include/llvm/Transforms/Scalar.h

 //===----------------------------------------------------------------------===//
 //
 // Float2Int - Demote floats to ints where possible.
 //
 FunctionPass *createFloat2IntPass();
 
 //===----------------------------------------------------------------------===//
 //
+// HeapToStack - Convert heap allocations to stack allocations where possible.
+//
+FunctionPass *createHeapToStackPass();
+
+//===----------------------------------------------------------------------===//
+//
 // NaryReassociate - Simplify n-ary operations by reassociation.
 //
 FunctionPass *createNaryReassociatePass();
 
 //===----------------------------------------------------------------------===//
 //
 // LoopDistribute - Distribute loops.
 //

include/llvm/Transforms/Scalar/HeapToStack.h

+//===-- HeapToStack.h - Heap-to-Stack Conversion ----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass moves small allocations from the heap to the stack.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCALAR_HEAPTOSTACK_H
+#define LLVM_TRANSFORMS_SCALAR_HEAPTOSTACK_H
+
+#include "llvm/IR/Function.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+/// Move instructions into successor blocks when possible.
+class HeapToStackPass : public PassInfoMixin<HeapToStackPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+}
+
+#endif // LLVM_TRANSFORMS_SCALAR_HEAPTOSTACK_H

lib/Passes/PassBuilder.cpp

 #include "llvm/Transforms/Scalar/CorrelatedValuePropagation.h"
 #include "llvm/Transforms/Scalar/DCE.h"
 #include "llvm/Transforms/Scalar/DeadStoreElimination.h"
 #include "llvm/Transforms/Scalar/DivRemPairs.h"
 #include "llvm/Transforms/Scalar/EarlyCSE.h"
 #include "llvm/Transforms/Scalar/Float2Int.h"
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Scalar/GuardWidening.h"
+#include "llvm/Transforms/Scalar/HeapToStack.h"
 #include "llvm/Transforms/Scalar/IVUsersPrinter.h"
 #include "llvm/Transforms/Scalar/IndVarSimplify.h"
 #include "llvm/Transforms/Scalar/InductiveRangeCheckElimination.h"
 #include "llvm/Transforms/Scalar/InstSimplifyPass.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Scalar/LICM.h"
 #include "llvm/Transforms/Scalar/LoopAccessAnalysisPrinter.h"
 #include "llvm/Transforms/Scalar/LoopDataPrefetch.h"
PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
 
   // Form SSA out of local memory accesses after breaking apart aggregates into
   // scalars.
   FPM.addPass(SROA());
 
   // Catch trivial redundancies
   FPM.addPass(EarlyCSEPass(EnableEarlyCSEMemSSA));
 
+  FPM.addPass(HeapToStackPass());
+
   // Hoisting of scalars and load expressions.
   if (EnableGVNHoist)
     FPM.addPass(GVNHoistPass());
 
   // Global value numbering based sinking.
   if (EnableGVNSink) {
     FPM.addPass(GVNSinkPass());
     FPM.addPass(SimplifyCFGPass());

lib/Passes/PassRegistry.def

 FUNCTION_PASS("early-cse-memssa", EarlyCSEPass(/*UseMemorySSA=*/true))
 FUNCTION_PASS("ee-instrument", EntryExitInstrumenterPass(/*PostInlining=*/false))
 FUNCTION_PASS("post-inline-ee-instrument", EntryExitInstrumenterPass(/*PostInlining=*/true))
 FUNCTION_PASS("gvn-hoist", GVNHoistPass())
 FUNCTION_PASS("instcombine", InstCombinePass())
 FUNCTION_PASS("instsimplify", InstSimplifyPass())
 FUNCTION_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 FUNCTION_PASS("float2int", Float2IntPass())
+FUNCTION_PASS("heap-to-stack", HeapToStackPass())
 FUNCTION_PASS("no-op-function", NoOpFunctionPass())
 FUNCTION_PASS("libcalls-shrinkwrap", LibCallsShrinkWrapPass())
 FUNCTION_PASS("loweratomic", LowerAtomicPass())
 FUNCTION_PASS("lower-expect", LowerExpectIntrinsicPass())
 FUNCTION_PASS("lower-guard-intrinsic", LowerGuardIntrinsicPass())
 FUNCTION_PASS("guard-widening", GuardWideningPass())
 FUNCTION_PASS("gvn", GVN())
 FUNCTION_PASS("loop-simplify", LoopSimplifyPass())

lib/Transforms/IPO/PassManagerBuilder.cpp

MPM.add(
         createPGOIndirectCallPromotionLegacyPass(false, !PGOSampleUse.empty()));
 }
 void PassManagerBuilder::addFunctionSimplificationPasses(
     legacy::PassManagerBase &MPM) {
   // Start of function pass.
   // Break up aggregate allocas, using SSAUpdater.
   MPM.add(createSROAPass());
   MPM.add(createEarlyCSEPass(EnableEarlyCSEMemSSA)); // Catch trivial redundancies
+  MPM.add(createHeapToStackPass());
   if (EnableGVNHoist)
     MPM.add(createGVNHoistPass());
   if (EnableGVNSink) {
     MPM.add(createGVNSinkPass());
     MPM.add(createCFGSimplificationPass());
   }
 
   // Speculative execution if the target has divergent branches; otherwise nop.

lib/Transforms/Scalar/CMakeLists.txt

add_llvm_library(LLVMScalarOpts
   DivRemPairs.cpp
   EarlyCSE.cpp
   FlattenCFGPass.cpp
   Float2Int.cpp
   GuardWidening.cpp
   GVN.cpp
   GVNHoist.cpp
   GVNSink.cpp
+  HeapToStack.cpp
   IVUsersPrinter.cpp
   InductiveRangeCheckElimination.cpp
   IndVarSimplify.cpp
   InferAddressSpaces.cpp
   InstSimplifyPass.cpp
   JumpThreading.cpp
   LICM.cpp
   LoopAccessAnalysisPrinter.cpp

lib/Transforms/Scalar/HeapToStack.cpp

+//===-- HeapToStack.cpp - Code HeapToStack -------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass moves instructions into successor blocks, when possible, so that
+// they aren't executed on paths where their results aren't needed.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/HeapToStack.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/CaptureTracking.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/MemoryBuiltins.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Scalar.h"
+#include <set>
+using namespace llvm;
+
+#define DEBUG_TYPE "heap-to-stack"
+
+STATISTIC(NumAConverted, "Number of converted allocations");
+STATISTIC(NumDConverted, "Number of converted deallocations");
+
+static cl::opt<bool> EnableHeapToStack("enable-heap-to-stack-conversion",
+                                       cl::init(true), cl::Hidden);
+static cl::opt<bool> BeAggressive("aggressive-heap-to-stack-conversion",
+                                  cl::init(false), cl::Hidden);
+static cl::opt<bool> BeVeryAggressive("very-aggressive-heap-to-stack-conversion",
+                                  cl::init(false), cl::Hidden);
+static cl::opt<int> MaxHeapToStackSize("max-heap-to-stack-size",
+                                       cl::init(1024), cl::Hidden);
+
+static bool convertAllocations(Function &F, DominatorTree &DT,
+                               LoopInfo &LI, const TargetLibraryInfo &TLI,
+                               bool &ChangedCFG) {
+  auto &DL = F.getParent()->getDataLayout();
+  bool MadeChange = false;
+  ChangedCFG = false;
+
+  if (!EnableHeapToStack)
+    return MadeChange;
+
+  // First, find allocation and deallocation functions. If we don't have at
+  // least one of both of these, then we're done.
+  SmallPtrSet<Instruction *, 4> MallocCalls, FreeCalls;
+
+  for (auto &BB : F)
+    for (auto &I : BB) {
+      if (isMallocLikeFn(&I, &TLI))
+        if (auto *CSize = dyn_cast<ConstantInt>(I.getOperand(0)))
+          if (CSize->getValue().sle(MaxHeapToStackSize)) {
+            MallocCalls.insert(&I);
+            LLVM_DEBUG(dbgs() << "H2S: Initial malloc call: " << I << "\n");
+          }
+
+      if (isFreeCall(&I, &TLI)) {
+        FreeCalls.insert(&I);
+        LLVM_DEBUG(dbgs() << "H2S: Initial free call: " << I << "\n");
+      }
+    }
+
+  if (MallocCalls.empty() || FreeCalls.empty())
+    return MadeChange;
+
+  // A map of each malloc call to the set of associated free calls.
+  DenseMap<Instruction *, SmallPtrSet<Instruction *, 4>> FreesForMalloc;
+
+  bool ShouldRefine = false;
+  do {
+    // For each call to free, the underlying objects of its arguments must be
+    // the already-collected set of calls to malloc. If not, remove the free.
+
+    FreesForMalloc.clear();
+    SmallVector<Instruction *, 4> BadFreeCalls;
+    for (auto *FreeCall : FreeCalls) {
+      SmallVector<Value *, 4> FreeUOs;
+      GetUnderlyingObjects(FreeCall->getOperand(0), FreeUOs, DL, &LI);
+
+      // Note that the call to GetUnderlyingObjects looks through pointer
+      // arithmetic, and that's okay, but we're relying on the fact that it is
+      // UB to pass an offset-adjusted pointer to free (thus, and offset
+      // applied must be reverted by the time the value gets to the call to
+      // free).
+
+      for (auto *UO : FreeUOs) {
+        if (!isa<Instruction>(UO) ||
+            !MallocCalls.count(cast<Instruction>(UO))) {
+          BadFreeCalls.push_back(FreeCall);
+          LLVM_DEBUG(dbgs() << "H2S: Bad free call: " << *FreeCall << " UO: " <<
+                               *UO << "\n");
+          break;
+        }
+
+        FreesForMalloc[cast<Instruction>(UO)].insert(FreeCall);
+      }
+    }
+
+    for (auto *BadFreeCall : BadFreeCalls)
+      FreeCalls.erase(BadFreeCall);
+
+    if (FreeCalls.empty())
+      return MadeChange;
+
+    // For each malloc call, we need to know if there's a path to the exit
+    // which does not contain one of the associated calls to free. If there is,
+    // then we can't convert this malloc.
+    SmallVector<Instruction *, 4> BadMallocCalls;
+    for (auto *MallocCall : MallocCalls) {
+      if (FreesForMalloc[MallocCall].empty()) {
+        BadMallocCalls.push_back(MallocCall);
+        LLVM_DEBUG(dbgs() << "H2S: Bad malloc call (no frees): " <<
+                             *MallocCall << "\n");
+        continue;
+      }
+
+      // Note: We need to be careful about the case where an infinite loop is
+      // followed by a free. In this case, the free is dead, and the infinite
+      // loop might allow some called function, or another thread, to free the
+      // data elsewhere. Currently, things that cause an infinite loop to be
+      // well defined (e.g., atomics, I/O) will also cause
+      // isGuaranteedToTransferExecutionToSuccessor to return false,
+      // so we're okay on that front.
+
+      struct WLEntry {
+        BasicBlock::iterator I;
+        BasicBlock *BB;
+        std::set<Instruction *> Ptrs;
+
+        WLEntry(BasicBlock::iterator I, BasicBlock *BB, Instruction *Ptr)
+          : I(I), BB(BB) {
+          Ptrs.insert(Ptr);
+        }
+
+        WLEntry(BasicBlock::iterator I, BasicBlock *BB,
+                std::set<Instruction *> &Ptrs)
+          : I(I), BB(BB), Ptrs(Ptrs) {}
+
+        bool operator == (const WLEntry &WLE) const {
+          return I == WLE.I && BB == WLE.BB && Ptrs == WLE.Ptrs;
+        }
+
+        bool operator < (const WLEntry &WLE) const {
+          if (std::less<BasicBlock *>()(BB, WLE.BB))
+            return true;
+          else if (std::less<BasicBlock *>()(WLE.BB, BB))
+            return false;
+
+          if (std::less<Instruction *>()((Instruction *) I.getNodePtr(),
+                                         (Instruction *) WLE.I.getNodePtr()))
+            return true;
+          else if (std::less<Instruction *>()((Instruction *) WLE.I.getNodePtr(),
+                                              (Instruction *) I.getNodePtr()))
+            return false;
+
+          return Ptrs < WLE.Ptrs;
+        }
+      };
+
+      SmallVector<WLEntry, 32> Worklist;
+      Worklist.push_back(WLEntry(std::next(BasicBlock::iterator(MallocCall)),
+                                 MallocCall->getParent(), MallocCall));
+
+      SmallSet<WLEntry, 32> Visited;
+
+      while (!Worklist.empty()) {
+        bool FoundExit = false, FoundFree = false;
+        WLEntry WLE = Worklist.pop_back_val();
+        if (!Visited.insert(WLE).second)
+          continue;
+
+        auto Ptrs = WLE.Ptrs;
+        for (auto I = WLE.I, IE = WLE.BB->end(); I != IE; ++I) {
+          if (isFreeCall(&*I, &TLI) &&
+              Ptrs.count((Instruction *) I->getOperand(0))) {
+            FoundFree = true;
+            LLVM_DEBUG(dbgs() << "H2S: Paired malloc call: " << *MallocCall <<
+                                  " with free at: " << *I << "\n");
+            break;
+          }
+
+          if (auto *GEPI = dyn_cast<GetElementPtrInst>(I)) {
+            if (Ptrs.count((Instruction *) GEPI->getPointerOperand()))
+              Ptrs.insert(&*I);
+          } else if (I->getOpcode() == Instruction::BitCast ||
+                     I->getOpcode() == Instruction::AddrSpaceCast) {
+            if (Ptrs.count((Instruction *) I->getOperand(0)))
+              Ptrs.insert(&*I);
+          } else if (auto CS = CallSite(&*I)) {
+            if (Ptrs.count((Instruction *) CS.getReturnedArgOperand()))
+              Ptrs.insert(&*I);
+          } else if (auto *SI = dyn_cast<SelectInst>(I)) {
+            if (Ptrs.count((Instruction *) SI->getTrueValue()) ||
+                Ptrs.count((Instruction *) SI->getFalseValue()))
+              Ptrs.insert(&*I);
+          }
+
+          if (BeAggressive) {
+            auto NotAnExit = [&](const Instruction *I) {
+              if (const auto *CRI = dyn_cast<CleanupReturnInst>(I))
+                return !CRI->unwindsToCaller();
+              if (const auto *CatchSwitch = dyn_cast<CatchSwitchInst>(I))
+                return !CatchSwitch->unwindsToCaller();
+              if (isa<ResumeInst>(I))
+                return false;
+              if (isa<ReturnInst>(I))
+                return false;
+              if (isa<UnreachableInst>(I))
+                return false;
+
+              if (auto CS = ImmutableCallSite(I))
+                if (!CS.doesNotThrow() && !BeVeryAggressive)
+                  return false;
+
+              return true;              
+            };
+
+            if (!NotAnExit(&*I))
+              FoundExit = true;
+          } else {
+            if (!isGuaranteedToTransferExecutionToSuccessor(&*I)) {
+              // Might return, etc. before we found something that might free the
+              // memory.
+              FoundExit = true;
+            }
+          }
+
+          if (FoundExit) {
+            // The exit is a problem only if the pointer is captured...
+            bool AnyCaptured = false;
+            for (auto *Ptr : Ptrs)
+              // TODO: Add an OBB cache to the capturing query.
+              if (PointerMayBeCapturedBefore(Ptr, /* ReturnCaptures */ true,
+                                             /* StoreCaptures */ true, &*I,
+                                             &DT, /* IncludeI */ true)) {
+                LLVM_DEBUG(dbgs() << "H2S: Pointer: " << *Ptr <<
+                                     " from malloc call: " << *MallocCall <<
+                                     " captured before, or at, "
+                                     "the potential exit at: " << *I << "\n");
+                AnyCaptured = true;
+                break;
+              }
+
+            if (!AnyCaptured) {
+              FoundExit = false;
+            } else {
+              LLVM_DEBUG(dbgs() << "H2S: Bad malloc call: " << *MallocCall <<
+                                   " found potential exit at: " << *I << "\n");
+              break;
+            }
+          }
+        }
+
+        if (FoundExit) {
+          BadMallocCalls.push_back(MallocCall);
+          break;
+        } else if (!FoundFree) {
+          for (auto *SBB : successors(WLE.BB)) {
+            auto SPtrs = Ptrs;
+            for (PHINode &PN : SBB->phis())
+              if (Ptrs.count((Instruction *)
+                               PN.getIncomingValueForBlock(WLE.BB)))
+                SPtrs.insert(&PN);
+
+            Worklist.push_back(WLEntry(SBB->begin(), SBB, SPtrs));
+          }
+        }
+      }
+    }
+
+    if (!BadMallocCalls.empty())
+      ShouldRefine = true;
+
+    for (auto *BadMallocCall : BadMallocCalls)
+      MallocCalls.erase(BadMallocCall);
+
+    if (MallocCalls.empty())
+      return MadeChange;
+  } while (ShouldRefine);
+
+  MadeChange = true;
+
+  for (auto *MallocCall : MallocCalls) {
+    LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall << "\n");
+
+    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);
+
+    if (AI->getType() != MallocCall->getType()) {
+      auto *BC = new BitCastInst(AI, MallocCall->getType());
+      BC->insertAfter(AI);
+      AI = BC;
+    }
+
+    MallocCall->replaceAllUsesWith(AI);
+
+    if (auto *II = dyn_cast<InvokeInst>(MallocCall)) {
+      auto *NBB = II->getNormalDest();
+      BranchInst::Create(NBB, MallocCall->getParent());
+      MallocCall->eraseFromParent();
+      ChangedCFG = true;
+    } else {
+      MallocCall->eraseFromParent();
+    }
+
+    ++NumAConverted;
+  }
+
+  for (auto *FreeCall : FreeCalls) {
+    LLVM_DEBUG(dbgs() << "H2S: Removing free call: " << *FreeCall << "\n");
+
+    FreeCall->eraseFromParent();
+    ++NumDConverted;
+  }
+
+  return MadeChange;
+}
+
+PreservedAnalyses HeapToStackPass::run(Function &F, FunctionAnalysisManager &AM) {
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
+  auto &LI = AM.getResult<LoopAnalysis>(F);
+  auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
+
+  bool ChangedCFG;
+  if (!convertAllocations(F, DT, LI, TLI, ChangedCFG))
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  if (!ChangedCFG)
+    PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+
+namespace {
+  class HeapToStackLegacyPass : public FunctionPass {
+  public:
+    static char ID; // Pass identification
+    HeapToStackLegacyPass() : FunctionPass(ID) {
+      initializeHeapToStackLegacyPassPass(*PassRegistry::getPassRegistry());
+    }
+
+    bool runOnFunction(Function &F) override {
+      auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+      auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+      auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
+
+      bool ChangedCFG;
+      return convertAllocations(F, DT, LI, TLI, ChangedCFG);
+    }
+
+    void getAnalysisUsage(AnalysisUsage &AU) const override {
+      FunctionPass::getAnalysisUsage(AU);
+      AU.addRequired<TargetLibraryInfoWrapperPass>();
+      AU.addRequired<DominatorTreeWrapperPass>();
+      AU.addRequired<LoopInfoWrapperPass>();
+      AU.addPreserved<DominatorTreeWrapperPass>();
+      AU.addPreserved<LoopInfoWrapperPass>();
+    }
+  };
+} // end anonymous namespace
+
+char HeapToStackLegacyPass::ID = 0;
+static const char H2SName[] = "Heap-to-Stack Conversion";
+
+INITIALIZE_PASS_BEGIN(HeapToStackLegacyPass, DEBUG_TYPE, H2SName, false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_END(HeapToStackLegacyPass, DEBUG_TYPE, H2SName, false, false)
+
+FunctionPass *llvm::createHeapToStackPass() { return new HeapToStackLegacyPass(); }

lib/Transforms/Scalar/Scalar.cpp

void llvm::initializeScalarOpts(PassRegistry &Registry) {
   initializeSinkingLegacyPassPass(Registry);
   initializeTailCallElimPass(Registry);
   initializeSeparateConstOffsetFromGEPPass(Registry);
   initializeSpeculativeExecutionLegacyPassPass(Registry);
   initializeStraightLineStrengthReducePass(Registry);
   initializePlaceBackedgeSafepointsImplPass(Registry);
   initializePlaceSafepointsPass(Registry);
   initializeFloat2IntLegacyPassPass(Registry);
+  initializeHeapToStackLegacyPassPass(Registry);
   initializeLoopDistributeLegacyPass(Registry);
   initializeLoopLoadEliminationPass(Registry);
   initializeLoopSimplifyCFGLegacyPassPass(Registry);
   initializeLoopVersioningPassPass(Registry);
   initializeEntryExitInstrumenterPass(Registry);
   initializePostInlineEntryExitInstrumenterPass(Registry);
 }