Use nofree argument attributr hor heap-to-stack conversion

[sstefan1]

Summary

Use nofree argument attribute for heap-to-stack conversion.

[jdoerfert]

LGTM, one nit.

llvm/lib/Transforms/IPO/Attributor.cpp

4159

NOTE can be removed as well, I think.

[arsenm]

Typo hor in commit message

[jdoerfert]

In D70140#1743338, @arsenm wrote:

Typo hor in commit message

Two typos even. Thanks for noticing!

[sstefan1]

In D70140#1743354, @jdoerfert wrote:

In D70140#1743338, @arsenm wrote:

Typo hor in commit message

Two typos even. Thanks for noticing!

Thanks! Will double check before I commit

llvm/lib/Transforms/IPO/Attributor.cpp

4159

I'll remove it before I commit.

llvm/lib/Transforms/IPO/Attributor.cpp

while (!Worklist.empty()) {
                      << "[H2S] escaping store to memory: " << *UserI << "\n");
           ValidUsesOnly = false;
         } else {
           // A store into the malloc'ed memory is fine.
         }
         continue;
       }
 
-      // NOTE: Right now, if a function that has malloc pointer as an argument
-      // frees memory, we assume that the malloc pointer is freed.
-
-      // TODO: Add nofree callsite argument attribute to indicate that pointer
-      // argument is not freed.
       if (auto *CB = dyn_cast<CallBase>(UserI)) {
         if (!CB->isArgOperand(U))
           continue;
 
         if (CB->isLifetimeStartOrEnd())
           continue;
 
         // Record malloc.
         if (isFreeCall(UserI, TLI)) {
           if (MustUse) {
             FreesForMalloc[&I].insert(
                 cast<Instruction>(const_cast<User *>(UserI)));
           } else {
             LLVM_DEBUG(dbgs() << "[H2S] free potentially on different mallocs: "
                               << *UserI << "\n");
             ValidUsesOnly = false;
           }
           continue;
         }
 
-        // If a function does not free memory we are fine
-        const auto &NoFreeAA =
-            A.getAAFor<AANoFree>(*this, IRPosition::callsite_function(*CB));
-
         unsigned ArgNo = CB->getArgOperandNo(U);
+
         const auto &NoCaptureAA = A.getAAFor<AANoCapture>(
             *this, IRPosition::callsite_argument(*CB, ArgNo));
 
-        if (!NoCaptureAA.isAssumedNoCapture() || !NoFreeAA.isAssumedNoFree()) {
+        // If a callsite argument use is nofree, we are fine.
+        const auto &ArgNoFreeAA = A.getAAFor<AANoFree>(
+            *this, IRPosition::callsite_argument(*CB, ArgNo));
+
+        if (!NoCaptureAA.isAssumedNoCapture() || !ArgNoFreeAA.isAssumedNoFree()) {
           LLVM_DEBUG(dbgs() << "[H2S] Bad user: " << *UserI << "\n");
           ValidUsesOnly = false;
         }
         continue;
       }
 
       if (isa<GetElementPtrInst>(UserI) || isa<BitCastInst>(UserI) ||
           isa<PHINode>(UserI) || isa<SelectInst>(UserI)) {

llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll

 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
 
+; CHECK: @nofree_arg_only(i8* nocapture nofree %p1, i8* nocapture %p2)
+define void @nofree_arg_only(i8* %p1, i8* %p2) {
+  tail call void @free(i8* %p2)
+  tail call void @nofree_func(i8* %p1)
+  ret void
+}
+
 ; 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* noalias nocapture %1)
   tail call void @nocapture_func_frees_pointer(i8* %1)
   tail call void (...) @func_throws()
define void @test3() {
   ; CHECK: %1 = alloca i8, i64 4
   ; 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
 }
 
+define void @test3a(i8* %p) {
+  %1 = tail call noalias i8* @malloc(i64 4)
+  ; 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* @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)
define void @test4() {
   ; CHECK-NEXT: @nofree_func(i8* noalias nocapture nofree %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) {
+define void @test5(i32, i8* %p) {
   %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 @nofree_arg_only(i8* %2, i8* %p)
   tail call void @free(i8* %2)
   ; CHECK-NOT: @free(i8* %2)
   br label %6
 
 6:                                                ; preds = %5, %4
   ret void
 }