[Coroutine] Allocas used by StoreInst does not always escape

[lxfind]

Summary

In the existing logic, for a given alloca, as long as its pointer value is stored into another location, it's considered as escaped.
This is a bit too conservative. Specifically, in non-optimized build mode, it's often to have patterns of code that first store an alloca somewhere and then load it right away.
These used should be handled without conservatively marking them escaped.

This patch tracks how the memory location where an alloca pointer is stored into is being used. As long as we only try to load from that location and nothing else, we can still
consider the original alloca not escaping and keep it on the stack instead of putting it on the frame.

[ChuanqiXu]

Good. The intuition and implementation make sense.

[lxfind]

Landed by accident. Reverted and reopening for review.

[junparser]

I'm not sure whether we should fix this case by case. Maybe we should use memoryssa to handle this.

[lxfind]

In D91305#2392934, @junparser wrote:

I'm not sure whether we should fix this case by case. Maybe we should use memoryssa to handle this.

In general yes I agree with you that we don't want to handle case by case and make this over-complicated. However there are good reasons that this patch is quite necessary.

1. When optimizations are enabled, there are usually pretty accurate lifetime information, and hence we don't usually even need the alloca visitor to track them. So the alloca visitor is crucial when optimizations are turned off and there is no lifetime information. When optimizations are turned off, because there isn't mem2reg, almost all data accesses are done through the patterns of stores followed with loads. What I found is that in such case, when optimizations are turned off, almost all allocas are put on the frame because they are all used in store instructions, making this a bit useless. This patch is a key to make the alloca tracking effective and so that majority of allocas can be moved out of frame. Hence the effect of this patch is quite significant comparing to the effort. It's unlikely that I will add any more logic to track other cases.
2. This is also a continuation of D90990, which is to avoid frame-use-after-destroy. What I found is that, even I fixed the front-end such that the use of temporaries after a call to await_suspend() are contained locally that does not go across suspension points, because they are all involved in this store/load pattern, they are still being put on the frame, causing memory corruptions in non-optimized mode. This patch can fix that issue completely.

[junparser]

In D91305#2392970, @lxfind wrote:

In D91305#2392934, @junparser wrote:

I'm not sure whether we should fix this case by case. Maybe we should use memoryssa to handle this.

In general yes I agree with you that we don't want to handle case by case and make this over-complicated. However there are good reasons that this patch is quite necessary.

1. When optimizations are enabled, there are usually pretty accurate lifetime information, and hence we don't usually even need the alloca visitor to track them. So the alloca visitor is crucial when optimizations are turned off and there is no lifetime information. When optimizations are turned off, because there isn't mem2reg, almost all data accesses are done through the patterns of stores followed with loads. What I found is that in such case, when optimizations are turned off, almost all allocas are put on the frame because they are all used in store instructions, making this a bit useless. This patch is a key to make the alloca tracking effective and so that majority of allocas can be moved out of frame. Hence the effect of this patch is quite significant comparing to the effort. It's unlikely that I will add any more logic to track other cases.
2. This is also a continuation of D90990, which is to avoid frame-use-after-destroy. What I found is that, even I fixed the front-end such that the use of temporaries after a call to await_suspend() are contained locally that does not go across suspension points, because they are all involved in this store/load pattern, they are still being put on the frame, causing memory corruptions in non-optimized mode. This patch can fix that issue completely.

Do we really need handle this in non-optimized mode? @bruno I'm afraid that this patch may affect the debugging quality fixed in D90772, @lxfind, have you ever checked this?

[lxfind]

Do we really need handle this in non-optimized mode? @bruno I'm afraid that this patch may affect the debugging quality fixed in D90772, @lxfind, have you ever checked this?

I believe we do. As I mentioned in the second reason above, there should be absolutely no frame access after the call to await_suspend() until the symmetric transfer, since the frame may have been destroyed in await_suspend(). However based on the current way of determining whether an alloca should live in the frame, almost all allocas will be put on the frame in O0 because of the common use of store instructions. Such issue leads to TSAN failures since TSAN usually runs in O0 without lifetime intrinsics.

This patch doesn't affects debugability. In general, whether or not an alloca should be put on the frame should never affect debugability. If so that would be a problem that needs to be fixed alone.

[junparser]

Thanks for the patch！

[lxfind]

In D91305#2395900, @junparser wrote:

Thanks for the patch！

Thank you for the review!

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
   }
 
   void visitSelectInst(SelectInst &I) {
     enqueueUsers(I);
     handleAlias(I);
   }
 
   void visitStoreInst(StoreInst &SI) {
-    // Base visit function will handle escape setting.
-    Base::visitStoreInst(SI);
-
     // Regardless whether the alias of the alloca is the value operand or the
     // pointer operand, we need to assume the alloca is been written.
     handleMayWrite(SI);
+
+    if (SI.getValueOperand() != U->get())
+      return;
+
+    // We are storing the pointer into a memory location, potentially escaping.
+    // As an optimization, we try to detect simple cases where it doesn't
+    // actually escape, for example:
+    //   %ptr = alloca ..
+    //   %addr = alloca ..
+    //   store %ptr, %addr
+    //   %x = load %addr
+    //   ..
+    // If %addr is only used by loading from it, we could simply treat %x as
+    // another alias of %ptr, and not considering %ptr being escaped.
+    auto IsSimpleStoreThenLoad = [&]() {
+      auto *AI = dyn_cast<AllocaInst>(SI.getPointerOperand());
+      // If the memory location we are storing to is not an alloca, it
+      // could be an alias of some other memory locations, which is difficult
+      // to analyze.
+      if (!AI)
+        return false;
+      // StoreAliases contains aliases of the memory location stored into.
+      SmallVector<Instruction *, 4> StoreAliases = {AI};
+      while (!StoreAliases.empty()) {
+        Instruction *I = StoreAliases.back();
+        StoreAliases.pop_back();
+        for (User *U : I->users()) {
+          // If we are loading from the memory location, we are creating an
+          // alias of the original pointer.
+          if (auto *LI = dyn_cast<LoadInst>(U)) {
+            enqueueUsers(*LI);
+            handleAlias(*LI);
+            continue;
+          }
+          // If we are overriding the memory location, the pointer certainly
+          // won't escape.
+          if (auto *S = dyn_cast<StoreInst>(U))
+            if (S->getPointerOperand() == I)
+              continue;
+          if (auto *II = dyn_cast<IntrinsicInst>(U))
+            if (II->isLifetimeStartOrEnd())
+              continue;
+          // BitCastInst creats aliases of the memory location being stored
+          // into.
+          if (auto *BI = dyn_cast<BitCastInst>(U)) {
+            StoreAliases.push_back(BI);
+            continue;
+          }
+          return false;
+        }
+      }
+
+      return true;
+    };
+
+    if (!IsSimpleStoreThenLoad())
+      PI.setEscaped(&SI);
   }
 
   // All mem intrinsics modify the data.
   void visitMemIntrinsic(MemIntrinsic &MI) { handleMayWrite(MI); }
 
   void visitBitCastInst(BitCastInst &BC) {
     Base::visitBitCastInst(BC);
     handleAlias(BC);

llvm/test/Transforms/Coroutines/coro-alloca-06.ll

+; Test that in some simple cases allocas will not live on the frame even
+; though their pointers are stored.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+%handle = type { i8* }
+
+define i8* @f() "coroutine.presplit"="1" {
+entry:
+  %0 = alloca %"handle", align 8
+  %1 = alloca %"handle"*, align 8
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  br label %tricky
+
+tricky:
+  %2 = call i8* @await_suspend()
+  %3 = getelementptr inbounds %"handle", %"handle"* %0, i32 0, i32 0
+  store i8* %2, i8** %3, align 8
+  %4 = bitcast %"handle"** %1 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 8, i8* %4)
+  store %"handle"* %0, %"handle"** %1, align 8
+  %5 = load %"handle"*, %"handle"** %1, align 8
+  %6 = getelementptr inbounds %"handle", %"handle"* %5, i32 0, i32 0
+  %7 = load i8*, i8** %6, align 8
+  %8 = bitcast %"handle"** %1 to i8*
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* %8)
+  br label %finish
+
+finish:
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume
+  i8 1, label %cleanup]
+resume:
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; CHECK:        %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i1 }
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca [[HANDLE:%.*]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca %handle*, align 8
+
+; CHECK:         [[TMP2:%.*]] = call i8* @await_suspend()
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds [[HANDLE]], %handle* [[TMP0]], i32 0, i32 0
+; CHECK-NEXT:    store i8* [[TMP2]], i8** [[TMP3]], align 8
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast %handle** [[TMP1]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 8, i8* [[TMP4]])
+; CHECK-NEXT:    store %handle* [[TMP0]], %handle** [[TMP1]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 8, i8* [[TMP4]])
+;
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
+
+declare i8* @await_suspend()
+declare void @print(i32* nocapture)
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)