From c4499e333337ef67f30d81eefd044816a8cf6f29 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Sat, 28 Mar 2020 12:45:59 +0530
Subject: [PATCH] [Attributor] Make attributor aware of aligned_alloc for heap
 to stack conversion

Make the attributor pass aware of aligned_alloc for converting heap
allocations to stack ones.

Depends on D76971.

Differential Revision: https://reviews.llvm.org/D76974
---
 llvm/lib/Transforms/IPO/Attributor.cpp           | 27 +++++++++++++++++++-----
 llvm/test/Transforms/Attributor/heap_to_stack.ll | 24 +++++++++++++++++++--
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 02fff40..8877a5f 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5030,6 +5030,7 @@ struct AAHeapToStackImpl : public AAHeapToStack {
       LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall
                         << "\n");
 
+      MaybeAlign Alignment;
       Constant *Size;
       if (isCallocLikeFn(MallocCall, TLI)) {
         auto *Num = cast<ConstantInt>(MallocCall->getOperand(0));
@@ -5037,13 +5038,19 @@ struct AAHeapToStackImpl : public AAHeapToStack {
         APInt TotalSize = SizeT->getValue() * Num->getValue();
         Size =
             ConstantInt::get(MallocCall->getOperand(0)->getType(), TotalSize);
+      } else if (isAlignedAllocLikeFn(MallocCall, TLI)) {
+        Size = cast<ConstantInt>(MallocCall->getOperand(1));
+        Alignment = MaybeAlign(cast<ConstantInt>(MallocCall->getOperand(0))
+                                   ->getValue()
+                                   .getZExtValue());
       } else {
         Size = cast<ConstantInt>(MallocCall->getOperand(0));
       }
 
       unsigned AS = cast<PointerType>(MallocCall->getType())->getAddressSpace();
-      Instruction *AI = new AllocaInst(Type::getInt8Ty(F->getContext()), AS,
-                                       Size, "", MallocCall->getNextNode());
+      Instruction *AI =
+          new AllocaInst(Type::getInt8Ty(F->getContext()), AS, Size, Alignment,
+                         "", MallocCall->getNextNode());
 
       if (AI->getType() != MallocCall->getType())
         AI = new BitCastInst(AI, MallocCall->getType(), "malloc_bc",
@@ -5175,8 +5182,9 @@ ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
       return true;
 
     bool IsMalloc = isMallocLikeFn(&I, TLI);
+    bool IsAlignedAllocLike = isAlignedAllocLikeFn(&I, TLI);
     bool IsCalloc = !IsMalloc && isCallocLikeFn(&I, TLI);
-    if (!IsMalloc && !IsCalloc) {
+    if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc) {
       BadMallocCalls.insert(&I);
       return true;
     }
@@ -5188,6 +5196,14 @@ ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
             MallocCalls.insert(&I);
             return true;
           }
+    } else if (IsAlignedAllocLike && isa<ConstantInt>(I.getOperand(0))) {
+      // Only if the alignment and sizes are constant.
+      if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
+        if (Size->getValue().ule(MaxHeapToStackSize))
+          if (UsesCheck(I) || FreeCheck(I)) {
+            MallocCalls.insert(&I);
+            return true;
+          }
     } else if (IsCalloc) {
       bool Overflow = false;
       if (auto *Num = dyn_cast<ConstantInt>(I.getOperand(0)))
@@ -5219,8 +5235,9 @@ struct AAHeapToStackFunction final : public AAHeapToStackImpl {
 
   /// See AbstractAttribute::trackStatistics().
   void trackStatistics() const override {
-    STATS_DECL(MallocCalls, Function,
-               "Number of malloc calls converted to allocas");
+    STATS_DECL(
+        MallocCalls, Function,
+        "Number of malloc/calloc/aligned_alloc calls converted to allocas");
     for (auto *C : MallocCalls)
       if (!BadMallocCalls.count(C))
         ++BUILD_STAT_NAME(MallocCalls, Function);
diff --git a/llvm/test/Transforms/Attributor/heap_to_stack.ll b/llvm/test/Transforms/Attributor/heap_to_stack.ll
index 8dc9b94..369a661 100644
--- a/llvm/test/Transforms/Attributor/heap_to_stack.ll
+++ b/llvm/test/Transforms/Attributor/heap_to_stack.ll
@@ -76,6 +76,26 @@ define void @test3a(i8* %p) {
   ret void
 }
 
+declare noalias i8* @aligned_alloc(i64, i64)
+
+define void @test3b(i8* %p) {
+  %1 = tail call noalias i8* @aligned_alloc(i64 32, i64 128)
+  ; CHECK: %1 = alloca i8, i64 128, align 32
+  ; CHECK-NEXT: tail call void @nofree_arg_only
+  tail call void @nofree_arg_only(i8* %1, i8* %p)
+  ; CHECK-NOT: @free(i8* %1)
+  tail call void @free(i8* %1)
+  ret void
+}
+
+; leave alone non-constant alignments.
+define void @test3c(i64 %alignment) {
+  %1 = tail call noalias i8* @aligned_alloc(i64 %alignment, i64 128)
+  ; CHECK: tail call noalias i8* @aligned_alloc
+  tail call void @free(i8* %1)
+  ret void
+}
+
 declare noalias i8* @calloc(i64, i64)
 
 define void @test0() {
@@ -90,7 +110,7 @@ define void @test0() {
   ret void
 }
 
-; TEST 4 
+; TEST 4
 define void @test4() {
   %1 = tail call noalias i8* @malloc(i64 4)
   ; CHECK: %1 = alloca i8, i64 4
@@ -219,7 +239,7 @@ define i32 @test_lifetime() {
   ret i32 %3
 }
 
-; TEST 11 
+; TEST 11
 
 define void @test11() {
   %1 = tail call noalias i8* @malloc(i64 4)
-- 
2.7.4