[llvm/llvm-project][Coroutines] ABI Object

* Adds the ABI object hierarchy for each type of ABI
* Replace some uses of the Shape.ABI enum with calls to the ABI object
  methods.

llvm/include/llvm/Transforms/Coroutines/CoroSplit.h

@@ -21,19 +21,26 @@
 
 namespace llvm {
 
+namespace coro {
+class BaseABI;
+class Shape;
+} // namespace coro
+
 struct CoroSplitPass : PassInfoMixin<CoroSplitPass> {
-  const std::function<bool(Instruction &)> MaterializableCallback;
+  // BaseABITy generates an instance of a coro ABI.
+  using BaseABITy = std::function<coro::BaseABI *(Function &, coro::Shape &)>;
 
   CoroSplitPass(bool OptimizeFrame = false);
   CoroSplitPass(std::function<bool(Instruction &)> MaterializableCallback,
-                bool OptimizeFrame = false)
-      : MaterializableCallback(MaterializableCallback),
-        OptimizeFrame(OptimizeFrame) {}
+                bool OptimizeFrame = false);
 
   PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR);
   static bool isRequired() { return true; }
 
+  // Generator for an ABI transformer
+  BaseABITy GenABI;
+
   // Would be true if the Optimization level isn't O0.
   bool OptimizeFrame;
 };

llvm/lib/Transforms/Coroutines/ABI.h

@@ -0,0 +1,115 @@
+//===- ABI.h - Coroutine ABI Transformers ---------------------*- C++ -*---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// This file declares the pass that analyzes a function for coroutine intrs and
+// a transformer class that contains methods for handling different steps of
+// coroutine lowering.
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_TRANSFORMS_COROUTINES_ABI_H
+#define LIB_TRANSFORMS_COROUTINES_ABI_H
+
+#include "CoroShape.h"
+#include "MaterializationUtils.h"
+#include "SuspendCrossingInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+
+namespace llvm {
+
+class Function;
+
+namespace coro {
+
+// This interface/API is to provide an object oriented way to implement ABI
+// functionality. This is intended to replace use of switch(ABI) to perform
+// ABI specific behavior. The ABIs (e.g. Switch, Async, Retcon{Once}) are the
+// common ABIs. However, specific users often need to modify the behavior of
+// these, such as for C++20 or Swift. This can be accomplished by inheriting
+// one of the common ABIs and overriding one or more of the methods to create
+// a custom ABI. The custom ABI is specified with the coro.begin.custom.abi
+// intrinsic instead of the coro.begin intrinsic by providing an i32 in the
+// last argument. This is used to lookup a generator for the custom ABI from
+// a set of generators provided to the CoroSplitPass constructor.
+//
+class LLVM_LIBRARY_VISIBILITY BaseABI {
+public:
+  BaseABI(Function &F, Shape &S)
+      : F(F), Shape(S), IsMaterializable(coro::isTriviallyMaterializable) {}
+
+  BaseABI(Function &F, coro::Shape &S,
+          std::function<bool(Instruction &)> IsMaterializable)
+      : F(F), Shape(S), IsMaterializable(IsMaterializable) {}
+
+  // Initialize the coroutine ABI
+  virtual void init() = 0;
+
+  // Allocate the coroutine frame and do spill/reload as needed.
+  virtual void buildCoroutineFrame(TargetTransformInfo &TTI);
+
+  // Perform the function splitting according to the ABI.
+  virtual void splitCoroutine(Function &F, coro::Shape &Shape,
+                              SmallVectorImpl<Function *> &Clones,
+                              TargetTransformInfo &TTI) = 0;
+
+  Function &F;
+  coro::Shape &Shape;
+
+  // Callback used by coro::BaseABI::buildCoroutineFrame for rematerialization.
+  // It is provided to coro::doMaterializations(..).
+  std::function<bool(Instruction &I)> IsMaterializable;
+};
+
+class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI {
+public:
+  SwitchABI(Function &F, coro::Shape &S) : BaseABI(F, S) {}
+
+  SwitchABI(Function &F, coro::Shape &S,
+            std::function<bool(Instruction &)> IsMaterializable)
+      : BaseABI(F, S, IsMaterializable) {}
+
+  void init() override;
+
+  void splitCoroutine(Function &F, coro::Shape &Shape,
+                      SmallVectorImpl<Function *> &Clones,
+                      TargetTransformInfo &TTI) override;
+};
+
+class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI {
+public:
+  AsyncABI(Function &F, coro::Shape &S) : BaseABI(F, S) {}
+
+  AsyncABI(Function &F, coro::Shape &S,
+           std::function<bool(Instruction &)> IsMaterializable)
+      : BaseABI(F, S, IsMaterializable) {}
+
+  void init() override;
+
+  void splitCoroutine(Function &F, coro::Shape &Shape,
+                      SmallVectorImpl<Function *> &Clones,
+                      TargetTransformInfo &TTI) override;
+};
+
+class LLVM_LIBRARY_VISIBILITY AnyRetconABI : public BaseABI {
+public:
+  AnyRetconABI(Function &F, coro::Shape &S) : BaseABI(F, S) {}
+
+  AnyRetconABI(Function &F, coro::Shape &S,
+               std::function<bool(Instruction &)> IsMaterializable)
+      : BaseABI(F, S, IsMaterializable) {}
+
+  void init() override;
+
+  void splitCoroutine(Function &F, coro::Shape &Shape,
+                      SmallVectorImpl<Function *> &Clones,
+                      TargetTransformInfo &TTI) override;
+};
+
+} // end namespace coro
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_COROUTINES_ABI_H

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -15,6 +15,7 @@
 // the value into the coroutine frame.
 //===----------------------------------------------------------------------===//
 
+#include "ABI.h"
 #include "CoroInternal.h"
 #include "MaterializationUtils.h"
 #include "SpillUtils.h"
@@ -2054,11 +2055,9 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
   rewritePHIs(F);
 }
 
-void coro::buildCoroutineFrame(
-    Function &F, Shape &Shape,
-    const std::function<bool(Instruction &)> &MaterializableCallback) {
+void coro::BaseABI::buildCoroutineFrame(TargetTransformInfo &TTI) {
   SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds);
-  doRematerializations(F, Checker, MaterializableCallback);
+  doRematerializations(F, Checker, IsMaterializable);
 
   const DominatorTree DT(F);
   if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&

llvm/lib/Transforms/Coroutines/CoroInternal.h

@@ -55,9 +55,6 @@ struct LowererBase {
 bool defaultMaterializable(Instruction &V);
 void normalizeCoroutine(Function &F, coro::Shape &Shape,
                         TargetTransformInfo &TTI);
-void buildCoroutineFrame(
-    Function &F, Shape &Shape,
-    const std::function<bool(Instruction &)> &MaterializableCallback);
 CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                              TargetTransformInfo &TTI,
                              ArrayRef<Value *> Arguments, IRBuilder<> &);

llvm/lib/Transforms/Coroutines/CoroShape.h

@@ -50,15 +50,53 @@ enum class ABI {
 // Holds structural Coroutine Intrinsics for a particular function and other
 // values used during CoroSplit pass.
 struct LLVM_LIBRARY_VISIBILITY Shape {
-  CoroBeginInst *CoroBegin;
+  CoroBeginInst *CoroBegin = nullptr;
   SmallVector<AnyCoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
   SmallVector<CoroAlignInst *, 2> CoroAligns;
   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
-  SmallVector<CallInst *, 2> SwiftErrorOps;
   SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
   SmallVector<CallInst *, 2> SymmetricTransfers;
 
+  // Values invalidated by bailout() and finalize()
+  SmallVector<CoroFrameInst *, 8> CoroFrames;
+  SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+
+  // Values invalidated by replaceSwiftErrorOps
+  SmallVector<CallInst *, 2> SwiftErrorOps;
+
+  Shape() = default;
+  explicit Shape(Function &F, bool OptimizeFrame)
+      : OptimizeFrame(OptimizeFrame) {
+    analyze(F);
+  }
+
+  void clear() {
+    CoroBegin = nullptr;
+    CoroEnds.clear();
+    CoroSizes.clear();
+    CoroAligns.clear();
+    CoroSuspends.clear();
+    CoroAwaitSuspends.clear();
+    SymmetricTransfers.clear();
+
+    CoroFrames.clear();
+    UnusedCoroSaves.clear();
+
+    SwiftErrorOps.clear();
+
+    FrameTy = nullptr;
+    FramePtr = nullptr;
+    AllocaSpillBlock = nullptr;
+  }
+
+  // Scan the function and collect the above intrinsics for later processing
+  void analyze(Function &F);
+  // If for some reason, we were not able to find coro.begin, bailout.
+  void bailout(Function &F);
+  // Remove orphaned and unnecessary intrinsics
+  void finalize();
+
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
     enum {
@@ -76,11 +114,13 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
 
   coro::ABI ABI;
 
-  StructType *FrameTy;
+  // Values below are used by ABIs for lowering. Ideally, these would live in
+  // the ABI objects along with the helpers that use them.
+  StructType *FrameTy = nullptr;
   Align FrameAlign;
-  uint64_t FrameSize;
-  Value *FramePtr;
-  BasicBlock *AllocaSpillBlock;
+  uint64_t FrameSize = 0;
+  Value *FramePtr = nullptr;
+  BasicBlock *AllocaSpillBlock = nullptr;
 
   /// This would only be true if optimization are enabled.
   bool OptimizeFrame;
@@ -94,6 +134,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     unsigned IndexOffset;
     bool HasFinalSuspend;
     bool HasUnwindCoroEnd;
+    size_t FinalSuspendIndex;
   };
 
   struct RetconLoweringStorage {
@@ -233,13 +274,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   ///
   /// \param CG - if non-null, will be updated for the new call
   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
-
-  Shape() = default;
-  explicit Shape(Function &F, bool OptimizeFrame = false)
-      : OptimizeFrame(OptimizeFrame) {
-    buildFrom(F);
-  }
-  void buildFrom(Function &F);
 };
 
 } // end namespace coro