[llvm] [LoopRotate][coroutines] Avoid hoisting addresses of thread-local variables outside loops in coroutines (PR #81937)

[Alan Zhao via llvm-commits]

https://github.com/alanzhao1 updated https://github.com/llvm/llvm-project/pull/81937

>From 045dc1f1909d3ce560dd4b39b80bb7c0f5569088 Mon Sep 17 00:00:00 2001
From: Alan Zhao <ayzhao at google.com>
Date: Thu, 15 Feb 2024 14:43:42 -0800
Subject: [PATCH 1/2] [LoopRotate][coroutines] Avoid caching addresses of
 thread-local variables outside loops in coroutines

Because loops in coroutines may have a co_await statement that
reschedules the coroutine to another thread, we cannot cache addresses
of thread-local variables obtained inside a loop by moving the
computation of thoes addresses outside a loop.

Since LLVM doesn't have a model for coroutine memory accesses, this
patch fixes this bug by disabling this optimization for coroutines in
the same way as https://reviews.llvm.org/D135550 and
https://reviews.llvm.org/D151774.
---
 .../Transforms/Utils/LoopRotationUtils.cpp    | 10 +++++-
 llvm/test/Transforms/LoopRotate/coroutine.ll  | 34 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopRotate/coroutine.ll

diff --git a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
index ec59a077302037..fb75f1622513df 100644
--- a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
@@ -612,7 +612,15 @@ bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
       // memory (without proving that the loop doesn't write).
       if (L->hasLoopInvariantOperands(Inst) && !Inst->mayReadFromMemory() &&
           !Inst->mayWriteToMemory() && !Inst->isTerminator() &&
-          !isa<DbgInfoIntrinsic>(Inst) && !isa<AllocaInst>(Inst)) {
+          !isa<DbgInfoIntrinsic>(Inst) && !isa<AllocaInst>(Inst) &&
+          // FIXME: It is not safe to cache the value of these instructions in
+          // coroutines, as the addresses of otherwise eligible variables (e.g.
+          // thread-local variables and errno) may change if the coroutine is
+          // resumed in a different thread. Therefore, we disable this
+          // optimization for correctness. However, this may block other correct
+          // optimizations. This should be reverted once we have a better model
+          // for memory access in coroutines.
+          !Inst->getFunction()->isPresplitCoroutine()) {
 
         if (LoopEntryBranch->getParent()->IsNewDbgInfoFormat &&
             !NextDbgInsts.empty()) {
diff --git a/llvm/test/Transforms/LoopRotate/coroutine.ll b/llvm/test/Transforms/LoopRotate/coroutine.ll
new file mode 100644
index 00000000000000..6dbaab2ecd7097
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/coroutine.ll
@@ -0,0 +1,34 @@
+; RUN: opt -S -passes=loop-rotate < %s | FileCheck %s
+
+declare void @bar1()
+
+ at threadlocalint = thread_local global i32 0, align 4
+
+define void @foo() #0 {
+; CHECK-LABEL: entry:
+; CHECK: call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @threadlocalint)
+; CHECK: br {{.*}} label %cond.end
+entry:
+  br label %while.cond
+
+while.cond:
+  %1 = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @threadlocalint)
+  %2 = load i32, ptr %1, align 4
+  %cmp = icmp eq i32 %2, 0
+  br i1 %cmp, label %cond.end, label %cond.false
+
+cond.false:
+  call void @bar1()
+  unreachable
+
+; The address of threadlocalint must not be cached outside loops in presplit
+; coroutines.
+; CHECK-LABEL: cond.end:
+; CHECK: call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @threadlocalint)
+; CHECK: br {{.*}} label %cond.end
+cond.end:
+  call void @bar1()
+  br label %while.cond
+}
+
+attributes #0 = { presplitcoroutine }

>From bffd43dd44b34424e3286f68c58c1a4e6af153f6 Mon Sep 17 00:00:00 2001
From: Alan Zhao <ayzhao at google.com>
Date: Tue, 20 Feb 2024 10:47:43 -0800
Subject: [PATCH 2/2] update coroutine.ll test

---
 llvm/test/Transforms/LoopRotate/coroutine.ll | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/test/Transforms/LoopRotate/coroutine.ll b/llvm/test/Transforms/LoopRotate/coroutine.ll
index 9d7bf595e21cdc..4544ac5e72bbe5 100644
--- a/llvm/test/Transforms/LoopRotate/coroutine.ll
+++ b/llvm/test/Transforms/LoopRotate/coroutine.ll
@@ -21,8 +21,9 @@ define void @foo() #0 {
 ; CHECK-NEXT:    ret void
 ; CHECK:       cond.end:
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr [[TMP0]], align 4
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[TMP2]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @threadlocalint)
+; CHECK-NEXT:    [[TMP3:%.*]] = load i32, ptr [[TMP2]], align 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[TMP3]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[COND_END]], label [[WHILE_COND_COND_FALSE_CRIT_EDGE:%.*]]
 ;
 entry: