[InstructionCombining] Replace small heap allocations with local variables

[xbolva00]

Summary

Requires constant size to be allocated

After inlining we may got code like this:

int main(void) {

char *ptr = malloc(80* sizeof(char)); // SIZE

strcpy(ptr, "test");
puts(ptr);

free(ptr);

}

if SIZE < threshold (currently 256)

-> create something like "char ptr[size]"

Fires 30+ in MultiSource benchmark.

Tests would be added later, after an initial discussion.

[efriedma]

The biggest problem I see with this is the potential to cause a stack overflow; you've specified a limit of 100 bytes per allocation, but there's no limit to the number of allocations. Not sure what the right solution is; LLVM doesn't really have good heuristics here.

lib/Transforms/InstCombine/InstructionCombining.cpp

2278

Please make this an option, so it can be tweaked for debugging.

2310

This isn't sufficient to guarantee malloc will be called once; getLoopFor won't find irreducible loops.

2325

You probably don't want to insert an alloca in the middle of a function; allocas outside the entry block are treated as dynamic allocations, so they're slower, and won't work the way you want with code that uses llvm.stackrestore.

[xbolva00]

Maybe introduce local cache to track:
Function --- Number of malloc -> alloca?

Or allow it only once per function? Or again some treshold value.

lib/Transforms/InstCombine/InstructionCombining.cpp

2278

Command line option for opt? Ok.

2325

Ok, I will fix it.

[xbolva00]

lib/Transforms/InstCombine/InstructionCombining.cpp

2310

How can we fix it and detect them?

I read some comments about irreducible loops in getNearestLoop...

[xbolva00]

Addressed review comments

[efriedma]

lib/Transforms/InstCombine/InstructionCombining.cpp

2277

You can't use a static variable like this; LLVM can run on multiple threads in a program at the same time.

2310

There's a helper containsIrreducibleCFG.

[xbolva00]

• Detect IrreducibleCFG

[xbolva00]

Maybe the code is ready?

So I can create some tests.

[efriedma]

lib/Transforms/InstCombine/InstCombineInternal.h

256

This still doesn't really do what you want: opt -instcombine -instcombine will ignore the limit (and we effectively do that in a lot of places). Maybe you could solve that by moving the transform into its own pass? Not sure; I'll think about it over the weekend and get back to you.

lib/Transforms/InstCombine/InstructionCombining.cpp

2307

This check isn't sufficient to ensure some other function doesn't call free() on the pointer (actually, this doesn't even handle the case where there are two calls to free() along different codepaths).

[xbolva00]

lib/Transforms/InstCombine/InstCombineInternal.h

256

Sure, thanks

lib/Transforms/InstCombine/InstructionCombining.cpp

2307

Other function...

void free_ptr(void **p);
void capture(void*p);

void small_alloc_3(void) {

void *s = malloc(10);
capture(s);
free(s);

}

void small_alloc_4(void) {

void *s = malloc(10);
free_ptr(&s);
free(s);

}

These case are not optimized (PointerMayBeCaptured solves it)

2307

But yes, the second case is a problem.. Not sure what to do.

void small_alloc_bad2(int n, int i) {

char *s = malloc(10);
puts(s);
if (n)
        free(s);

if (i)
        free(s);

}

[xbolva00]

lib/Transforms/InstCombine/InstructionCombining.cpp

2307

Check uses of "s"? If there is only one "free" ?

[efriedma]

lib/Transforms/InstCombine/InstructionCombining.cpp

2307

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.

[xbolva00]

• Check number of frees

[xbolva00]

lib/Transforms/InstCombine/InstructionCombining.cpp

2307

Yes hmm. Maybe extend CaptureTracker with to handle this with a new parameter bool ArgCaptures?

[xbolva00]

lib/Transforms/InstCombine/InstructionCombining.cpp

2276

Or maybe we can turn larger allocations into global variables also? + allocations with unknown (not constant) size?

2294

this looks bad, any advice how to do it properly?

Maybe just remove all free calls in Malloc's uses?

[xbolva00]

lib/Transforms/InstCombine/InstructionCombining.cpp

2276

Ah, possible data race on that global var..

[xbolva00]

Friendly ping :)

lib/Transforms/InstCombine/InstCombineInternal.h

 /// This file provides internal interfaces used to implement the InstCombine.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TRANSFORMS_INSTCOMBINE_INSTCOMBINEINTERNAL_H
 #define LLVM_LIB_TRANSFORMS_INSTCOMBINE_INSTCOMBINEINTERNAL_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetFolder.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
private:
   OptimizationRemarkEmitter &ORE;
 
   // Optional analyses. When non-null, these can both be used to do better
   // combining and will be updated to reflect any changes.
   LoopInfo *LI;
 
   bool MadeIRChange = false;
 
+  MapVector<Function *, unsigned> HeapAllocEliminations;
+
 public:
   InstCombiner(InstCombineWorklist &Worklist, BuilderTy &Builder,
                bool MinimizeSize, bool ExpensiveCombines, AliasAnalysis *AA,
                AssumptionCache &AC, TargetLibraryInfo &TLI, DominatorTree &DT,
                OptimizationRemarkEmitter &ORE, const DataLayout &DL,
                LoopInfo *LI)
       : Worklist(Worklist), Builder(Builder), MinimizeSize(MinimizeSize),
         ExpensiveCombines(ExpensiveCombines), AA(AA), AC(AC), TLI(TLI), DT(DT),
public:
   Instruction *visitExtractValueInst(ExtractValueInst &EV);
   Instruction *visitLandingPadInst(LandingPadInst &LI);
   Instruction *visitVAStartInst(VAStartInst &I);
   Instruction *visitVACopyInst(VACopyInst &I);
 
   /// Specify what to return for unhandled instructions.
   Instruction *visitInstruction(Instruction &I) { return nullptr; }
 
+  /// Eliminate small allocations
+  bool eliminateSmallAllocations(Value *Op, CallInst &FI);
+
   /// True when DB dominates all uses of DI except UI.
   /// UI must be in the same block as DI.
   /// The routine checks that the DI parent and DB are different.
   bool dominatesAllUses(const Instruction *DI, const Instruction *UI,
                         const BasicBlock *DB) const;
 
   /// Try to replace select with select operand SIOpd in SI-ICmp sequence.
   bool replacedSelectWithOperand(SelectInst *SI, const ICmpInst *Icmp,

lib/Transforms/InstCombine/InstructionCombining.cpp

 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
 #include "llvm-c/Initialization.h"
 #include "llvm-c/Transforms/InstCombine.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/None.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetFolder.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"

 
 STATISTIC(NumCombined , "Number of insts combined");
 STATISTIC(NumConstProp, "Number of constant folds");
 STATISTIC(NumDeadInst , "Number of dead inst eliminated");
 STATISTIC(NumSunkInst , "Number of instructions sunk");
 STATISTIC(NumExpand,    "Number of expansions");
 STATISTIC(NumFactor   , "Number of factorizations");
 STATISTIC(NumReassoc  , "Number of reassociations");
+STATISTIC(NumMallocElim, "Number of mallocs eliminated");
 DEBUG_COUNTER(VisitCounter, "instcombine-visit",
               "Controls which instructions are visited");
 
 static cl::opt<bool>
 EnableExpensiveCombines("expensive-combines",
                         cl::desc("Enable expensive instruction combines"));
 
 static cl::opt<unsigned>
 MaxArraySize("instcombine-maxarray-size", cl::init(1024),
              cl::desc("Maximum array size considered when doing a combine"));
 
 // FIXME: Remove this flag when it is no longer necessary to convert
 // llvm.dbg.declare to avoid inaccurate debug info. Setting this to false
 // increases variable availability at the cost of accuracy. Variables that
 // cannot be promoted by mem2reg or SROA will be described as living in memory
 // for their entire lifetime. However, passes like DSE and instcombine can
 // delete stores to the alloca, leading to misleading and inaccurate debug
 // information. This flag can be removed when those passes are fixed.
 static cl::opt<unsigned> ShouldLowerDbgDeclare("instcombine-lower-dbg-declare",
                                                cl::Hidden, cl::init(true));
 
+static cl::opt<unsigned>
+    HeapAllocSize("heap-alloc-size", cl::init(256),
+                 cl::desc("Maximum size of heap allocation to be replaced with alloca"));
+
+static cl::opt<unsigned>
+    HeapAllocCount("heap-alloc-count", cl::init(1),
+                  cl::desc("Maximum count of heap allocation to be replaced "
+                           "with alloca per function"));
+
 Value *InstCombiner::EmitGEPOffset(User *GEP) {
   return llvm::EmitGEPOffset(&Builder, DL, GEP);
 }
 
 /// Return true if it is desirable to convert an integer computation from a
 /// given bit width to a new bit width.
 /// We don't want to convert from a legal to an illegal type or from a smaller
 /// to a larger illegal type. A width of '1' is always treated as a legal type
if (SuccBB != (Pred == ICmpInst::ICMP_EQ ? TrueBB : FalseBB))
     return nullptr;
   assert(FreeInstrBB == (Pred == ICmpInst::ICMP_EQ ? FalseBB : TrueBB) &&
          "Broken CFG: missing edge from predecessor to successor");
 
   FI.moveBefore(TI);
   return &FI;
 }
 
+bool InstCombiner::eliminateSmallAllocations(Value *Op, CallInst &FI) {
+  Function *F = FI.getFunction();
+
+  Function *FreeFn = FI.getCalledFunction();
+  if (!FreeFn)
+    return false;
+
+  LibFunc Func;
+  if (!TLI.getLibFunc(*FreeFn, Func) || !TLI.has(Func))
+    return false;
+
+  if (Func != LibFunc_free)
+    return false;
+
+  CallInst *Malloc = dyn_cast<CallInst>(Op);
+  if (!Malloc)
+    return false;
+
+  bool FoundFree = false;
+  for (User *I : Malloc->users()) {
+    auto * CI = dyn_cast<CallInst>(I);
+    if (CI && CI->getCalledFunction() == FI.getCalledFunction()) {
+      // Check for free on multiple code paths
+      if (FoundFree)
+        return false;
+      FoundFree = true;
+    }
+  }
+
+  // Is the inner call really malloc?
+  Function *MallocFn = Malloc->getCalledFunction();
+  if (!MallocFn)
+    return false;
+
+  if (!TLI.getLibFunc(*MallocFn, Func) || !TLI.has(Func))
+    return false;
+
+  if (Func != LibFunc_malloc)
+    return false;
+
+  // Check if alloced pointer is captured
+  if (PointerMayBeCaptured(Malloc, true, true))
+    return false;
+
+  BasicBlock *BB = Malloc->getParent();
+  // Do not optimize if malloc is in loop
+  if (LI->getLoopFor(BB))
+    return false;
+
+  ReversePostOrderTraversal<BasicBlock *> RPOT(&*BB);
+  if (containsIrreducibleCFG<BasicBlock *>(RPOT, *LI))
+    return false;
+
+  // Check for constant size
+  ConstantInt *Size = dyn_cast<ConstantInt>(Malloc->getOperand(0));
+  if (!Size)
+    return false;
+
+  if (HeapAllocEliminations[F] == HeapAllocCount)
+    return false;
+
+  uint64_t N = Size->getZExtValue();
+  if (N > HeapAllocSize)
+    return false;
+
+  Builder.SetInsertPoint(&F->getEntryBlock(), F->getEntryBlock().begin());
+  Value *LocAlloc = Builder.CreateAlloca(Builder.getInt8Ty(), Size);
+  Malloc->replaceAllUsesWith(LocAlloc);
+  Malloc->eraseFromParent();
+
+  ++NumMallocElim;
+  ++HeapAllocEliminations[F];
+
+  return true;
+}
+
 Instruction *InstCombiner::visitFree(CallInst &FI) {
   Value *Op = FI.getArgOperand(0);
 
+  // Small "malloc"ation -> create new local variable
+  if (eliminateSmallAllocations(Op, FI))
+    return eraseInstFromFunction(FI);
+
   // free undef -> unreachable.
   if (isa<UndefValue>(Op)) {
     // Insert a new store to null because we cannot modify the CFG here.
     Builder.CreateStore(ConstantInt::getTrue(FI.getContext()),
                         UndefValue::get(Type::getInt1PtrTy(FI.getContext())));
     return eraseInstFromFunction(FI);
   }