[llvm] [Coroutines][NFC] Refactor CoroCloner (PR #116885)

Tyler Nowicki via llvm-commits llvm-commits at lists.llvm.org
Wed Nov 20 06:33:50 PST 2024


https://github.com/TylerNowicki updated https://github.com/llvm/llvm-project/pull/116885

>From 69431dc2bc5c5bbe1fabfc8a65e399a5cf51af71 Mon Sep 17 00:00:00 2001
From: tnowicki <tyler.nowicki at amd.com>
Date: Tue, 19 Nov 2024 17:59:24 -0500
Subject: [PATCH 1/2] [Coroutines][NFC] Refactor CoroCloner

* Move CoroCloner to its own header. For now the header is located in
  lib/Transforms/Coroutines
* Change private to protected to allow inheritance
* Create CoroSwitchCloner and move some of the switch specific code into
  this cloner.
* This continus previous work done to CoroFrame. The goal is to isolate
  the cloner code required for each ABI as much as is reasonable. The
  base class should include a useful set of helpers and base
  implementations. More changes will follow.
---
 llvm/lib/Transforms/Coroutines/CoroCloner.h  | 147 ++++++++++++++++++
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 149 ++-----------------
 2 files changed, 163 insertions(+), 133 deletions(-)
 create mode 100644 llvm/lib/Transforms/Coroutines/CoroCloner.h

diff --git a/llvm/lib/Transforms/Coroutines/CoroCloner.h b/llvm/lib/Transforms/Coroutines/CoroCloner.h
new file mode 100644
index 00000000000000..e4e3976924e59a
--- /dev/null
+++ b/llvm/lib/Transforms/Coroutines/CoroCloner.h
@@ -0,0 +1,147 @@
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// Helper class for splitting a coroutine into separate functions. For example
+// the returned-continuation coroutine is split into separate continuation
+// functions.
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/TimeProfiler.h"
+#include "llvm/Transforms/Coroutines/ABI.h"
+#include "llvm/Transforms/Coroutines/CoroInstr.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
+
+using namespace llvm;
+
+class CoroCloner {
+public:
+  enum class Kind {
+    /// The shared resume function for a switch lowering.
+    SwitchResume,
+
+    /// The shared unwind function for a switch lowering.
+    SwitchUnwind,
+
+    /// The shared cleanup function for a switch lowering.
+    SwitchCleanup,
+
+    /// An individual continuation function.
+    Continuation,
+
+    /// An async resume function.
+    Async,
+  };
+
+protected:
+  Function &OrigF;
+  const Twine &Suffix;
+  coro::Shape &Shape;
+  Kind FKind;
+  IRBuilder<> Builder;
+  TargetTransformInfo &TTI;
+
+  ValueToValueMapTy VMap;
+  Function *NewF = nullptr;
+  Value *NewFramePtr = nullptr;
+
+
+  /// The active suspend instruction; meaningful only for continuation and async
+  /// ABIs.
+  AnyCoroSuspendInst *ActiveSuspend = nullptr;
+
+  /// Create a cloner for a continuation lowering.
+  CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
+             Function *NewF, AnyCoroSuspendInst *ActiveSuspend,
+             TargetTransformInfo &TTI)
+      : OrigF(OrigF), Suffix(Suffix), Shape(Shape),
+        FKind(Shape.ABI == coro::ABI::Async ? Kind::Async : Kind::Continuation),
+        Builder(OrigF.getContext()), TTI(TTI), NewF(NewF), ActiveSuspend(ActiveSuspend)  {
+    assert(Shape.ABI == coro::ABI::Retcon ||
+           Shape.ABI == coro::ABI::RetconOnce || Shape.ABI == coro::ABI::Async);
+    assert(NewF && "need existing function for continuation");
+    assert(ActiveSuspend && "need active suspend point for continuation");
+  }
+
+public:
+  CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
+      Kind FKind, TargetTransformInfo &TTI) :
+    OrigF(OrigF), Suffix(Suffix), Shape(Shape), FKind(FKind), Builder(OrigF.getContext()), TTI(TTI) {
+    }
+
+  virtual ~CoroCloner() { }
+
+  /// Create a clone for a continuation lowering.
+  static Function *createClone(Function &OrigF, const Twine &Suffix,
+                               coro::Shape &Shape, Function *NewF,
+                               AnyCoroSuspendInst *ActiveSuspend,
+                               TargetTransformInfo &TTI) {
+    assert(Shape.ABI == coro::ABI::Retcon ||
+           Shape.ABI == coro::ABI::RetconOnce || Shape.ABI == coro::ABI::Async);
+    TimeTraceScope FunctionScope("CoroCloner");
+
+    CoroCloner Cloner(OrigF, Suffix, Shape, NewF, ActiveSuspend, TTI);
+    Cloner.create();
+    return Cloner.getFunction();
+  }
+
+  Function *getFunction() const {
+    assert(NewF != nullptr && "declaration not yet set");
+    return NewF;
+  }
+
+  virtual void create();
+
+protected:
+  bool isSwitchDestroyFunction() {
+    switch (FKind) {
+    case Kind::Async:
+    case Kind::Continuation:
+    case Kind::SwitchResume:
+      return false;
+    case Kind::SwitchUnwind:
+    case Kind::SwitchCleanup:
+      return true;
+    }
+    llvm_unreachable("Unknown CoroCloner::Kind enum");
+  }
+
+  void replaceEntryBlock();
+  Value *deriveNewFramePointer();
+  void replaceRetconOrAsyncSuspendUses();
+  void replaceCoroSuspends();
+  void replaceCoroEnds();
+  void replaceSwiftErrorOps();
+  void salvageDebugInfo();
+  void handleFinalSuspend();
+};
+
+
+class CoroSwitchCloner : public CoroCloner {
+protected:
+  /// Create a cloner for a switch lowering.
+  CoroSwitchCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
+                   Kind FKind, TargetTransformInfo &TTI)
+      : CoroCloner(OrigF, Suffix, Shape, FKind, TTI) { }
+
+  void create() override;
+
+public:
+  /// Create a clone for a switch lowering.
+  static Function *createClone(Function &OrigF, const Twine &Suffix,
+                               coro::Shape &Shape, Kind FKind,
+                               TargetTransformInfo &TTI) {
+    assert(Shape.ABI == coro::ABI::Switch);
+    TimeTraceScope FunctionScope("CoroCloner");
+
+    CoroSwitchCloner Cloner(OrigF, Suffix, Shape, FKind, TTI);
+    Cloner.create();
+    return Cloner.getFunction();
+  }
+};
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 55951e54518bd2..a4406f8364bfd9 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -20,6 +20,7 @@
 
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "CoroInternal.h"
+#include "CoroCloner.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/ADT/STLExtras.h"
@@ -44,10 +45,8 @@
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
-#include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
-#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
@@ -61,17 +60,13 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/PrettyStackTrace.h"
-#include "llvm/Support/TimeProfiler.h"
 #include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/Coroutines/ABI.h"
-#include "llvm/Transforms/Coroutines/CoroInstr.h"
 #include "llvm/Transforms/Coroutines/MaterializationUtils.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
-#include "llvm/Transforms/Utils/ValueMapper.h"
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -82,122 +77,6 @@ using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
-namespace {
-
-/// A little helper class for building
-class CoroCloner {
-public:
-  enum class Kind {
-    /// The shared resume function for a switch lowering.
-    SwitchResume,
-
-    /// The shared unwind function for a switch lowering.
-    SwitchUnwind,
-
-    /// The shared cleanup function for a switch lowering.
-    SwitchCleanup,
-
-    /// An individual continuation function.
-    Continuation,
-
-    /// An async resume function.
-    Async,
-  };
-
-private:
-  Function &OrigF;
-  Function *NewF;
-  const Twine &Suffix;
-  coro::Shape &Shape;
-  Kind FKind;
-  ValueToValueMapTy VMap;
-  IRBuilder<> Builder;
-  Value *NewFramePtr = nullptr;
-
-  /// The active suspend instruction; meaningful only for continuation and async
-  /// ABIs.
-  AnyCoroSuspendInst *ActiveSuspend = nullptr;
-
-  TargetTransformInfo &TTI;
-
-  /// Create a cloner for a switch lowering.
-  CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
-             Kind FKind, TargetTransformInfo &TTI)
-      : OrigF(OrigF), NewF(nullptr), Suffix(Suffix), Shape(Shape), FKind(FKind),
-        Builder(OrigF.getContext()), TTI(TTI) {
-    assert(Shape.ABI == coro::ABI::Switch);
-  }
-
-  /// Create a cloner for a continuation lowering.
-  CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
-             Function *NewF, AnyCoroSuspendInst *ActiveSuspend,
-             TargetTransformInfo &TTI)
-      : OrigF(OrigF), NewF(NewF), Suffix(Suffix), Shape(Shape),
-        FKind(Shape.ABI == coro::ABI::Async ? Kind::Async : Kind::Continuation),
-        Builder(OrigF.getContext()), ActiveSuspend(ActiveSuspend), TTI(TTI) {
-    assert(Shape.ABI == coro::ABI::Retcon ||
-           Shape.ABI == coro::ABI::RetconOnce || Shape.ABI == coro::ABI::Async);
-    assert(NewF && "need existing function for continuation");
-    assert(ActiveSuspend && "need active suspend point for continuation");
-  }
-
-public:
-  /// Create a clone for a switch lowering.
-  static Function *createClone(Function &OrigF, const Twine &Suffix,
-                               coro::Shape &Shape, Kind FKind,
-                               TargetTransformInfo &TTI) {
-    TimeTraceScope FunctionScope("CoroCloner");
-
-    CoroCloner Cloner(OrigF, Suffix, Shape, FKind, TTI);
-    Cloner.create();
-    return Cloner.getFunction();
-  }
-
-  /// Create a clone for a continuation lowering.
-  static Function *createClone(Function &OrigF, const Twine &Suffix,
-                               coro::Shape &Shape, Function *NewF,
-                               AnyCoroSuspendInst *ActiveSuspend,
-                               TargetTransformInfo &TTI) {
-    TimeTraceScope FunctionScope("CoroCloner");
-
-    CoroCloner Cloner(OrigF, Suffix, Shape, NewF, ActiveSuspend, TTI);
-    Cloner.create();
-    return Cloner.getFunction();
-  }
-
-  Function *getFunction() const {
-    assert(NewF != nullptr && "declaration not yet set");
-    return NewF;
-  }
-
-  void create();
-
-private:
-  bool isSwitchDestroyFunction() {
-    switch (FKind) {
-    case Kind::Async:
-    case Kind::Continuation:
-    case Kind::SwitchResume:
-      return false;
-    case Kind::SwitchUnwind:
-    case Kind::SwitchCleanup:
-      return true;
-    }
-    llvm_unreachable("Unknown CoroCloner::Kind enum");
-  }
-
-  void replaceEntryBlock();
-  Value *deriveNewFramePointer();
-  void replaceRetconOrAsyncSuspendUses();
-  void replaceCoroSuspends();
-  void replaceCoroEnds();
-  void replaceSwiftErrorOps();
-  void salvageDebugInfo();
-  void handleFinalSuspend();
-};
-
-} // end anonymous namespace
-
 // FIXME:
 // Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape
 // and it is known that other transformations, for example, sanitizers
@@ -985,11 +864,7 @@ static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,
 /// Clone the body of the original function into a resume function of
 /// some sort.
 void CoroCloner::create() {
-  // Create the new function if we don't already have one.
-  if (!NewF) {
-    NewF = createCloneDeclaration(OrigF, Shape, Suffix,
-                                  OrigF.getParent()->end(), ActiveSuspend);
-  }
+  assert(NewF);
 
   // Replace all args with dummy instructions. If an argument is the old frame
   // pointer, the dummy will be replaced by the new frame pointer once it is
@@ -1213,12 +1088,20 @@ void CoroCloner::create() {
 
   // Salvage debug info that points into the coroutine frame.
   salvageDebugInfo();
+}
+
+void CoroSwitchCloner::create() {
+  // Create a new function matching the original type
+  NewF = createCloneDeclaration(OrigF, Shape, Suffix,
+      OrigF.getParent()->end(), ActiveSuspend);
+
+  // Clone the function
+  CoroCloner::create();
 
   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
   // to suppress deallocation code.
-  if (Shape.ABI == coro::ABI::Switch)
-    coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
-                          /*Elide=*/FKind == CoroCloner::Kind::SwitchCleanup);
+  coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
+      /*Elide=*/FKind == CoroCloner::Kind::SwitchCleanup);
 }
 
 static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
@@ -1495,11 +1378,11 @@ struct SwitchCoroutineSplitter {
     // setting new entry block and replacing coro.suspend an appropriate value
     // to force resume or cleanup pass for every suspend point.
     createResumeEntryBlock(F, Shape);
-    auto *ResumeClone = CoroCloner::createClone(
+    auto *ResumeClone = CoroSwitchCloner::createClone(
         F, ".resume", Shape, CoroCloner::Kind::SwitchResume, TTI);
-    auto *DestroyClone = CoroCloner::createClone(
+    auto *DestroyClone = CoroSwitchCloner::createClone(
         F, ".destroy", Shape, CoroCloner::Kind::SwitchUnwind, TTI);
-    auto *CleanupClone = CoroCloner::createClone(
+    auto *CleanupClone = CoroSwitchCloner::createClone(
         F, ".cleanup", Shape, CoroCloner::Kind::SwitchCleanup, TTI);
 
     postSplitCleanup(*ResumeClone);

>From 3e54d333b9ee004866afe0be44ee67072cd00ea3 Mon Sep 17 00:00:00 2001
From: tnowicki <tyler.nowicki at amd.com>
Date: Tue, 19 Nov 2024 18:08:48 -0500
Subject: [PATCH 2/2] [Coroutines] Clang format

---
 llvm/lib/Transforms/Coroutines/CoroCloner.h  | 15 +++++++--------
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp |  8 ++++----
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Coroutines/CoroCloner.h b/llvm/lib/Transforms/Coroutines/CoroCloner.h
index e4e3976924e59a..5f861d60facbb5 100644
--- a/llvm/lib/Transforms/Coroutines/CoroCloner.h
+++ b/llvm/lib/Transforms/Coroutines/CoroCloner.h
@@ -51,7 +51,6 @@ class CoroCloner {
   Function *NewF = nullptr;
   Value *NewFramePtr = nullptr;
 
-
   /// The active suspend instruction; meaningful only for continuation and async
   /// ABIs.
   AnyCoroSuspendInst *ActiveSuspend = nullptr;
@@ -62,7 +61,8 @@ class CoroCloner {
              TargetTransformInfo &TTI)
       : OrigF(OrigF), Suffix(Suffix), Shape(Shape),
         FKind(Shape.ABI == coro::ABI::Async ? Kind::Async : Kind::Continuation),
-        Builder(OrigF.getContext()), TTI(TTI), NewF(NewF), ActiveSuspend(ActiveSuspend)  {
+        Builder(OrigF.getContext()), TTI(TTI), NewF(NewF),
+        ActiveSuspend(ActiveSuspend) {
     assert(Shape.ABI == coro::ABI::Retcon ||
            Shape.ABI == coro::ABI::RetconOnce || Shape.ABI == coro::ABI::Async);
     assert(NewF && "need existing function for continuation");
@@ -71,11 +71,11 @@ class CoroCloner {
 
 public:
   CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
-      Kind FKind, TargetTransformInfo &TTI) :
-    OrigF(OrigF), Suffix(Suffix), Shape(Shape), FKind(FKind), Builder(OrigF.getContext()), TTI(TTI) {
-    }
+             Kind FKind, TargetTransformInfo &TTI)
+      : OrigF(OrigF), Suffix(Suffix), Shape(Shape), FKind(FKind),
+        Builder(OrigF.getContext()), TTI(TTI) {}
 
-  virtual ~CoroCloner() { }
+  virtual ~CoroCloner() {}
 
   /// Create a clone for a continuation lowering.
   static Function *createClone(Function &OrigF, const Twine &Suffix,
@@ -122,13 +122,12 @@ class CoroCloner {
   void handleFinalSuspend();
 };
 
-
 class CoroSwitchCloner : public CoroCloner {
 protected:
   /// Create a cloner for a switch lowering.
   CoroSwitchCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
                    Kind FKind, TargetTransformInfo &TTI)
-      : CoroCloner(OrigF, Suffix, Shape, FKind, TTI) { }
+      : CoroCloner(OrigF, Suffix, Shape, FKind, TTI) {}
 
   void create() override;
 
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index a4406f8364bfd9..11b6dcf6f29a3b 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -19,8 +19,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
-#include "CoroInternal.h"
 #include "CoroCloner.h"
+#include "CoroInternal.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/ADT/STLExtras.h"
@@ -1092,8 +1092,8 @@ void CoroCloner::create() {
 
 void CoroSwitchCloner::create() {
   // Create a new function matching the original type
-  NewF = createCloneDeclaration(OrigF, Shape, Suffix,
-      OrigF.getParent()->end(), ActiveSuspend);
+  NewF = createCloneDeclaration(OrigF, Shape, Suffix, OrigF.getParent()->end(),
+                                ActiveSuspend);
 
   // Clone the function
   CoroCloner::create();
@@ -1101,7 +1101,7 @@ void CoroSwitchCloner::create() {
   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
   // to suppress deallocation code.
   coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
-      /*Elide=*/FKind == CoroCloner::Kind::SwitchCleanup);
+                        /*Elide=*/FKind == CoroCloner::Kind::SwitchCleanup);
 }
 
 static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {



More information about the llvm-commits mailing list