[llvm] [CoroSplit] AllocaUseVisitor visits insertvalue/insertelement (PR #156788)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Weibo He (NewSigma)

<details>
<summary>Changes</summary>

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
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)

``````````

</details>


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