[Coroutines] Move OptimizeFrame out of Shape (#111017)

* OptimizeFrame is not really a part of the Coroutine Shape info, rather
it is specifically for the addFieldForAllocas method called indirectly
by buildCoroutineFrame.
* This patch passes OptimizeFrame directly to buildCoroutineFrame so it
can be provided to addFieldForAllocas instead of keeping it in the
Shape.

Co-authored-by: tnowicki <[email protected]>

llvm/lib/Transforms/Coroutines/ABI.h

@@ -41,7 +41,7 @@ class LLVM_LIBRARY_VISIBILITY BaseABI {
   virtual void init() = 0;
 
   // Allocate the coroutine frame and do spill/reload as needed.
-  virtual void buildCoroutineFrame();
+  virtual void buildCoroutineFrame(bool OptimizeFrame);
 
   // Perform the function splitting according to the ABI.
   virtual void splitCoroutine(Function &F, coro::Shape &Shape,

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -234,7 +234,7 @@ class FrameTypeBuilder {
   /// Side Effects: Because We sort the allocas, the order of allocas in the
   /// frame may be different with the order in the source code.
   void addFieldForAllocas(const Function &F, FrameDataInfo &FrameData,
-                          coro::Shape &Shape);
+                          coro::Shape &Shape, bool OptimizeFrame);
 
   /// Add a field to this structure.
   [[nodiscard]] FieldIDType addField(Type *Ty, MaybeAlign MaybeFieldAlignment,
@@ -336,7 +336,8 @@ void FrameDataInfo::updateLayoutIndex(FrameTypeBuilder &B) {
 
 void FrameTypeBuilder::addFieldForAllocas(const Function &F,
                                           FrameDataInfo &FrameData,
-                                          coro::Shape &Shape) {
+                                          coro::Shape &Shape,
+                                          bool OptimizeFrame) {
   using AllocaSetType = SmallVector<AllocaInst *, 4>;
   SmallVector<AllocaSetType, 4> NonOverlapedAllocas;
 
@@ -350,7 +351,7 @@ void FrameTypeBuilder::addFieldForAllocas(const Function &F,
     }
   });
 
-  if (!Shape.OptimizeFrame) {
+  if (!OptimizeFrame) {
     for (const auto &A : FrameData.Allocas) {
       AllocaInst *Alloca = A.Alloca;
       NonOverlapedAllocas.emplace_back(AllocaSetType(1, Alloca));
@@ -860,7 +861,8 @@ static void buildFrameDebugInfo(Function &F, coro::Shape &Shape,
 //     ... spills ...
 //   };
 static StructType *buildFrameType(Function &F, coro::Shape &Shape,
-                                  FrameDataInfo &FrameData) {
+                                  FrameDataInfo &FrameData,
+                                  bool OptimizeFrame) {
   LLVMContext &C = F.getContext();
   const DataLayout &DL = F.getDataLayout();
   StructType *FrameTy = [&] {
@@ -905,7 +907,7 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
 
   // Because multiple allocas may own the same field slot,
   // we add allocas to field here.
-  B.addFieldForAllocas(F, FrameData, Shape);
+  B.addFieldForAllocas(F, FrameData, Shape, OptimizeFrame);
   // Add PromiseAlloca to Allocas list so that
   // 1. updateLayoutIndex could update its index after
   // `performOptimizedStructLayout`
@@ -2056,7 +2058,7 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
   rewritePHIs(F);
 }
 
-void coro::BaseABI::buildCoroutineFrame() {
+void coro::BaseABI::buildCoroutineFrame(bool OptimizeFrame) {
   SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds);
   doRematerializations(F, Checker, IsMaterializable);
 
@@ -2087,7 +2089,7 @@ void coro::BaseABI::buildCoroutineFrame() {
 
   // Build frame
   FrameDataInfo FrameData(Spills, Allocas);
-  Shape.FrameTy = buildFrameType(F, Shape, FrameData);
+  Shape.FrameTy = buildFrameType(F, Shape, FrameData, OptimizeFrame);
   Shape.FramePtr = Shape.CoroBegin;
   // For now, this works for C++ programs only.
   buildFrameDebugInfo(F, Shape, FrameData);

llvm/lib/Transforms/Coroutines/CoroShape.h

@@ -112,9 +112,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   Value *FramePtr = nullptr;
   BasicBlock *AllocaSpillBlock = nullptr;
 
-  /// This would only be true if optimization are enabled.
-  bool OptimizeFrame;
-
   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
@@ -265,8 +262,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
 
   Shape() = default;
-  explicit Shape(Function &F, bool OptimizeFrame = false)
-      : OptimizeFrame(OptimizeFrame) {
+  explicit Shape(Function &F) {
     SmallVector<CoroFrameInst *, 8> CoroFrames;
     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
 

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -2053,7 +2053,8 @@ void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
 }
 
 static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
-                             coro::BaseABI &ABI, TargetTransformInfo &TTI) {
+                             coro::BaseABI &ABI, TargetTransformInfo &TTI,
+                             bool OptimizeFrame) {
   PrettyStackTraceFunction prettyStackTrace(F);
 
   auto &Shape = ABI.Shape;
@@ -2064,7 +2065,7 @@ static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   simplifySuspendPoints(Shape);
 
   normalizeCoroutine(F, Shape, TTI);
-  ABI.buildCoroutineFrame();
+  ABI.buildCoroutineFrame(OptimizeFrame);
   replaceFrameSizeAndAlignment(Shape);
 
   bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
@@ -2273,7 +2274,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     // unreachable blocks before collecting intrinsics into Shape.
     removeUnreachableBlocks(F);
 
-    coro::Shape Shape(F, OptimizeFrame);
+    coro::Shape Shape(F);
     if (!Shape.CoroBegin)
       continue;
 
@@ -2283,7 +2284,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
 
     SmallVector<Function *, 4> Clones;
     auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
-    doSplitCoroutine(F, Clones, *ABI, TTI);
+    doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame);
     CurrentSCC = &updateCallGraphAfterCoroutineSplit(
         *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);