[TRE] Do not mark alloca's as escaped if the function only writes to argmem

[caojoshua]

Summary

TRE currently reasons that if a function does not write to any memory, a
passed alloca cannot escape. We strengthen this by not escaping for functions
that write only to argmem memory. An alloca is escaped if its address is
written to a location outside of the function, such as a global. If the
function only writes to arguments, there is no escaping. For example:

unsigned *foo(unsigned *n) {
    *n = 12;
    return n;
}

foo() writes to memory, but only to argmem memory. If n is a passed
alloca, it does not escape.

[caojoshua]

Update commit message

[caojoshua]

As a side note, while looking through the code, I didn't understand this:

// If it can write to memory, it can leak the alloca value.

Even if an alloca is written to in a called function, I don't see why the callsite needs to be escaped. I also don't see why we care about the alloca's value, we only care about the pointer to the alloca. I think maybe the comment is concerned about the called function writing copying the alloca pointer to a global for example, but that we should not be concerned because we already checked for nocapture

[efriedma]

A call cannot be marked tail if it may read or write to an alloca in the same function. There are basically two ways to to "read or write": either the pointer is passed as an argument, or it's captured.

For the function @test16 in the given testcase, neither call can be marked tail. The call to @use_readonly can obviously read the alloca. The call to @use_readonly can also capture the address of the alloca, so the call to @noarg can also read the alloca.

[caojoshua]

Update this patch to not mark alloca's as escaped even if their call does not write to memory.

I added some more pre-committed tests for memory(none|read) functions. After this patch, they are no longer marked as tail calls

[efriedma]

There are basically two ways a pointer can be "captured": it can be written to memory, or it can be returned.

The code here is trying to reason more deeply about memory(none) functions: the memory(none) implies no writes to memory, so the only way a pointer can be captured is by returning it. The code follows the return value to figure out if it's actually used. So the tail markings in test17 and test18 should be fine, I think.

[caojoshua]

In D149563#4310611, @efriedma wrote:

There are basically two ways a pointer can be "captured": it can be written to memory, or it can be returned.

The code here is trying to reason more deeply about memory(none) functions: the memory(none) implies no writes to memory, so the only way a pointer can be captured is by returning it. The code follows the return value to figure out if it's actually used. So the tail markings in test17 and test18 should be fine, I think.

OK, I think I'm starting to understand this better. Looking at the function attributes doc, I think we can slightly strengthen this to not escape for memory(read, argmem: readwrite). Right not, we always escape if the callee can write to memory, but we actually don't need to escape if the callee is only writing to argmems. Please see these examples: https://godbolt.org/z/qzj76rhah. Does this make sense?

[caojoshua]

Rework the patch and commit message completey based on discussion. Don't escape allocas that write only to argmems.

[efriedma]

I don't understand the distinction you're trying to draw for argument vs. non-argument memory. Writing a pointer into memory passed as an argument counts as "captured" for the purpose of this analysis; it doesn't track memory, only values in registers.

[caojoshua]

Stepping back, the purpose of the analysis is to determine which allocas 'escape'. If an alloca does not escape, we cannot mark a call as a tail call. The definition of escaped is not defined in the source. For the purpose of deciding which calls can be tail called, I would say an alloca is escaped if its address is accessible from outside of the function, i.e. a global. In this case, we would not be able to make a tail call, because the alloca would be lost when the caller's frame is overwritten.

In general, I think there is a lot of missed opportunities for tail calls in TRE. In this example:

unsigned *foo(unsigned *n) { // define dso_local ptr @foo(ptr noundef returned writeonly %n) memory(argmem: write)
    *n = 12;
    return n;
}

void bar() {
    unsigned n; // lets say this is an alloca, and does not become a register
    foo(&n);
    could_be_tail_called();

n is captured when calling foo() because n is returned. Since foo() writes to memory, TRE marks n as escaped, and no functions can be marked as a tail call. Looking at the function, I don't see how the n alloca escapes, and it should not prevent any tail calls.

I think my current assumptions regarding function attributes in this patch are incorrect. I'll abandon this patch. If I revisit this idea, I'll back it up with data so I can make a stronger case.

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

while (!Worklist.empty()) {
       case Instruction::Invoke: {
         auto &CB = cast<CallBase>(*I);
         // If the alloca-derived argument is passed byval it is not an escape
         // point, or a use of an alloca. Calling with byval copies the contents
         // of the alloca into argument registers or stack slots, which exist
         // beyond the lifetime of the current frame.
         if (CB.isArgOperand(U) && CB.isByValArgument(CB.getArgOperandNo(U)))
           continue;
-        bool IsNocapture =
-            CB.isDataOperand(U) && CB.doesNotCapture(CB.getDataOperandNo(U));
-        callUsesLocalStack(CB, IsNocapture);
-        if (IsNocapture) {
+        AllocaUsers.insert(&CB);
+        if (!CB.isDataOperand(U) || CB.doesNotCapture(CB.getDataOperandNo(U)))
           // If the alloca-derived argument is passed in as nocapture, then it
           // can't propagate to the call's return. That would be capturing.
           continue;
-        }
+        if ((CB.getMemoryEffects().getModRef(MemoryEffects::ArgMem) &
+             ModRefInfo::Mod) != ModRefInfo::NoModRef)
+          // The alloca escapes if the function writes to non argument memory,
+          // such as globals.
+          EscapePoints.insert(&CB);
         break;
       }
       case Instruction::Load: {
         // The result of a load is not alloca-derived (unless an alloca has
         // otherwise escaped, but this is a local analysis).
         continue;
       }
       case Instruction::Store: {
while (!Worklist.empty()) {
         EscapePoints.insert(I);
         break;
       }
 
       AddUsesToWorklist(I);
     }
   }
 
-  void callUsesLocalStack(CallBase &CB, bool IsNocapture) {
-    // Add it to the list of alloca users.
-    AllocaUsers.insert(&CB);
-
-    // If it's nocapture then it can't capture this alloca.
-    if (IsNocapture)
-      return;
-
-    // If it can write to memory, it can leak the alloca value.
-    if (!CB.onlyReadsMemory())
-      EscapePoints.insert(&CB);
-  }
-
   SmallPtrSet<Instruction *, 32> AllocaUsers;
   SmallPtrSet<Instruction *, 32> EscapePoints;
 };
 }
 
 static bool markTails(Function &F, OptimizationRemarkEmitter *ORE) {
   if (F.callsFunctionThatReturnsTwice())
     return false;

llvm/test/Transforms/TailCallElim/function_readwrite_attrs.ll

entry:
   ret void
 }
 
 define void @test_use_argmem_none() {
 ; CHECK-LABEL: define void @test_use_argmem_none() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    call void @use_memory_argmem_none(ptr [[X]])
-; CHECK-NEXT:    call void @noarg()
+; CHECK-NEXT:    tail call void @noarg()
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %x = alloca i32, align 4
   call void @use_memory_argmem_none(ptr %x)
   call void @noarg()
   ret void
 }
 
 define void @test_use_argmem_read() {
 ; CHECK-LABEL: define void @test_use_argmem_read() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    call void @use_memory_argmem_read(ptr [[X]])
-; CHECK-NEXT:    call void @noarg()
+; CHECK-NEXT:    tail call void @noarg()
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %x = alloca i32, align 4
   call void @use_memory_argmem_read(ptr %x)
   call void @noarg()
   ret void
 }