Fix PR40710: Outlined Function has token parameter but isn't an intrinsic

[hiraditya]

Summary

Code extractor outlines function with token parameter which isn't an intrinsic. This patch:

• Adds a verifier to check valid parameter types.
• Allows hotcoldsplit to bail out once input parameters are calculated (before splitting)
• Adds a sanity check using the verifier in debug mode (after splitting)

https://github.com/llvm/llvm-project/issues/40056

[compnerd]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

344 ↗

(On Diff #222327)

For some reason it feels like an odd name to me. What do you think of something like validateDataDependencies or something along those lines? I think the fact that input is slightly fuzzy is why it just feels odd, but, I can understand where that comes from too.

llvm/test/Transforms/HotColdSplit/token-arg.ll

2

Might not be a bad idea to ensure that the transformation is idempotent - that is to have checks on the output IR

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

344 ↗

(On Diff #222327)

Sounds good.

llvm/test/Transforms/HotColdSplit/token-arg.ll

2

sure

[hiraditya]

Comments from @compnerd

[vsk]

Thanks for working on this!

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

151–154

The validation done by validateInputDataDependencies should be moved into isEligible. This way, extractCodeRegion will simply fail on non-eligible regions -- this is something clients of CodeExtractor can already handle.

Adding a new validation API may leave some clients, like PartialInliner, unfixed.

[hiraditya]

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

151–154

Makes sense, i'll update this.

[hiraditya]

Addressed comments from @vsk

• Move input parameter validation to isEligible().

This now requires recomputing inputs in extractCodeRegion because PHI nodes, and return blocks are split. The computation cost of input is same as before because previously the input parameter would be computed on the client side (caller of extractCodeRegion) and introduce extra overhead to be taken care of.

[vsk]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

598

Why not set a bit inside of the CodeExtractor instance the first time findInputsOutputs runs? Then there's no need to duplicate the 'findInputs' logic.

[hiraditya]

Do we have further feedback for this patch?

[vsk]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1388

On closer inspection, I don't think my suggestion for storing an 'are inputs eligible' bit inside of CodeExtractor after findInputsOutputs runs the first time is a great idea (and I see that you haven't followed that advice here). The problem with that approach is that it requires clients to run findInputsOutputs, not modify any IR, and then call extractCodeRegion. That happens to be how clients use CodeExtractor today, but the new requirement could be brittle.

That said, I don't think the solution can be to run findInputsOutputs a third time. In addition to being expensive, it creates some confusion over the need for the operation to be repeated three times.

I think there's a simpler and cheaper alternative: run findInputsOutputs() exactly once within extractCodeRegion(). Note that splitReturnBlocks() doesn't change the input set. In severSplitPHINodesOfEntry(), you can add phis from the old header block to the input set if they have uses in the extraction region. Similarly, severSplitPHINodesOfExits() can be modified to remove exit block phis from the output set, and to add any newly created split phis.

To recap, please call findInputsOutputs() just once at the start of extractCodeRegion() and delete the subsequent call. Any input validation can happen up front (e.g. findInputsOutputs could return an llvm::Error indicating whether validation succeeded).

[hiraditya]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1388

That said, I don't think the solution can be to run findInputsOutputs a third time. In addition to being expensive, it creates some confusion over the need for the operation to be repeated three times.

One of the reasons why verification of input parameters could be a separate function and not called from isEligible here. The clients can make sanity checks before calling the extractCodeRegion. In general, any optimization pass, would have analysis phase and a transformation phase. IMHO This function should just extract the region and not do anything else. But there's already some check in place here which isn't ideal. The goal of this diff was to fix the ICE. I'm totally fine with repairing CodeExtractor in a separate patch.

[hiraditya]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1388

To recap, please call findInputsOutputs() just once at the start of extractCodeRegion() and delete the subsequent call.
Any input validation can happen up front (e.g. findInputsOutputs could return an llvm::Error indicating whether
validation succeeded).

Maybe i misunderstood but keeping the first call as above, and removing the later call (line 1442) results in the following failure. Moving findInputsOutputs anywhere before splitting PHI, or splitting return could result in miscompilation.

Failing Tests (4):

LLVM :: Transforms/HotColdSplit/duplicate-phi-preds-crash.ll
LLVM :: Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll
LLVM :: Transforms/HotColdSplit/split-phis-in-exit-blocks.ll
LLVM :: Transforms/HotColdSplit/succ-block-with-self-edge.ll

[hiraditya]

Essentially did this:

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 952263c0b1f..49fa07a3648 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -1336,13 +1336,13 @@ void CodeExtractor::calculateNewCallTerminatorWeights(
 
 Function *CodeExtractor::extractCodeRegion() {
   ValueSet inputs, outputs, SinkingCands, HoistingCands;
-  findInputsOutputs(inputs, outputs, SinkingCands, true);
+  findInputsOutputs(inputs, outputs, SinkingCands);
 
   if (!isEligible(inputs))
     return nullptr;
   // After sanity check, clear inputs as it will be recomputed
   // after CFG is prepared for splitting.
-  inputs.clear();
+  //inputs.clear();
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
@@ -1439,7 +1439,7 @@ Function *CodeExtractor::extractCodeRegion() {
   assert(HoistingCands.empty() || CommonExit);
 
   // Find inputs to, outputs from the code region.
-  findInputsOutputs(inputs, outputs, SinkingCands);
+  //findInputsOutputs(inputs, outputs, SinkingCands);
 
   // Now sink all instructions which only have non-phi uses inside the region.

and got the failures:

Failing Tests (4):

LLVM :: Transforms/HotColdSplit/duplicate-phi-preds-crash.ll
LLVM :: Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll
LLVM :: Transforms/HotColdSplit/split-phis-in-exit-blocks.ll
LLVM :: Transforms/HotColdSplit/succ-block-with-self-edge.ll

[vsk]

One of the reasons why verification of input parameters could be a separate function and not called from isEligible here. The clients can make sanity checks before calling the extractCodeRegion. In general, any optimization pass, would have analysis phase and a transformation phase. IMHO This function should just extract the region and not do anything else. But there's already some check in place here which isn't ideal. The goal of this diff was to fix the ICE. I'm totally fine with repairing CodeExtractor in a separate patch.

I'm not sure I follow. I think we're on the same page about wanting to fix CodeExtractor's handling of token inputs, but I'm not sure why a separate patch would be needed to do that.

Maybe i misunderstood but keeping the first call as above, and removing the later call (line 1442) results in the following failure. Moving findInputsOutputs anywhere before splitting PHI, or splitting return could result in miscompilation.

Oh, on this point I think we're not on the same page :). In addition to moving the call to findInputsOutputs to beginning of extractCodeRegion (where input validation must happen), I'm suggesting that we modify extractCodeRegion to update its input and output value sets as it modifies IR. I think I've outlined all the steps needed to do that in my previous comment. This would be simpler and cheaper than recomputing inputs & outputs from scratch a third time -- a side benefit is that it would probably make the phi splitting code clearer as well.

[hiraditya]

rebase against latest changes

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

public:
     /// Returns false when verifier finds errors. AssumptionCache is passed as
     /// parameter to make this function stateless.
     static bool verifyAssumptionCache(const Function& F, AssumptionCache *AC);
 
     /// Test whether this code extractor is eligible.
     ///
     /// Based on the blocks used when constructing the code extractor,
     /// determine whether it is eligible for extraction.
     /// 
     /// Checks that varargs handling (with vastart and vaend) is only done in
     /// the outlined blocks.
-    bool isEligible() const;
+    bool isEligible(const ValueSet &inputs) const;
 
     /// Compute the set of input values and output values for the code.
     ///
     /// These can be used either when performing the extraction or to evaluate
     /// the expected size of a call to the extracted function. Note that this
     /// work cannot be cached between the two as once we decide to extract
     /// a code sequence, that sequence is modified, including changing these
     /// sets, before extraction occurs. These modifications won't have any
     /// significant impact on the cost however.
     void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
-                           const ValueSet &Allocas) const;
+                           const ValueSet &Allocas,
+                           bool InputsOnly = false) const;
 
     /// Check if life time marker nodes can be hoisted/sunk into the outline
     /// region.
     ///
     /// Returns true if it is safe to do the code motion.
     bool
     isLegalToShrinkwrapLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
                                        Instruction *AllocaAddr) const;
private:
     LifetimeMarkerInfo
     getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
                        Instruction *Addr, BasicBlock *ExitBlock) const;
 
     void severSplitPHINodesOfEntry(BasicBlock *&Header);
     void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
     void splitReturnBlocks();
 
+    /// Verify if args to the new function are valid.
+    /// Returns false when any argument is invalid.
+    bool validateInputDataDependencies(const ValueSet &inputs) const;
+
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,
                                 BasicBlock *newRootNode, BasicBlock *newHeader,
                                 Function *oldFunction, Module *M);
 
     void moveCodeToFunction(Function *newFunction);
 

llvm/lib/Transforms/Utils/CodeExtractor.cpp

for (unsigned I = 0, E = Bitcasts.size(); I != E; ++I) {
         LLVM_DEBUG(dbgs() << "Sinking bitcast-of-alloca: " << *BitcastAddr
                           << "\n");
         SinkCands.insert(BitcastAddr);
       }
     }
   }
 }
 
-bool CodeExtractor::isEligible() const {
+bool CodeExtractor::isEligible(const ValueSet &inputs) const {
   if (Blocks.empty())
     return false;
   BasicBlock *Header = *Blocks.begin();
   Function *F = Header->getParent();
 
   // For functions with varargs, check that varargs handling is only done in the
   // outlined function, i.e vastart and vaend are only used in outlined blocks.
   if (AllowVarArgs && F->getFunctionType()->isVarArg()) {
     auto containsVarArgIntrinsic = [](const Instruction &I) {
       if (const CallInst *CI = dyn_cast<CallInst>(&I))
         if (const Function *Callee = CI->getCalledFunction())
           return Callee->getIntrinsicID() == Intrinsic::vastart ||
                  Callee->getIntrinsicID() == Intrinsic::vaend;
       return false;
     };
 
     for (auto &BB : *F) {
       if (Blocks.count(&BB))
         continue;
       if (llvm::any_of(BB, containsVarArgIntrinsic))
         return false;
     }
   }
-  return true;
+  return validateInputDataDependencies(inputs);
 }
 
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
-                                      const ValueSet &SinkCands) const {
+                                      const ValueSet &SinkCands,
+                                      bool InputsOnly) const {
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
     for (Instruction &II : *BB) {
       for (auto &OI : II.operands()) {
         Value *V = OI;
         if (!SinkCands.count(V) && definedInCaller(Blocks, V))
           Inputs.insert(V);
       }
-
+      if (InputsOnly)
+        return;
       for (User *U : II.users())
         if (!definedInRegion(Blocks, U)) {
           Outputs.insert(&II);
           break;
         }
     }
   }
 }
for (unsigned I = 0, E = BranchDist.Weights.size(); I < E; ++I) {
     BranchProbability BP(Weight.Amount, BranchDist.Total);
     BPI->setEdgeProbability(CodeReplacer, Weight.TargetNode.Index, BP);
   }
   TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
-Function *
-CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC) {
-  if (!isEligible())
+Function *CodeExtractor::extractCodeRegion() {
+  ValueSet inputs, outputs, SinkingCands, HoistingCands;
+  findInputsOutputs(inputs, outputs, SinkingCands, true);
+
+  if (!isEligible(inputs))
     return nullptr;
+  // After sanity check, clear inputs as it will be recomputed
+  // after CFG is prepared for splitting.
+  inputs.clear();
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
   Function *oldFunction = header->getParent();
 
   // Calculate the entry frequency of the new function before we change the root
   //   block.
any_of(Blocks, [&BranchI](const BasicBlock *BB) {
           return false;
         BranchI->setDebugLoc(I.getDebugLoc());
         return true;
       });
     });
   }
   newFuncRoot->getInstList().push_back(BranchI);
 
-  ValueSet inputs, outputs, SinkingCands, HoistingCands;
   BasicBlock *CommonExit = nullptr;
   findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
   assert(HoistingCands.empty() || CommonExit);
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs, SinkingCands);
 
   // Now sink all instructions which only have non-phi uses inside the region.
Function *CodeExtractor::extractCodeRegion() {
   LLVM_DEBUG(if (verifyFunction(*newFunction, &errs())) {
     newFunction->dump();
     report_fatal_error("verification of newFunction failed!");
   });
   LLVM_DEBUG(if (verifyFunction(*oldFunction))
              report_fatal_error("verification of oldFunction failed!"));
   LLVM_DEBUG(if (AC && verifyAssumptionCache(*oldFunction, AC))
              report_fatal_error("Stale Asumption cache for old Function!"));
+  LLVM_DEBUG(if (!validateInputDataDependencies(inputs))
+             report_fatal_error("verification of newFunction args failed!"));
   return newFunction;
 }
 
 bool CodeExtractor::verifyAssumptionCache(const Function& F,
                                           AssumptionCache *AC) {
   for (auto AssumeVH : AC->assumptions()) {
     CallInst *I = cast<CallInst>(AssumeVH);
     if (I->getFunction() != &F)
       return true;
   }
   return false;
 }
+
+bool
+CodeExtractor::validateInputDataDependencies(const ValueSet &inputs) const {
+  for (const Value *Arg : inputs) {
+    Type *Ty = Arg->getType();
+    if (!Ty->isFirstClassType() || Ty->isMetadataTy() || Ty->isTokenTy())
+      return false;
+  }
+  return true;
+}

llvm/test/Transforms/HotColdSplit/token-arg.ll

+; RUN: opt -S -hotcoldsplit < %s | FileCheck %s
+
+; CHECK-LABEL: @zot
+; CHECK-LABEL: bb
+; CHECK: invoke void @barney()
+; CHECK-LABEL: bb3
+; CHECK: call void @barney()
+; CHECK-LABEL: bb4
+; CHECK: call void @barney.1()
+; CHECK-NOT: .cold
+
+define void @zot() personality i8* bitcast (i32 (...)* @bar to i8*) {
+bb:
+  invoke void @barney()
+          to label %bb1 unwind label %bb2
+
+bb1:                                              ; preds = %bb
+  ret void
+
+bb2:                                              ; preds = %bb
+  %tmp = cleanuppad within none []
+  br label %bb3
+
+bb3:                                              ; preds = %bb2
+  call void @barney() [ "funclet"(token %tmp) ]
+  br label %bb4
+
+bb4:                                              ; preds = %bb3
+  call void @barney.1() [ "funclet"(token %tmp) ]
+  unreachable
+}
+
+declare void @barney()
+
+declare i32 @bar(...)
+
+declare void @barney.1()

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

)invalid",
                                                 Err, Ctx));
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 3> Candidates{ getBlockByName(Func, "header"),
                                            getBlockByName(Func, "body1"),
                                            getBlockByName(Func, "body2") };
 
   CodeExtractor CE(Candidates);
-  EXPECT_TRUE(CE.isEligible());
+  SetVector<Value *> Inputs, Outputs, Sinks;
+  CE.findInputsOutputs(Inputs, Outputs, Sinks, true);
+  EXPECT_TRUE(CE.isEligible(Inputs));
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit = getBlockByName(Func, "notExtracted");
   BasicBlock *ExitSplit = getBlockByName(Outlined, "notExtracted.split");
   // Ensure that PHI in exit block has only one incoming value (from code
   // replacer block).
TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "extracted1"),
     getBlockByName(Func, "extracted2")
   };
 
   CodeExtractor CE(ExtractedBlocks);
-  EXPECT_TRUE(CE.isEligible());
+  SetVector<Value *> Inputs, Outputs, Sinks;
+  CE.findInputsOutputs(Inputs, Outputs, Sinks, true);
+  EXPECT_TRUE(CE.isEligible(Inputs));
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit1 = getBlockByName(Func, "exit1");
   BasicBlock *Exit2 = getBlockByName(Func, "exit2");
   // Ensure that PHIs in exits are not splitted (since that they have only one
   // incoming value from extracted region).
TEST(CodeExtractor, StoreOutputInvokeResultAfterEHPad) {
   SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "catch"),
     getBlockByName(Func, "invoke.cont2"),
     getBlockByName(Func, "invoke.cont3"),
     getBlockByName(Func, "lpad2")
   };
 
   CodeExtractor CE(ExtractedBlocks);
-  EXPECT_TRUE(CE.isEligible());
+  SetVector<Value *> Inputs, Outputs, Sinks;
+  CE.findInputsOutputs(Inputs, Outputs, Sinks, true);
+  EXPECT_TRUE(CE.isEligible(Inputs));
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined, &errs()));
   EXPECT_FALSE(verifyFunction(*Func, &errs()));
 }
 
TEST(CodeExtractor, StoreOutputInvokeResultInExitStub) {
   )invalid",
                                                 Err, Ctx));
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "entry"),
                                        getBlockByName(Func, "lpad") };
 
   CodeExtractor CE(Blocks);
-  EXPECT_TRUE(CE.isEligible());
+  SetVector<Value *> Inputs, Outputs, Sinks;
+  CE.findInputsOutputs(Inputs, Outputs, Sinks, true);
+  EXPECT_TRUE(CE.isEligible(Inputs));
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
   )ir",
                                                 Err, Ctx));
 
   assert(M && "Could not parse module?");
   Function *Func = M->getFunction("test");
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "if.else") };
   AssumptionCache AC(*Func);
   CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, &AC);
-  EXPECT_TRUE(CE.isEligible());
+  SetVector<Value *> Inputs, Outputs, Sinks;
+  CE.findInputsOutputs(Inputs, Outputs, Sinks, true);
+  EXPECT_TRUE(CE.isEligible(Inputs));
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
   EXPECT_FALSE(CE.verifyAssumptionCache(*Func, &AC));
 }
 } // end anonymous namespace