[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

20

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

147

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

152

Typo at the end (have been?)

[dstuttard]

Thanks for the speedy review Jannik
Changes made

[jsilvanus]

llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp

149

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 {
 
 // RematGraph is used to construct a DAG for rematerializable instructions
 // When the constructor is invoked with a candidate instruction (which is
 // materializable) it builds a DAG of materializable instructions from that
 // point.
 // Typically, for each instruction identified as re-materializable across a
 // suspend point, a RematGraph will be created.
 struct RematGraph {
   // Each RematNode in the graph contains the edges to instructions providing
   // operands in the current node.
   struct RematNode {
     Instruction *Node;
     SmallVector<RematNode *> Operands;
     RematNode() = default;
     RematNode(Instruction *V) : Node(V) {}
   };
 
   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(const 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))
       return;
 
     // 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->Operands.push_back(Remats[D].get());
         continue;
       }
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);
 

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,
+    const 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)));
     }

llvm/unittests/Transforms/CMakeLists.txt

+add_subdirectory(Coroutines)
 add_subdirectory(IPO)
 add_subdirectory(Scalar)
 add_subdirectory(Utils)
 add_subdirectory(Vectorize)

llvm/unittests/Transforms/Coroutines/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  Analysis
+  AsmParser
+  Core
+  Coroutines
+  Passes
+  Support
+  TargetParser
+  TransformUtils
+  )
+
+add_llvm_unittest(CoroTests
+  ExtraRematTest.cpp
+  )
+
+target_link_libraries(CoroTests PRIVATE LLVMTestingSupport)
+
+set_property(TARGET CoroTests PROPERTY FOLDER "Tests/UnitTests/TransformTests")

llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp

+//===- ExtraRematTest.cpp - Coroutines unit tests -------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Testing/Support/Error.h"
+#include "llvm/Transforms/Coroutines/CoroSplit.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+struct ExtraRematTest : public testing::Test {
+  LLVMContext Ctx;
+  ModulePassManager MPM;
+  PassBuilder PB;
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+
+  ExtraRematTest() {
+    PB.registerModuleAnalyses(MAM);
+    PB.registerCGSCCAnalyses(CGAM);
+    PB.registerFunctionAnalyses(FAM);
+    PB.registerLoopAnalyses(LAM);
+    PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  }
+
+  BasicBlock *getBasicBlockByName(Function *F, StringRef Name) const {
+    for (BasicBlock &BB : *F) {
+      if (BB.getName() == Name)
+        return &BB;
+    }
+    return nullptr;
+  }
+
+  CallInst *getCallByName(BasicBlock *BB, StringRef Name) const {
+    for (Instruction &I : *BB) {
+      if (CallInst *CI = dyn_cast<CallInst>(&I))
+        if (CI->getCalledFunction()->getName() == Name)
+          return CI;
+    }
+    return nullptr;
+  }
+
+  void ParseAssembly(const StringRef IR) {
+    SMDiagnostic Error;
+    M = parseAssemblyString(IR, Error, Context);
+    std::string errMsg;
+    raw_string_ostream os(errMsg);
+    Error.print("", os);
+
+    // A failure here means that the test itself is buggy.
+    if (!M)
+      report_fatal_error(os.str().c_str());
+  }
+};
+
+StringRef Text = R"(
+    define ptr @f(i32 %n) presplitcoroutine {
+    entry:
+      %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 ptr @llvm.coro.begin(token %id, ptr %alloc)
+
+      %inc1 = add i32 %n, 1
+      %val2 = call i32 @should.remat(i32 %inc1)
+      %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+      switch i8 %sp1, label %suspend [i8 0, label %resume1
+                                      i8 1, label %cleanup]
+    resume1:
+      %inc2 = add i32 %val2, 1
+      %sp2 = call i8 @llvm.coro.suspend(token none, i1 false)
+      switch i8 %sp1, label %suspend [i8 0, label %resume2
+                                      i8 1, label %cleanup]
+
+    resume2:
+      call void @print(i32 %val2)
+      call void @print(i32 %inc2)
+      br label %cleanup
+
+    cleanup:
+      %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+      call void @free(ptr %mem)
+      br label %suspend
+    suspend:
+      call i1 @llvm.coro.end(ptr %hdl, i1 0)
+      ret ptr %hdl
+    }
+
+    declare ptr @llvm.coro.free(token, ptr)
+    declare i32 @llvm.coro.size.i32()
+    declare i8  @llvm.coro.suspend(token, i1)
+    declare void @llvm.coro.resume(ptr)
+    declare void @llvm.coro.destroy(ptr)
+
+    declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+    declare i1 @llvm.coro.alloc(token)
+    declare ptr @llvm.coro.begin(token, ptr)
+    declare i1 @llvm.coro.end(ptr, i1)
+
+    declare i32 @should.remat(i32)
+
+    declare noalias ptr @malloc(i32)
+    declare void @print(i32)
+    declare void @free(ptr)
+  )";
+
+// Materializable callback with extra rematerialization
+bool ExtraMaterializable(Instruction &I) {
+  if (isa<CastInst>(&I) || isa<GetElementPtrInst>(&I) ||
+      isa<BinaryOperator>(&I) || isa<CmpInst>(&I) || isa<SelectInst>(&I))
+    return true;
+
+  if (auto *CI = dyn_cast<CallInst>(&I)) {
+    auto *CalledFunc = CI->getCalledFunction();
+    if (CalledFunc && CalledFunc->getName().startswith("should.remat"))
+      return true;
+  }
+
+  return false;
+}
+
+TEST_F(ExtraRematTest, TestCoroRematDefault) {
+  ParseAssembly(Text);
+
+  ASSERT_TRUE(M);
+
+  CGSCCPassManager CGPM;
+  CGPM.addPass(CoroSplitPass());
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  MPM.run(*M, MAM);
+
+  // Verify that extra rematerializable instruction has been rematerialized
+  Function *F = M->getFunction("f.resume");
+  ASSERT_TRUE(F) << "could not find split function f.resume";
+
+  BasicBlock *Resume1 = getBasicBlockByName(F, "resume1");
+  ASSERT_TRUE(Resume1)
+      << "could not find expected BB resume1 in split function";
+
+  // With default materialization the intrinsic should not have been
+  // rematerialized
+  CallInst *CI = getCallByName(Resume1, "should.remat");
+  ASSERT_FALSE(CI);
+}
+
+TEST_F(ExtraRematTest, TestCoroRematWithCallback) {
+  ParseAssembly(Text);
+
+  ASSERT_TRUE(M);
+
+  CGSCCPassManager CGPM;
+  CGPM.addPass(
+      CoroSplitPass(std::function<bool(Instruction &)>(ExtraMaterializable)));
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  MPM.run(*M, MAM);
+
+  // Verify that extra rematerializable instruction has been rematerialized
+  Function *F = M->getFunction("f.resume");
+  ASSERT_TRUE(F) << "could not find split function f.resume";
+
+  BasicBlock *Resume1 = getBasicBlockByName(F, "resume1");
+  ASSERT_TRUE(Resume1)
+      << "could not find expected BB resume1 in split function";
+
+  // With callback the extra rematerialization of the function should have
+  // happened
+  CallInst *CI = getCallByName(Resume1, "should.remat");
+  ASSERT_TRUE(CI);
+}
+} // namespace