[CoroSplit] AllocaUseVisitor visits insertvalue/insertelement

[NewSigma]

Pointers to allocas might be escaped by users of insertvalue/insertelement. AllocaUseVisitor should visit these instructions so that CoroSplit can successfully determine which allocas should live on the frame.

[llvmbot]

@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-llvm-transforms

Author: Weibo He (NewSigma)

Changes

Pointers to allocas might be escaped by users of insertvalue/insertelement. AllocaUseVisitor should visit these pointers so that CoroSplit can successfully determine which allocas should live on the frame.

---

Full diff: https://github.com/llvm/llvm-project/pull/156788.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/SpillUtils.cpp (+10)
• (added) llvm/test/Transforms/Coroutines/coro-alloca-09.ll (+62)

diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
index d5d60a320b521..098dc9917ceaf 100644
--- a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
+++ b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
@@ -183,6 +183,16 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
     handleAlias(I);
   }
 
+  void visitInsertElementInst(InsertElementInst &I) {
+    enqueueUsers(I);
+    handleAlias(I);
+  }
+
+  void visitInsertValueInst(InsertValueInst &I)  {
+    enqueueUsers(I);
+    handleAlias(I);
+  }
+
   void visitStoreInst(StoreInst &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.
diff --git a/llvm/test/Transforms/Coroutines/coro-alloca-09.ll b/llvm/test/Transforms/Coroutines/coro-alloca-09.ll
new file mode 100644
index 0000000000000..2539811f46b7c
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-alloca-09.ll
@@ -0,0 +1,62 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Tests that CoroSplit can succesfully determine allocas should live on the frame
+; if their aliases are used across suspension points through insertvalue/insertelement.
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+define ptr @f(i1 %n) presplitcoroutine {
+; CHECK-LABEL: define ptr @f(
+; CHECK-SAME: i1 [[N:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @f.resumers)
+; CHECK-NEXT:    [[MEM:%.*]] = call ptr @malloc(i32 56)
+; CHECK-NEXT:    [[HDL:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[MEM]])
+; CHECK-NEXT:    store ptr @f.resume, ptr [[HDL]], align 8
+; CHECK-NEXT:    [[DESTROY_ADDR:%.*]] = getelementptr inbounds nuw [[F_FRAME:%.*]], ptr [[HDL]], i32 0, i32 1
+; CHECK-NEXT:    store ptr @f.destroy, ptr [[DESTROY_ADDR]], align 8
+; CHECK-NEXT:    [[X_RELOAD_ADDR:%.*]] = getelementptr inbounds [[F_FRAME]], ptr [[HDL]], i32 0, i32 2
+; CHECK-NEXT:    [[ALIAS_SPILL_ADDR:%.*]] = getelementptr inbounds [[F_FRAME]], ptr [[HDL]], i32 0, i32 3
+; CHECK-NEXT:    store i64 0, ptr [[X_RELOAD_ADDR]], align 4
+; CHECK-NEXT:    store i64 0, ptr [[ALIAS_SPILL_ADDR]], align 4
+; CHECK-NEXT:    [[X_ALIAS:%.*]] = insertvalue [1 x ptr] poison, ptr [[X_RELOAD_ADDR]], 0
+; CHECK-NEXT:    [[X_ALIAS_SPILL_ADDR:%.*]] = getelementptr inbounds [[F_FRAME]], ptr [[HDL]], i32 0, i32 4
+; CHECK-NEXT:    store [1 x ptr] [[X_ALIAS]], ptr [[X_ALIAS_SPILL_ADDR]], align 8
+; CHECK-NEXT:    [[Y_ALIAS:%.*]] = insertelement <1 x ptr> poison, ptr [[ALIAS_SPILL_ADDR]], i32 0
+; CHECK-NEXT:    [[Y_ALIAS_SPILL_ADDR:%.*]] = getelementptr inbounds [[F_FRAME]], ptr [[HDL]], i32 0, i32 5
+; CHECK-NEXT:    store <1 x ptr> [[Y_ALIAS]], ptr [[Y_ALIAS_SPILL_ADDR]], align 8
+; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds nuw [[F_FRAME]], ptr [[HDL]], i32 0, i32 6
+; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR1]], align 1
+; CHECK-NEXT:    ret ptr [[HDL]]
+;
+entry:
+  %x = alloca i64
+  %y = alloca i64
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %mem = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %mem)
+  store i64 0, ptr %x
+  store i64 0, ptr %y
+  %x.alias = insertvalue [1 x ptr] poison, ptr %x, 0
+  %y.alias = insertelement <1 x ptr> poison, ptr %y, i32 0
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume
+  i8 1, label %cleanup]
+resume:
+  call void @use1([1 x ptr] %x.alias)
+  call void @use2(<1 x ptr> %y.alias)
+  br label %cleanup
+
+cleanup:
+  %mem1 = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem1)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret ptr %hdl
+}
+
+declare void @use1([1 x ptr])
+declare void @use2(<1 x ptr>)
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[NewSigma]

I think this should fix #149604, but we should not close it before reverting the previous workaround.

[zmodem]

I'm not familiar with AllocaUseVisitor and handleAlias(), but how do we know that we're not missing other visit functions?

It's not clear to me what criteria is used to decide which instructions need a visit function and call to handleAlias(), and which do not.

For example, I noticed that visitPtrToIntInst() isn't implemented. Couldn't an alloca escape through that too?

[ChuanqiXu9]

I'm not familiar with AllocaUseVisitor and handleAlias(), but how do we know that we're not missing other visit functions?

I think we don't know. The idea is to avoid putting things on the coroutine frame as much as possible. This was originally for optimization. But later we find it is needed since we won't have the alloca before coro.begin.

It's not clear to me what criteria is used to decide which instructions need a visit function and call to handleAlias(), and which do not.

It is basically an escape analysis.

For example, I noticed that visitPtrToIntInst() isn't implemented. Couldn't an alloca escape through that too?

If it is allowed to create a pointer from an int (I believe so), yes, it may be a problem.

[NewSigma]

I'm not familiar with AllocaUseVisitor and handleAlias(), but how do we know that we're not missing other visit functions?

All visit functions will call visitInstruction of InstVisitor by default. I scaned the code and I think we have handled all paths to visitInstruction. We can add an assertion to visitInstruction if further validation is desired.

It's not clear to me what criteria is used to decide which instructions need a visit function and call to handleAlias(), and which do not.

It is basically an escape analysis.

For example, I noticed that visitPtrToIntInst() isn't implemented.

No, it is handled by PtrUseVisitor, the base class of AllocaUseVisitor.

[zmodem]

But wasn't insertvalue/insertelement also already handled by PtrUseVisitor (via its base, InstVisitor). Why do they need special handling (handleAlias) and for example ptrtoint doesn't?

[NewSigma]

But wasn't insertvalue/insertelement also already handled by PtrUseVisitor (via its base, InstVisitor).

Hmm, I think InstVisitor is basically a no-op. We should override the definition in a derived class.

Why do they need special handling (handleAlias) and for example ptrtoint doesn't?

If we were to use ptrtoint before coro.begin, I think it would be reasonable to apply handleAlias to it.