[LoopInfo] Pointer to stack object may not be loop invariant in a coroutine function (#149936)

weiguozhi · web-flow · commit 5e8779220069 · 2025-08-09T14:20:19.000-07:00
A coroutine function may be split to ramp function and resume function, and they have different stack frames, so a pointer to stack objects may have different addresses depending on where it is used, so it's not a loop invariant. It temporarily fixes #149604.
diff --git a/llvm/include/llvm/Analysis/LoopInfo.h b/llvm/include/llvm/Analysis/LoopInfo.h
@@ -59,11 +59,12 @@ class LLVM_ABI Loop : public LoopBase<BasicBlock, Loop> {
   };
 
   /// Return true if the specified value is loop invariant.
-  bool isLoopInvariant(const Value *V) const;
+  bool isLoopInvariant(const Value *V, bool HasCoroSuspendInst = false) const;
 
   /// Return true if all the operands of the specified instruction are loop
   /// invariant.
-  bool hasLoopInvariantOperands(const Instruction *I) const;
+  bool hasLoopInvariantOperands(const Instruction *I,
+                                bool HasCoroSuspendInst = false) const;
 
   /// If the given value is an instruction inside of the loop and it can be
   /// hoisted, do so to make it trivially loop-invariant.
diff --git a/llvm/include/llvm/Transforms/Utils/LoopUtils.h b/llvm/include/llvm/Transforms/Utils/LoopUtils.h
@@ -185,7 +185,8 @@ LLVM_ABI bool hoistRegion(DomTreeNode *, AAResults *, LoopInfo *,
                           TargetLibraryInfo *, Loop *, MemorySSAUpdater &,
                           ScalarEvolution *, ICFLoopSafetyInfo *,
                           SinkAndHoistLICMFlags &, OptimizationRemarkEmitter *,
-                          bool, bool AllowSpeculation);
+                          bool, bool AllowSpeculation,
+                          bool HasCoroSuspendInst = false);
 
 /// Return true if the induction variable \p IV in a Loop whose latch is
 /// \p LatchBlock would become dead if the exit test \p Cond were removed.
diff --git a/llvm/lib/Analysis/LoopInfo.cpp b/llvm/lib/Analysis/LoopInfo.cpp
@@ -58,14 +58,26 @@ static cl::opt<bool, true>
 // Loop implementation
 //
 
-bool Loop::isLoopInvariant(const Value *V) const {
-  if (const Instruction *I = dyn_cast<Instruction>(V))
-    return !contains(I);
+bool Loop::isLoopInvariant(const Value *V, bool HasCoroSuspendInst) const {
+  if (const Instruction *I = dyn_cast<Instruction>(V)) {
+    // FIXME: this is semantically inconsistent. We're tracking a proper fix in
+    // issue #149604.
+    // If V is a pointer to stack object and L contains a coro.suspend function
+    // call, then V may not be loop invariant because the ramp function and
+    // resume function have different stack frames.
+    if (HasCoroSuspendInst && isa<AllocaInst>(I))
+      return false;
+    else
+      return !contains(I);
+  }
   return true; // All non-instructions are loop invariant
 }
 
-bool Loop::hasLoopInvariantOperands(const Instruction *I) const {
-  return all_of(I->operands(), [this](Value *V) { return isLoopInvariant(V); });
+bool Loop::hasLoopInvariantOperands(const Instruction *I,
+                                    bool HasCoroSuspendInst) const {
+  return all_of(I->operands(), [&](Value *V) {
+    return isLoopInvariant(V, HasCoroSuspendInst);
+  });
 }
 
 bool Loop::makeLoopInvariant(Value *V, bool &Changed, Instruction *InsertPt,
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -472,7 +472,7 @@ bool LoopInvariantCodeMotion::runOnLoop(Loop *L, AAResults *AA, LoopInfo *LI,
   if (Preheader)
     Changed |= hoistRegion(DT->getNode(L->getHeader()), AA, LI, DT, AC, TLI, L,
                            MSSAU, SE, &SafetyInfo, Flags, ORE, LoopNestMode,
-                           LicmAllowSpeculation);
+                           LicmAllowSpeculation, HasCoroSuspendInst);
 
   // Now that all loop invariants have been removed from the loop, promote any
   // memory references to scalars that we can.
@@ -881,7 +881,7 @@ bool llvm::hoistRegion(DomTreeNode *N, AAResults *AA, LoopInfo *LI,
                        ICFLoopSafetyInfo *SafetyInfo,
                        SinkAndHoistLICMFlags &Flags,
                        OptimizationRemarkEmitter *ORE, bool LoopNestMode,
-                       bool AllowSpeculation) {
+                       bool AllowSpeculation, bool HasCoroSuspendInst) {
   // Verify inputs.
   assert(N != nullptr && AA != nullptr && LI != nullptr && DT != nullptr &&
          CurLoop != nullptr && SafetyInfo != nullptr &&
@@ -914,11 +914,11 @@ bool llvm::hoistRegion(DomTreeNode *N, AAResults *AA, LoopInfo *LI,
       // TODO: It may be safe to hoist if we are hoisting to a conditional block
       // and we have accurately duplicated the control flow from the loop header
       // to that block.
-      if (CurLoop->hasLoopInvariantOperands(&I) &&
+      if (CurLoop->hasLoopInvariantOperands(&I, HasCoroSuspendInst) &&
           canSinkOrHoistInst(I, AA, DT, CurLoop, MSSAU, true, Flags, ORE) &&
-          isSafeToExecuteUnconditionally(
-              I, DT, TLI, CurLoop, SafetyInfo, ORE,
-              Preheader->getTerminator(), AC, AllowSpeculation)) {
+          isSafeToExecuteUnconditionally(I, DT, TLI, CurLoop, SafetyInfo, ORE,
+                                         Preheader->getTerminator(), AC,
+                                         AllowSpeculation)) {
         hoist(I, DT, CurLoop, CFH.getOrCreateHoistedBlock(BB), SafetyInfo,
               MSSAU, SE, ORE);
         HoistedInstructions.push_back(&I);
@@ -964,7 +964,7 @@ bool llvm::hoistRegion(DomTreeNode *N, AAResults *AA, LoopInfo *LI,
                SafetyInfo->doesNotWriteMemoryBefore(I, CurLoop);
       };
       if ((IsInvariantStart(I) || isGuard(&I)) &&
-          CurLoop->hasLoopInvariantOperands(&I) &&
+          CurLoop->hasLoopInvariantOperands(&I, HasCoroSuspendInst) &&
           MustExecuteWithoutWritesBefore(I)) {
         hoist(I, DT, CurLoop, CFH.getOrCreateHoistedBlock(BB), SafetyInfo,
               MSSAU, SE, ORE);
diff --git a/llvm/test/Transforms/LICM/licm-coroutine.ll b/llvm/test/Transforms/LICM/licm-coroutine.ll
@@ -0,0 +1,78 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=licm -S | FileCheck %s
+
+; %fca.0 and %fca.1 should not be hoisted out of the loop because the ramp
+; function and resume function have different stack frames, so %pointer1 and
+; %pointer2 have different values before and after @llvm.coro.suspend.
+
+define ptr @f(i32 %n) presplitcoroutine {
+; CHECK-LABEL: define ptr @f(
+; CHECK-SAME: i32 [[N:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[POINTER1:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    [[POINTER2:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+; CHECK-NEXT:    [[SIZE:%.*]] = call i32 @llvm.coro.size.i32()
+; CHECK-NEXT:    [[ALLOC:%.*]] = call ptr @malloc(i32 [[SIZE]])
+; CHECK-NEXT:    [[HDL:%.*]] = call noalias ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[N_VAL:%.*]] = phi i32 [ [[N]], %[[ENTRY]] ], [ [[INC:%.*]], %[[RESUME:.*]] ]
+; CHECK-NEXT:    [[INC]] = add nsw i32 [[N_VAL]], 1
+; CHECK-NEXT:    call void @print(i32 [[N_VAL]])
+; CHECK-NEXT:    [[TMP0:%.*]] = call i8 @llvm.coro.suspend(token none, i1 false)
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[SUSPEND_LOOPEXIT:.*]] [
+; CHECK-NEXT:      i8 0, label %[[RESUME]]
+; CHECK-NEXT:      i8 1, label %[[CLEANUP:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[RESUME]]:
+; CHECK-NEXT:    [[FCA_0:%.*]] = insertvalue [2 x ptr] poison, ptr [[POINTER1]], 0
+; CHECK-NEXT:    [[FCA_1:%.*]] = insertvalue [2 x ptr] [[FCA_0]], ptr [[POINTER2]], 1
+; CHECK-NEXT:    call void @foo([2 x ptr] [[FCA_1]])
+; CHECK-NEXT:    br label %[[LOOP]]
+; CHECK:       [[CLEANUP]]:
+; CHECK-NEXT:    [[MEM:%.*]] = call ptr @llvm.coro.free(token [[ID]], ptr [[HDL]])
+; CHECK-NEXT:    call void @free(ptr [[MEM]])
+; CHECK-NEXT:    br label %[[SUSPEND:.*]]
+; CHECK:       [[SUSPEND_LOOPEXIT]]:
+; CHECK-NEXT:    br label %[[SUSPEND]]
+; CHECK:       [[SUSPEND]]:
+; CHECK-NEXT:    [[UNUSED:%.*]] = call i1 @llvm.coro.end(ptr [[HDL]], i1 false, token none)
+; CHECK-NEXT:    ret ptr [[HDL]]
+;
+entry:
+  %pointer1 = alloca ptr
+  %pointer2 = alloca ptr
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call noalias ptr @llvm.coro.begin(token %id, ptr %alloc)
+  br label %loop
+
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  %inc = add nsw i32 %n.val, 1
+  call void @print(i32 %n.val)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+  i8 1, label %cleanup]
+
+resume:
+  %fca.0 = insertvalue [2 x ptr] poison, ptr %pointer1, 0
+  %fca.1 = insertvalue [2 x ptr] %fca.0, ptr %pointer2, 1
+  call void @foo([2 x ptr] %fca.1)
+  br label %loop
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+suspend:
+  %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none)
+  ret ptr %hdl
+}
+
+declare void @free(ptr)
+declare ptr @malloc(i32)
+declare void @print(i32)
+declare void @foo([2 x ptr])
