[Attributor] Make attributor aware of aligned_alloc for heap to stack conversion

[bondhugula]

Summary

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

Signed-off-by: Uday Bondhugula <[e-mail]>

Depends on D76971.

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

5053

Nit: Alignment

5206

Why a signed comparison?

[bondhugula]

llvm/lib/Transforms/IPO/Attributor.cpp

5206

Fixed, thanks.

[bondhugula]

Address review comments.

[bondhugula]

Ping reviewers - @jdoerfert - thanks!

[lebedev.ri]

llvm/lib/Transforms/IPO/Attributor.cpp

5033–5040

Hmm, what is the 'default' alignment for alloca?
Even for non-aligned_alloc, there's some alignment guarantees for the allocation functions..

[bondhugula]

llvm/lib/Transforms/IPO/Attributor.cpp

5033–5040

For alloca, by default the target will typically align to sizeof elt type boundary (per LangRef). For alloc functions without alignment args, it's all system dependent. Glibc's malloc will provide alignment in line with the largest primitive (non-vector) type supported on the system which is typically 16 bytes (fp80, fp128 are the largest).

[xbolva00]

llvm/lib/Transforms/IPO/Attributor.cpp

5033–5040

we could detect if GNU environment (glibc) is used for compilation so I am wondering if we can annotate malloc with such alignment (16B) if isGNU() == true.

[bondhugula]

llvm/lib/Transforms/IPO/Attributor.cpp

5033–5040

Even if it's the GNU env, the 16 byte alignment isn't guaranteed by malloc (although nearly always true on today's systems). It could be 8 bytes on really old ones - I don't know!

[lebedev.ri]

llvm/lib/Transforms/IPO/Attributor.cpp

5033–5040

While yes, we should be doing that, llvm is really too late for this, it must be done in clang, similarly to D73380.

[jdoerfert]

LGTM, @lebedev.ri @xbolva00 ?

[lebedev.ri]

No comments from me.

[xbolva00]

Ok too

llvm/lib/Transforms/IPO/Attributor.cpp

for (Instruction *MallocCall : MallocCalls) {
         LLVM_DEBUG(dbgs() << "H2S: Removing free call: " << *FreeCall << "\n");
         A.deleteAfterManifest(*FreeCall);
         HasChanged = ChangeStatus::CHANGED;
       }
 
       LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall
                         << "\n");
 
+      MaybeAlign Alignment;
       Constant *Size;
       if (isCallocLikeFn(MallocCall, TLI)) {
         auto *Num = cast<ConstantInt>(MallocCall->getOperand(0));
         auto *SizeT = cast<ConstantInt>(MallocCall->getOperand(1));
         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",
                              AI->getNextNode());
 
       A.changeValueAfterManifest(*MallocCall, *AI);
 
       if (auto *II = dyn_cast<InvokeInst>(MallocCall)) {
auto UsesCheck = [&](Instruction &I) {
     return ValidUsesOnly;
   };
 
   auto MallocCallocCheck = [&](Instruction &I) {
     if (BadMallocCalls.count(&I))
       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;
     }
 
     if (IsMalloc) {
       if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (Size->getValue().ule(MaxHeapToStackSize))
           if (UsesCheck(I) || FreeCheck(I)) {
             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)))
         if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
           if ((Size->getValue().umul_ov(Num->getValue(), Overflow))
                   .ule(MaxHeapToStackSize))
             if (!Overflow && (UsesCheck(I) || FreeCheck(I))) {
               MallocCalls.insert(&I);
ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
   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 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);
   }
 };
 
 /// ----------------------- Privatizable Pointers ------------------------------
 struct AAPrivatizablePtrImpl : public AAPrivatizablePtr {

llvm/test/Transforms/Attributor/heap_to_stack.ll

define void @test3a(i8* %p) {
   ; CHECK: %1 = alloca i8, i64 4
   ; 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
 }
 
+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() {
   %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* noalias nocapture nofree %1)
   tail call void @no_sync_func(i8* %1)
   ; CHECK-NOT: @free(i8* %1)
   tail call void @free(i8* %1)
   ret void
 }
 
-; TEST 4 
+; TEST 4
 define void @test4() {
   %1 = tail call noalias i8* @malloc(i64 4)
   ; CHECK: %1 = alloca i8, i64 4
   ; CHECK-NEXT: @nofree_func(i8* noalias nocapture nofree %1)
   tail call void @nofree_func(i8* %1)
   ret void
 }
 
define i32 @test_lifetime() {
   %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
 }
 
-; TEST 11 
+; TEST 11
 
 define void @test11() {
   %1 = tail call noalias i8* @malloc(i64 4)
   ; CHECK: test11
   ; CHECK-NEXT: alloc
   ; CHECK-NEXT: @sync_will_return(i8* %1)
   tail call void @sync_will_return(i8* %1)
   tail call void @free(i8* %1)