[Couroutines] Modify CoroFrame materializable into a callback

[dstuttard]

Summary

This change makes it possible to optionally provide a different callback to
determine if an instruction is materializable.

By default the behaviour is unchanged.

[sebastian-ne]

LGTM

[jsilvanus]

LGTM, one minor inline comment.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1935

Maybe pass the callback by reference, both here and for the following calls?

An std::function is quite large (32 bytes for libstdc++, 40 for MSVC, 48 for libc++).

[dstuttard]

Addressing reviewer comments

[dstuttard]

Rebased

[ChuanqiXu]

While this is innocent and no harmful, I prefer to land this after we found a usage for this someday. It wouldn't be a problem to do such a NFC change any times.

[dstuttard]

Rebase on fixed version

[dstuttard]

In D142621#4114357, @ChuanqiXu wrote:

While this is innocent and no harmful, I prefer to land this after we found a usage for this someday. It wouldn't be a problem to do such a NFC change any times.

We have a use-case for this change for one of our backends.

[dstuttard]

Address review commit (sorry Jannik - missed it before)

[dstuttard]

In D142621#4115184, @dstuttard wrote:

In D142621#4114357, @ChuanqiXu wrote:

While this is innocent and no harmful, I prefer to land this after we found a usage for this someday. It wouldn't be a problem to do such a NFC change any times.

We have a use-case for this change for one of our backends.

To expand on this - we have a tool based on LLVM that already adds extra materialization options. This change was to enable generic support to extend materialization determination so that we can simplify our code to not have patches on top of upstream llvm.
I can imagine other groups having similar requirements - none of which would necessarily end up back in core llvm.

[ChuanqiXu]

In D142621#4115723, @dstuttard wrote:

In D142621#4115184, @dstuttard wrote:

In D142621#4114357, @ChuanqiXu wrote:

While this is innocent and no harmful, I prefer to land this after we found a usage for this someday. It wouldn't be a problem to do such a NFC change any times.

We have a use-case for this change for one of our backends.

To expand on this - we have a tool based on LLVM that already adds extra materialization options. This change was to enable generic support to extend materialization determination so that we can simplify our code to not have patches on top of upstream llvm.
I can imagine other groups having similar requirements - none of which would necessarily end up back in core llvm.

Got it. But it is still not good to modify the upstream due the requirement from the downstream. (I know this happens in llvm occasionally. But it is still not good, right?)

I think it may better for you to contribute the other materialization options you mentioned. Or you can try to add a unittest if you don't want to contribute the downstream extension to public for any reason . Otherwise, even all people here get in consensus. It is possible that other developers found the coding here is redundant and optimize it with a NFC patch directly. (This happens a lot in llvm)

[dstuttard]

Add unit test demonstrating use of extra rematerialization callback

[dstuttard]

Got it. But it is still not good to modify the upstream due the requirement from the downstream. (I know this happens in llvm occasionally. But it is still not good, right?)

I think it may better for you to contribute the other materialization options you mentioned. Or you can try to add a unittest if you don't want to contribute the downstream extension to public for any reason . Otherwise, even all people here get in consensus. It is possible that other developers found the coding here is redundant and optimize it with a NFC patch directly. (This happens a lot in llvm)

Ok - I added a unit test for this - adding the other rematerializations isn't an option since they are quite project specific. (To be fair, I should probably have created a unit test in the first place).

[jsilvanus]

Thanks for the unit test!

llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp

19 ↗

(On Diff #496412)

Maybe put the struct (or everything) into an anonymous namespace?

146 ↗

(On Diff #496412)

I think this should use a gtest macro, e.g. ASSERT_TRUE. Same for the other asserts.

151 ↗

(On Diff #496412)

Typo at the end (have been?)

[dstuttard]

Thanks for the speedy review Jannik
Changes made

[jsilvanus]

llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp

148 ↗

(On Diff #496430)

I believe the intended way to pass custom messages is like this:

ASSERT_TRUE(F) << "could not find split function f.resume";

[dstuttard]

Fix up ASSERT_TRUE to pass messages correctly

[jsilvanus]

Thanks!

[ChuanqiXu]

LGTM. Thanks.

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

 
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 struct CoroSplitPass : PassInfoMixin<CoroSplitPass> {
-  CoroSplitPass(bool OptimizeFrame = false) : OptimizeFrame(OptimizeFrame) {}
+  const std::function<bool(Instruction &)> MaterializableCallback;
+
+  CoroSplitPass(bool OptimizeFrame = false);
+  CoroSplitPass(std::function<bool(Instruction &)> MaterializableCallback,
+                bool OptimizeFrame = false)
+      : MaterializableCallback(MaterializableCallback),
+        OptimizeFrame(OptimizeFrame) {}
 
   PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR);
   static bool isRequired() { return true; }
 
   // Would be true if the Optimization level isn't O0.
   bool OptimizeFrame;
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_COROUTINES_COROSPLIT_H

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

for (size_t I = 0; I < N; ++I) {
           LLVM_DEBUG(dump("SavedCons", SavedConsumes));
         }
       }
     }
   } while (Changed);
   LLVM_DEBUG(dump());
 }
 
-static bool materializable(Instruction &V);
-
 namespace {
 struct RematNode {
   Instruction *Node;
   SmallVector<RematNode *> Children;
   RematNode() = default;
   RematNode(Instruction *V) : Node(V) {}
 };
 
 struct RematGraph {
   RematNode *EntryNode;
   using RematNodeMap =
       SmallMapVector<Instruction *, std::unique_ptr<RematNode>, 8>;
   RematNodeMap Remats;
+  const std::function<bool(Instruction &)> MaterializableCallback;
   SuspendCrossingInfo &Checker;
 
-  RematGraph(Instruction *I, SuspendCrossingInfo &Checker) : Checker(Checker) {
+  RematGraph(std::function<bool(Instruction &)> MaterializableCallback,
+             Instruction *I, SuspendCrossingInfo &Checker)
+      : MaterializableCallback(MaterializableCallback), Checker(Checker) {
     std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);
     EntryNode = FirstNode.get();
     std::deque<std::unique_ptr<RematNode>> WorkList;
     addNode(std::move(FirstNode), WorkList, cast<User>(I));
     while (WorkList.size()) {
       std::unique_ptr<RematNode> N = std::move(WorkList.front());
       WorkList.pop_front();
       addNode(std::move(N), WorkList, cast<User>(I));
     }
   }
 
   void addNode(std::unique_ptr<RematNode> NUPtr,
                std::deque<std::unique_ptr<RematNode>> &WorkList,
                User *FirstUse) {
     RematNode *N = NUPtr.get();
     if (!Remats.count(N->Node)) {
       // We haven't see this node yet - add to the list
       Remats[N->Node] = std::move(NUPtr);
       for (auto &Def : N->Node->operands()) {
         Instruction *D = dyn_cast<Instruction>(Def.get());
-        if (!D || !materializable(*D) ||
+        if (!D || !MaterializableCallback(*D) ||
             !Checker.isDefinitionAcrossSuspend(*D, FirstUse))
           continue;
 
         if (Remats.count(D)) {
           // Already have this in the graph
           N->Children.push_back(Remats[D].get());
         } else {
           bool NoMatch = true;
           for (auto &I : WorkList) {
             if (I->Node == D) {
               NoMatch = false;
               N->Children.push_back(I.get());
               break;
             }
           }
+
           if (NoMatch) {
             // Create a new node
             std::unique_ptr<RematNode> ChildNode =
                 std::make_unique<RematNode>(D);
             N->Children.push_back(ChildNode.get());
             WorkList.push_back(std::move(ChildNode));
           }
         }
for (BasicBlock &BB : F)
     if (auto *PN = dyn_cast<PHINode>(&BB.front()))
       if (PN->getNumIncomingValues() > 1)
         WorkList.push_back(&BB);
 
   for (BasicBlock *BB : WorkList)
     rewritePHIs(*BB);
 }
 
+/// Default materializable callback
 // Check for instructions that we can recreate on resume as opposed to spill
 // the result into a coroutine frame.
-static bool materializable(Instruction &V) {
-  return isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
-         isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V);
+bool coro::defaultMaterializable(Instruction &V) {
+  return (isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
+          isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));
 }
 
 // Check for structural coroutine intrinsics that should not be spilled into
 // the coroutine frame.
 static bool isCoroutineStructureIntrinsic(Instruction &I) {
   return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I) ||
          isa<CoroSuspendInst>(&I);
 }
if (auto *I = dyn_cast<Instruction>(Storage))
       InsertPt = I->getInsertionPointAfterDef();
     else if (isa<Argument>(Storage))
       InsertPt = &*F->getEntryBlock().begin();
     if (InsertPt)
       DVI->moveBefore(InsertPt);
   }
 }
 
-static void doRematerializations(Function &F, SuspendCrossingInfo &Checker) {
+static void doRematerializations(
+    Function &F, SuspendCrossingInfo &Checker,
+    const std::function<bool(Instruction &)> MaterializableCallback) {
   SpillInfo Spills;
 
   // See if there are materializable instructions across suspend points
   // We record these as the starting point to also identify materializable
   // defs of uses in these operations
   for (Instruction &I : instructions(F)) {
-    if (!materializable(I))
+    if (!MaterializableCallback(I))
       continue;
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U))
         Spills[&I].push_back(cast<Instruction>(U));
   }
 
   // Process each of the identified rematerializable instructions
   // and add predecessor instructions that can also be rematerialized.
static void doRematerializations(
   for (auto &E : Spills) {
     for (Instruction *U : E.second) {
       // Don't process a user twice (this can happen if the instruction uses
       // more than one rematerializable def)
       if (AllRemats.count(U))
         continue;
 
       // Constructor creates the whole RematGraph for the given Use
-      auto RematUPtr = std::make_unique<RematGraph>(U, Checker);
+      auto RematUPtr =
+          std::make_unique<RematGraph>(MaterializableCallback, U, Checker);
 
       LLVM_DEBUG(dbgs() << "***** Next remat group *****\n";
                  ReversePostOrderTraversal<RematGraph *> RPOT(RematUPtr.get());
                  for (auto I = RPOT.begin(); I != RPOT.end();
                       ++I) { (*I)->Node->dump(); } dbgs()
                  << "\n";);
 
       AllRemats[U] = std::move(RematUPtr);
     }
   }
 
   // Rewrite materializable instructions to be materialized at the use
   // point.
   LLVM_DEBUG(dumpRemats("Materializations", AllRemats));
   rewriteMaterializableInstructions(AllRemats);
 }
 
-void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
+void coro::buildCoroutineFrame(
+    Function &F, Shape &Shape,
+    const std::function<bool(Instruction &)> MaterializableCallback) {
   // Don't eliminate swifterror in async functions that won't be split.
   if (Shape.ABI != coro::ABI::Async || !Shape.CoroSuspends.empty())
     eliminateSwiftError(F, Shape);
 
   if (Shape.ABI == coro::ABI::Switch &&
       Shape.SwitchLowering.PromiseAlloca) {
     Shape.getSwitchCoroId()->clearPromise();
   }
void coro::buildCoroutineFrame(
 
   // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
   // never has its definition separated from the PHI by the suspend point.
   rewritePHIs(F);
 
   // Build suspend crossing info.
   SuspendCrossingInfo Checker(F, Shape);
 
-  doRematerializations(F, Checker);
+  doRematerializations(F, Checker, MaterializableCallback);
 
   FrameDataInfo FrameData;
   SmallVector<CoroAllocaAllocInst*, 4> LocalAllocas;
   SmallVector<Instruction*, 4> DeadInstructions;
+
   if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
       Shape.ABI != coro::ABI::RetconOnce)
     sinkLifetimeStartMarkers(F, Shape, Checker);
 
   // Collect the spills for arguments and other not-materializable values.
   for (Argument &A : F.args())
     for (User *U : A.users())
       if (Checker.isDefinitionAcrossSuspend(A, U))

llvm/lib/Transforms/Coroutines/CoroInternal.h

struct LLVM_LIBRARY_VISIBILITY Shape {
   Shape() = default;
   explicit Shape(Function &F, bool OptimizeFrame = false)
       : OptimizeFrame(OptimizeFrame) {
     buildFrom(F);
   }
   void buildFrom(Function &F);
 };
 
-void buildCoroutineFrame(Function &F, Shape &Shape);
+bool defaultMaterializable(Instruction &V);
+void buildCoroutineFrame(
+    Function &F, Shape &Shape,
+    std::function<bool(Instruction &)> MaterializableCallback);
 CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                              ArrayRef<Value *> Arguments, IRBuilder<> &);
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

public:
     void print(raw_ostream &OS) const override {
       OS << "While splitting coroutine ";
       F.printAsOperand(OS, /*print type*/ false, F.getParent());
       OS << "\n";
     }
   };
 }
 
-static coro::Shape splitCoroutine(Function &F,
-                                  SmallVectorImpl<Function *> &Clones,
-                                  TargetTransformInfo &TTI,
-                                  bool OptimizeFrame) {
+static coro::Shape
+splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
+               TargetTransformInfo &TTI, bool OptimizeFrame,
+               std::function<bool(Instruction &)> MaterializableCallback) {
   PrettyStackTraceFunction prettyStackTrace(F);
 
   // The suspend-crossing algorithm in buildCoroutineFrame get tripped
   // up by uses in unreachable blocks, so remove them as a first pass.
   removeUnreachableBlocks(F);
 
   coro::Shape Shape(F, OptimizeFrame);
   if (!Shape.CoroBegin)
     return Shape;
 
   simplifySuspendPoints(Shape);
-  buildCoroutineFrame(F, Shape);
+  buildCoroutineFrame(F, Shape, MaterializableCallback);
   replaceFrameSizeAndAlignment(Shape);
 
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
     handleNoSuspendCoroutine(Shape);
   } else {
     switch (Shape.ABI) {

 static void addPrepareFunction(const Module &M,
                                SmallVectorImpl<Function *> &Fns,
                                StringRef Name) {
   auto *PrepareFn = M.getFunction(Name);
   if (PrepareFn && !PrepareFn->use_empty())
     Fns.push_back(PrepareFn);
 }
 
+CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
+    : MaterializableCallback(coro::defaultMaterializable),
+      OptimizeFrame(OptimizeFrame) {}
+
 PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
                                      CGSCCAnalysisManager &AM,
                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
   // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a
   //     non-zero number of nodes, so we assume that here and grab the first
   //     node's function's module.
   Module &M = *C.begin()->getFunction().getParent();
   auto &FAM =
PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
     LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
                       << "\n");
     F.setSplittedCoroutine();
 
     SmallVector<Function *, 4> Clones;
-    const coro::Shape Shape = splitCoroutine(
-        F, Clones, FAM.getResult<TargetIRAnalysis>(F), OptimizeFrame);
+    const coro::Shape Shape =
+        splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
+                       OptimizeFrame, MaterializableCallback);
     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
 
     if (!Shape.CoroSuspends.empty()) {
       // Run the CGSCC pipeline on the original and newly split functions.
       UR.CWorklist.insert(&C);
       for (Function *Clone : Clones)
         UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
     }