Split IndirectBr critical edges before PGO gen/use passes.

[hjyamauchi]

Summary

The PGO gen/use passes currently fail with an assert failure if there's a
critical edge whose source is an IndirectBr instruction and that edge
needs to be instrumented.

To avoid this in certain cases, split IndirectBr critical edges in the PGO
gen/use passes. This works for blocks with single indirectbr predecessors,
but not for those with multiple indirectbr predecessors (splitting an
IndirectBr critical edge isn't always possible.)

[efriedma]

The PGO gen/use passes currently fail with an assert failure if there's a critical edge whose source is an IndirectBr instruction.

findIBRPredecessor will fail for blocks with multiple multiple indirectbr predecessors. And then you're back at square one, with an assertion failure.

[davidxl]

1. As Eli pointed out, findIBRpredecessor may return null when there are multiple such predecessors, so perhaps create a new interface to return all such predecessors?

2. it is probably better to split out the refactoring part of the patch.

lib/Passes/PassBuilder.cpp

519

this is common to Use and Gen, so it is possible to move it into the previous SCC pipeline after InstCombinePass.

[hjyamauchi]

In D40699#941591, @efriedma wrote:

The PGO gen/use passes currently fail with an assert failure if there's a critical edge whose source is an IndirectBr instruction.

findIBRPredecessor will fail for blocks with multiple multiple indirectbr predecessors. And then you're back at square one, with an assertion failure.

True. This change (and SplitIndirectBrCriticalEdges) isn't intended to handle all possible cases (see the comment in BasicBlockUtils.h). It's an improvement in that it can handle single indirectbr predecessor cases.

[hjyamauchi]

lib/Passes/PassBuilder.cpp

519

Do you mean moving this to after line 498 above? If so, would it not work if isOptimizingForSize() is true?

[hjyamauchi]

Created https://reviews.llvm.org/D40750 for the refactoring part.

[hjyamauchi]

Rebased after D40750.

Changed the SplitIndirectBrCriticalEdges pass to a function pass.

Moved the SplitIndirectBrCriticalEdges pass into the previous SCC pipeline.

[efriedma]

It's an improvement in that it can handle single indirectbr predecessor cases.

It's not good to kick the can down the road like this. If an IR optimization pass is crashing, you should fix it by fixing that pass, not by changing the default pass pipeline to avoid generating IR which triggers the crash.

[davidxl]

This can be done in the PGO pass itself before the MST is formed. BFI has update interface such as setBlockFreq. The update can use simple heuristic to distribute the frequency.

[hjyamauchi]

Split indirectbr critical edges in the PGO gen/use passes.

[hjyamauchi]

Split only the IndirectBr critical edges that need to be split rather than
all. This avoids changing the indirectbr.ll test.

Minor cleanup.

[davidxl]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

567

Why not splitting the edges before MST is computed?

[hjyamauchi]

Detect indirectbr critical non-MST edges in CFGMST and if none, avoid iterating
over the edges in FuncPGOInstrumentation().

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

567

This is trying to not split unless it's necessary (that is, if it's a non-MST indirectbr critical edges.) For that, we need to compute the MST once before making a decision to split. We could split all indirectbr critical edges unconditionally before computing the MST, but I thought it may be somewhat nicer that way as it would minimize CFG changes (eg, the indirectbr.ll test previously had to be changed even though it didn't suffer from this assert failure.)

[davidxl]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

572

It seems cleaner to move the IndirectBrTargets collection inside MST's buildEdges method.

576

There seems no guarantee that the non-split indirectBR critical edges will be be in MST when cfgmst is recomputed.

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

576

Good point. I'll go back to splitting all indirectBR critical edges before computing the MST (the previous diff).

[hjyamauchi]

Going back to Diff 125662.

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

572

This comment doesn't apply any more.

[davidxl]

lgtm

[hjyamauchi]

Rebased.

[efriedma]

You still never answered the question of what happens when there's multiple indirectbrs whose edges can't be split.

[hjyamauchi]

I think it will be an assert failure like before. Like said above, this patch isn't intended to handle that case.

include/llvm/InitializePasses.h

 void initializeSimpleInlinerPass(PassRegistry&);
 void initializeSimpleLoopUnswitchLegacyPassPass(PassRegistry&);
 void initializeSingleLoopExtractorPass(PassRegistry&);
 void initializeSinkingLegacyPassPass(PassRegistry&);
 void initializeSjLjEHPreparePass(PassRegistry&);
 void initializeSlotIndexesPass(PassRegistry&);
 void initializeSpeculativeExecutionLegacyPassPass(PassRegistry&);
 void initializeSpillPlacementPass(PassRegistry&);
+void initializeSplitIndirectBrCriticalEdgesLegacyPassPass(PassRegistry&);
 void initializeStackColoringPass(PassRegistry&);
 void initializeStackMapLivenessPass(PassRegistry&);
 void initializeStackProtectorPass(PassRegistry&);
 void initializeStackSlotColoringPass(PassRegistry&);
 void initializeStraightLineStrengthReducePass(PassRegistry&);
 void initializeStripDeadDebugInfoPass(PassRegistry&);
 void initializeStripDeadPrototypesLegacyPassPass(PassRegistry&);
 void initializeStripDebugDeclarePass(PassRegistry&);

include/llvm/LinkAllPasses.h

ForcePassLinking() {
       (void) llvm::createDomPrinterPass();
       (void) llvm::createDomOnlyViewerPass();
       (void) llvm::createDomViewerPass();
       (void) llvm::createGCOVProfilerPass();
       (void) llvm::createPGOInstrumentationGenLegacyPass();
       (void) llvm::createPGOInstrumentationUseLegacyPass();
       (void) llvm::createPGOIndirectCallPromotionLegacyPass();
       (void) llvm::createPGOMemOPSizeOptLegacyPass();
+      (void) llvm::createSplitIndirectBrCriticalEdgesLegacyPass();
       (void) llvm::createInstrProfilingLegacyPass();
       (void) llvm::createFunctionImportPass();
       (void) llvm::createFunctionInliningPass();
       (void) llvm::createAlwaysInlinerLegacyPass();
       (void) llvm::createGlobalDCEPass();
       (void) llvm::createGlobalOptimizerPass();
       (void) llvm::createGlobalsAAWrapperPass();
       (void) llvm::createGuardWideningPass();

include/llvm/Transforms/Instrumentation.h

 // PGO Instrumention
 ModulePass *createPGOInstrumentationGenLegacyPass();
 ModulePass *
 createPGOInstrumentationUseLegacyPass(StringRef Filename = StringRef(""));
 ModulePass *createPGOIndirectCallPromotionLegacyPass(bool InLTO = false,
                                                      bool SamplePGO = false);
 FunctionPass *createPGOMemOPSizeOptLegacyPass();
 
+FunctionPass *createSplitIndirectBrCriticalEdgesLegacyPass();
+
 // Helper function to check if it is legal to promote indirect call \p Inst
 // to a direct call of function \p F. Stores the reason in \p Reason.
 bool isLegalToPromote(Instruction *Inst, Function *F, const char **Reason);
 
 // Helper function that transforms Inst (either an indirect-call instruction, or
 // an invoke instruction , to a conditional call to F. This is like:
 //     if (Inst.CalledValue == F)
 //        F(...);

include/llvm/Transforms/PGOInstrumentation.h

public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 void setProfMetadata(Module *M, Instruction *TI, ArrayRef<uint64_t> EdgeCounts,
                      uint64_t MaxCount);
 
 void setIrrLoopHeaderMetadata(Module *M, Instruction *TI, uint64_t Count);
 
+/// The split-indirectbr-critical-edges pass.
+class SplitIndirectBrCriticalEdgesPass :
+      public PassInfoMixin<SplitIndirectBrCriticalEdgesPass> {
+public:
+  SplitIndirectBrCriticalEdgesPass() {}
+
+  static StringRef name() { return "SplitIndirectBrCriticalEdgesPass"; }
+
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_PGOINSTRUMENTATION_H

lib/Passes/PassBuilder.cpp

if (!isOptimizingForSize(Level)) {
 
     CGPipeline.addPass(InlinerPass(IP));
 
     FunctionPassManager FPM;
     FPM.addPass(SROA());
     FPM.addPass(EarlyCSEPass());    // Catch trivial redundancies.
     FPM.addPass(SimplifyCFGPass()); // Merge & remove basic blocks.
     FPM.addPass(InstCombinePass()); // Combine silly sequences.
+    // Split the indirectbr critical edges before the PGO instrumentation gen
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    FPM.addPass(SplitIndirectBrCriticalEdgesPass());
     invokePeepholeEPCallbacks(FPM, Level);
 
     CGPipeline.addPass(createCGSCCToFunctionPassAdaptor(std::move(FPM)));
 
     MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPipeline)));
   }
 
   // Delete anything that is now dead to make sure that we don't instrument
   // dead code. Instrumentation can end up keeping dead code around and
   // dramatically increase code size.
   MPM.addPass(GlobalDCEPass());
 
   if (RunProfileGen) {
     MPM.addPass(PGOInstrumentationGen());
 
     FunctionPassManager FPM;
     FPM.addPass(createFunctionToLoopPassAdaptor(LoopRotatePass()));
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
 
     // Add the profile lowering pass.
     InstrProfOptions Options;
     if (!ProfileGenFile.empty())
       Options.InstrProfileOutput = ProfileGenFile;
     Options.DoCounterPromotion = true;
     MPM.addPass(InstrProfiling(Options));
   }
 
-  if (!ProfileUseFile.empty())
+  if (!ProfileUseFile.empty()) {
     MPM.addPass(PGOInstrumentationUse(ProfileUseFile));
   }
+}
 
 static InlineParams
 getInlineParamsFromOptLevel(PassBuilder::OptimizationLevel Level) {
   auto O3 = PassBuilder::O3;
   unsigned OptLevel = Level > O3 ? 2 : Level;
   unsigned SizeLevel = Level > O3 ? Level - O3 : 0;
   return getInlineParams(OptLevel, SizeLevel);
 }

lib/Passes/PassRegistry.def

 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(dbgs()))
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("sccp", SCCPPass())
 FUNCTION_PASS("simplify-cfg", SimplifyCFGPass())
 FUNCTION_PASS("sink", SinkingPass())
 FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("speculative-execution", SpeculativeExecutionPass())
 FUNCTION_PASS("spec-phis", SpeculateAroundPHIsPass())
+FUNCTION_PASS("split-indirectbr-critical-edges",
+              SplitIndirectBrCriticalEdgesPass())
 FUNCTION_PASS("sroa", SROA())
 FUNCTION_PASS("tailcallelim", TailCallElimPass())
 FUNCTION_PASS("unreachableblockelim", UnreachableBlockElimPass())
 FUNCTION_PASS("unroll", LoopUnrollPass())
 FUNCTION_PASS("verify", VerifierPass())
 FUNCTION_PASS("verify<domtree>", DominatorTreeVerifierPass())
 FUNCTION_PASS("verify<loops>", LoopVerifierPass())
 FUNCTION_PASS("verify<memoryssa>", MemorySSAVerifierPass())

lib/Transforms/IPO/PassManagerBuilder.cpp

if (OptLevel > 0 && SizeLevel == 0 && !DisablePreInliner &&
     // This should probably be lowered after performance testing.
     IP.HintThreshold = 325;
 
     MPM.add(createFunctionInliningPass(IP));
     MPM.add(createSROAPass());
     MPM.add(createEarlyCSEPass());             // Catch trivial redundancies
     MPM.add(createCFGSimplificationPass());    // Merge & remove BBs
     MPM.add(createInstructionCombiningPass()); // Combine silly seq's
+    // Split the indirectbr critical edges before the PGO instrumentation gen
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    MPM.add(createSplitIndirectBrCriticalEdgesLegacyPass());
     addExtensionsToPM(EP_Peephole, MPM);
   }
   if (EnablePGOInstrGen) {
     MPM.add(createPGOInstrumentationGenLegacyPass());
     // Add the profile lowering pass.
     InstrProfOptions Options;
     if (!PGOInstrGen.empty())
       Options.InstrProfileOutput = PGOInstrGen;

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

private:
 
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<BlockFrequencyInfoWrapperPass>();
   }
 };
 
+class SplitIndirectBrCriticalEdgesLegacyPass : public FunctionPass {
+public:
+  static char ID;
+
+  SplitIndirectBrCriticalEdgesLegacyPass() : FunctionPass(ID) {
+  }
+
+  StringRef getPassName() const override {
+    return "SplitIndirectBrCriticalEdgesPass";
+  }
+
+private:
+  bool runOnFunction(Function &F) override;
+};
+
 } // end anonymous namespace
 
 char PGOInstrumentationGenLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(PGOInstrumentationGenLegacyPass, "pgo-instr-gen",
                       "PGO instrumentation.", false, false)
 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)

 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                     "Read PGO instrumentation profile.", false, false)
 
 ModulePass *llvm::createPGOInstrumentationUseLegacyPass(StringRef Filename) {
   return new PGOInstrumentationUseLegacyPass(Filename.str());
 }
 
+char SplitIndirectBrCriticalEdgesLegacyPass::ID = 0;
+
+INITIALIZE_PASS_BEGIN(SplitIndirectBrCriticalEdgesLegacyPass,
+                      "split-indirectbr-critical-edges",
+                      "Split IndirectBr critical edges.", false, false)
+INITIALIZE_PASS_END(SplitIndirectBrCriticalEdgesLegacyPass,
+                    "split-indirectbr-critical-edges",
+                    "Split IndirectBr critical edges.", false, false)
+
+FunctionPass *llvm::createSplitIndirectBrCriticalEdgesLegacyPass() {
+  return new SplitIndirectBrCriticalEdgesLegacyPass();
+}
+
 namespace {
 
 /// \brief An MST based instrumentation for PGO
 ///
 /// Implements a Minimum Spanning Tree (MST) based instrumentation for PGO
 /// in the function level.
 struct PGOEdge {
   // This class implements the CFG edges. Note the CFG can be a multi-graph.
FuncPGOInstrumentation(
       std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
       bool CreateGlobalVar = false, BranchProbabilityInfo *BPI = nullptr,
       BlockFrequencyInfo *BFI = nullptr)
       : F(Func), ComdatMembers(ComdatMembers), ValueSites(IPVK_Last + 1),
         SIVisitor(Func), MIVisitor(Func), MST(F, BPI, BFI) {
     // This should be done before CFG hash computation.
     SIVisitor.countSelects(Func);
     MIVisitor.countMemIntrinsics(Func);
     NumOfPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
     NumOfPGOMemIntrinsics += MIVisitor.getNumOfMemIntrinsics();
     ValueSites[IPVK_IndirectCallTarget] = findIndirectCallSites(Func);
     ValueSites[IPVK_MemOPSize] = MIVisitor.findMemIntrinsics(Func);
 
     FuncName = getPGOFuncName(F);
     computeCFGHash();
     if (!ComdatMembers.empty())
       renameComdatFunction();
     DEBUG(dumpInfo("after CFGMST"));
 
     NumOfPGOBB += MST.BBInfos.size();
     for (auto &E : MST.AllEdges) {
       if (E->Removed)
         continue;
       NumOfPGOEdge++;
       if (!E->InMST)
         NumOfPGOInstrument++;
bool PGOInstrumentationUseLegacyPass::runOnModule(Module &M) {
   };
   auto LookupBFI = [this](Function &F) {
     return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
   };
 
   return annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI);
 }
 
+PreservedAnalyses
+SplitIndirectBrCriticalEdgesPass::run(Function &F,
+                                      FunctionAnalysisManager &AM) {
+  return SplitIndirectBrCriticalEdges(F)
+      ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
+
+bool SplitIndirectBrCriticalEdgesLegacyPass::runOnFunction(Function &F) {
+  if (skipFunction(F))
+    return false;
+  return SplitIndirectBrCriticalEdges(F);
+}
+
 static std::string getSimpleNodeName(const BasicBlock *Node) {
   if (!Node->getName().empty())
     return Node->getName();
 
   std::string SimpleNodeName;
   raw_string_ostream OS(SimpleNodeName);
   Node->printAsOperand(OS, false);
   return OS.str();

test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll

+
+; RUN: opt < %s -passes=split-indirectbr-critical-edges -S | FileCheck %s
+; RUN: opt < %s -passes='function(split-indirectbr-critical-edges),pgo-instr-gen' -S | FileCheck %s
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @bar(i32 %v) local_unnamed_addr #0 {
+entry:
+  %mul = shl nsw i32 %v, 1
+  ret i32 %mul
+}
+
+; Function Attrs: norecurse nounwind readonly uwtable
+define i32 @foo(i8* nocapture readonly %p) #1 {
+entry:
+  %targets = alloca [256 x i8*], align 16
+  %arrayidx1 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 93
+  store i8* blockaddress(@foo, %if.end), i8** %arrayidx1, align 8
+  br label %for.cond2
+
+for.cond2:                                        ; preds = %if.end, %for.cond2, %entry
+; CHECK: for.cond2:                                        ; preds = %.split1
+  %p.addr.0 = phi i8* [ %p, %entry ], [ %incdec.ptr5, %if.end ], [ %incdec.ptr, %for.cond2 ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %p.addr.0, i64 1
+  %0 = load i8, i8* %p.addr.0, align 1
+  %cond = icmp eq i8 %0, 93
+  br i1 %cond, label %if.end.preheader, label %for.cond2
+
+if.end.preheader:                                 ; preds = %for.cond2
+  br label %if.end
+
+if.end:                                           ; preds = %if.end.preheader, %if.end
+; CHECK: if.end:                                           ; preds = %.split1
+  %p.addr.1 = phi i8* [ %incdec.ptr5, %if.end ], [ %incdec.ptr, %if.end.preheader ]
+  %incdec.ptr5 = getelementptr inbounds i8, i8* %p.addr.1, i64 1
+  %1 = load i8, i8* %p.addr.1, align 1
+  %idxprom6 = zext i8 %1 to i64
+  %arrayidx7 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 %idxprom6
+  %2 = load i8*, i8** %arrayidx7, align 8
+  indirectbr i8* %2, [label %for.cond2, label %if.end]
+; CHECK: indirectbr i8* %2, [label %for.cond2, label %if.end]
+}