[CSSPGO] Unblocking optimizations by dangling pseudo probes.

[hoy]

Summary

This change fixes a couple places where the pseudo probe intrinsic blocks optimizations because they are not naturally removable. To unblock those optimizations, the blocking pseudo probes are moved out of the original blocks and tagged dangling, instead of allowing pseudo probes to be literally removed. The reason is that when the original block is removed, we won't be able to sample it. Instead of assigning it a zero weight, moving all its pseudo probes into another block and marking them dangling should allow the counts inference a chance to assign them a more reasonable weight. We have not seen counts quality degradation from our experiments.

The optimizations being unblocked are:

1. Removing conditional probes for if-converted branches. Conditional probes are tagged dangling when their homing branch arms are folded so that they will not be over-counted.
2. Unblocking jump threading from removing empty blocks. Pseudo probe prevents jump threading from removing logically empty blocks that only has one unconditional jump instructions.
3. Unblocking SimplifyCFG and MIR tail duplicate to thread empty blocks and blocks with redundant branch checks.

Since dangling probes are logically deleted, they should not consume any samples in LTO postLink. This can be achieved by setting their distribution factors to zero when dangled.

[hoy]

Removing test code that belongs to another patch.

[wmi]

A general question, do we care about where we should put dangling probe? Currently dangling probes are scattered in some predecessor/successor blocks they don't belong to. Looks like they can be anywhere. Wondering whether it is better to collect all of them to some place (entry block?). I am not sure about it either. Just want to know what you think.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

674–677

In which case SkipPseudoOp needs to be true and which case to be false?

[hoy]

In D97481#2589255, @wmi wrote:

A general question, do we care about where we should put dangling probe? Currently dangling probes are scattered in some predecessor/successor blocks they don't belong to. Looks like they can be anywhere. Wondering whether it is better to collect all of them to some place (entry block?). I am not sure about it either. Just want to know what you think.

That's a very good question. We don't care about where a dangling probe is placed technically. But in reality placing dangling probes around the original place they are from helps IR readability. It can help reason about why they are dangled.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

674–677

Good question. SkipPseudoOp should be true when we are sure that the probes should not be ignored. For example, we don't want to merge blocks that look the same except for their pseudo probes because the merge will make it hard to infer the counts for the original edges. On the contrary, SkipPseudoOp should be false when we are sure probes are movable or removable.

[wmi]

A general question, do we care about where we should put dangling probe? Currently dangling probes are scattered in some predecessor/successor blocks they don't belong to. Looks like they can be anywhere. Wondering whether it is better to collect all of them to some place (entry block?). I am not sure about it either. Just want to know what you think.

That's a very good question. We don't care about where a dangling probe is placed technically. But in reality placing dangling probes around the original place they are from helps IR readability. It can help reason about why they are dangled.

Ok, that is reasonable.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

674–677

Feel a little hesitant to add a sampleFDO specific param to a common used interface because people may not understand when the param should be true or when it should be false. Even using the default version sometimes can block optimizations unintentionally in sampleFDO mode, like what the change tries to fix for taildup or branchfolding.

pseudo hook itself has some ambiguity here. When compiler is trying to merge blocks with different pseudo hook, they are essential instructions which could affect the profile result. When compiler is trying to remove or thread empty block, pseudo hook is non essential. Could you give some actionable guidance as comments about when SkipPseudoOp should be true?

[wmi]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

674–677

Think about it a little more. If tail merge (merge blocks that look the same except for their pseudo probes) is causing trouble for pseudo hook based profile accuracy, is it possible to just directly turn off tail merge if pseudo hook is in use? In that way, we eliminate the ambiguity of pseudo hook and we can always treat it the same way as debug instructions in getFirstNonDebugInstr.

[hoy]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

674–677

That's a valid concern. Yeah, the semantics of pseudo probes as far as whether optimizations should be blocked is ambiguous and requires user to be judicious. We are trying to make a trade off between profile quality and code quality. I went through all the uses of these APIs and changed some of them to skip probes for better quality while not damaging profile quality. I'll put some comments based on my experience.

Regarding tail merge, sometimes it is useful to blocks with exactly the same code including probes. A pass typically does not handle a single pattern like block merge or block duplicate. For example, both the merge and duplicate is done by the jump threading pass. Disabling a part of a pass (especially a future pass) may not be done automatically and still requires user's judgement.

I agree that letting optimization developer be aware of sample profile quality may not be practical. So far the approach I'm taking is to keep a good profile quality while fixing code quality regression. I expect at some time we can flip over the default value of SkipPseudoOp for maximum code quality and we keep an eye on the profile quality. We can also make a knob that automatically changes the default value to favor either factor.

[hoy]

Adding comments for the usage of SkipPseudoOp.

[wmi]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

676–678

This should work most of the time when ..., except for certain cases ...

It sounds like "SkipPseudoOp = true" is the more common cases and only certain cases will use the default value. It will lead to the confusion whether user should use the default value or not for now. I know you mentioned you would flip the default value in the future. If you think user is better to choose "SkipPseudoOp = true" in common cases, it may be good to suggest that in the front and add a TODO for flipping the default value.

llvm/lib/Transforms/IPO/SampleProfile.cpp

538–539

Nit: should not consume

[hoy]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

676–678

Sure, comment added. I'd like to turn it on by default after a few more rounds until all optimizations that could hurt profile quality are understood.

llvm/lib/Transforms/IPO/SampleProfile.cpp

538–539

Fixed.

[hoy]

Addressing Wei's feedbacks.

[wmi]

LGTM.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

#endif
   /// copies all the successors FromMBB and remove all the successors from
   /// FromMBB).
   void transferSuccessors(MachineBasicBlock *FromMBB);
 
   /// Transfers all the successors, as in transferSuccessors, and update PHI
   /// operands in the successor blocks which refer to FromMBB to refer to this.
   void transferSuccessorsAndUpdatePHIs(MachineBasicBlock *FromMBB);
 
+  /// move all pseudo probes in this block to the end of /c ToMBB To and tag
+  /// them dangling.
+  void moveAndDanglePseudoProbes(MachineBasicBlock *ToMBB);
+
   /// Return true if any of the successors have probabilities attached to them.
   bool hasSuccessorProbabilities() const { return !Probs.empty(); }
 
   /// Return true if the specified MBB is a predecessor of this block.
   bool isPredecessor(const MachineBasicBlock *MBB) const;
 
   /// Return true if the specified MBB is a successor of this block.
   bool isSuccessor(const MachineBasicBlock *MBB) const;
#endif
   }
 
   /// Same getFirstTerminator but it ignores bundles and return an
   /// instr_iterator instead.
   instr_iterator getFirstInstrTerminator();
 
   /// Returns an iterator to the first non-debug instruction in the basic block,
   /// or end().
-  iterator getFirstNonDebugInstr();
-  const_iterator getFirstNonDebugInstr() const {
-    return const_cast<MachineBasicBlock *>(this)->getFirstNonDebugInstr();
+  iterator getFirstNonDebugInstr(bool SkipPseudoOp = false);
+  const_iterator getFirstNonDebugInstr(bool SkipPseudoOp = false) const {
+    return const_cast<MachineBasicBlock *>(this)->getFirstNonDebugInstr(
+        SkipPseudoOp);
   }
 
   /// Returns an iterator to the last non-debug instruction in the basic block,
   /// or end().
-  iterator getLastNonDebugInstr();
-  const_iterator getLastNonDebugInstr() const {
-    return const_cast<MachineBasicBlock *>(this)->getLastNonDebugInstr();
+  iterator getLastNonDebugInstr(bool SkipPseudoOp = false);
+  const_iterator getLastNonDebugInstr(bool SkipPseudoOp = false) const {
+    return const_cast<MachineBasicBlock *>(this)->getLastNonDebugInstr(
+        SkipPseudoOp);
   }
 
   /// Convenience function that returns true if the block ends in a return
   /// instruction.
   bool isReturnBlock() const {
     return !empty() && back().isReturn();
   }
 
public:
 
   MachineBasicBlock::iterator getInitial() { return I; }
 };
 
 /// Increment \p It until it points to a non-debug instruction or to \p End
 /// and return the resulting iterator. This function should only be used
 /// MachineBasicBlock::{iterator, const_iterator, instr_iterator,
 /// const_instr_iterator} and the respective reverse iterators.
-template<typename IterT>
-inline IterT skipDebugInstructionsForward(IterT It, IterT End) {
-  while (It != End && It->isDebugInstr())
+template <typename IterT>
+inline IterT skipDebugInstructionsForward(IterT It, IterT End,
+                                          bool SkipPseudoOp = false) {
+  while (It != End &&
+         (It->isDebugInstr() || (SkipPseudoOp && It->isPseudoProbe())))
     ++It;
   return It;
 }
 
 /// Decrement \p It until it points to a non-debug instruction or to \p Begin
 /// and return the resulting iterator. This function should only be used
 /// MachineBasicBlock::{iterator, const_iterator, instr_iterator,
 /// const_instr_iterator} and the respective reverse iterators.
-template<class IterT>
-inline IterT skipDebugInstructionsBackward(IterT It, IterT Begin) {
-  while (It != Begin && It->isDebugInstr())
+template <class IterT>
+inline IterT skipDebugInstructionsBackward(IterT It, IterT Begin,
+                                           bool SkipPseudoOp = false) {
+  while (It != Begin &&
+         (It->isDebugInstr() || (SkipPseudoOp && It->isPseudoProbe())))
     --It;
   return It;
 }
 
 /// Increment \p It, then continue incrementing it while it points to a debug
 /// instruction. A replacement for std::next.
-template <typename IterT> inline IterT next_nodbg(IterT It, IterT End) {
-  return skipDebugInstructionsForward(std::next(It), End);
+template <typename IterT>
+inline IterT next_nodbg(IterT It, IterT End, bool SkipPseudoOp = false) {
+  return skipDebugInstructionsForward(std::next(It), End, SkipPseudoOp);
 }
 
 /// Decrement \p It, then continue decrementing it while it points to a debug
 /// instruction. A replacement for std::prev.
-template <typename IterT> inline IterT prev_nodbg(IterT It, IterT Begin) {
-  return skipDebugInstructionsBackward(std::prev(It), Begin);
+template <typename IterT>
+inline IterT prev_nodbg(IterT It, IterT Begin, bool SkipPseudoOp = false) {
+  return skipDebugInstructionsBackward(std::prev(It), Begin, SkipPseudoOp);
 }
 
 /// Construct a range iterator which begins at \p It and moves forwards until
 /// \p End is reached, skipping any debug instructions.
 template <typename IterT>
-inline auto instructionsWithoutDebug(IterT It, IterT End) {
-  return make_filter_range(make_range(It, End), [](const MachineInstr &MI) {
-    return !MI.isDebugInstr();
+inline auto instructionsWithoutDebug(IterT It, IterT End,
+                                     bool SkipPseudoOp = false) {
+  return make_filter_range(make_range(It, End), [=](const MachineInstr &MI) {
+    return !MI.isDebugInstr() && !(SkipPseudoOp && MI.isPseudoProbe());
   });
 }
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MACHINEBASICBLOCK_H

llvm/include/llvm/IR/PseudoProbe.h

public:
   constexpr static uint8_t FullDistributionFactor = 100;
 };
 
 struct PseudoProbe {
   uint32_t Id;
   uint32_t Type;
   uint32_t Attr;
   float Factor;
+
+  bool isDangling() const {
+    return Attr & (uint32_t)PseudoProbeAttributes::Dangling;
+  }
 };
 
 Optional<PseudoProbe> extractProbe(const Instruction &Inst);
 
 void setProbeDistributionFactor(Instruction &Inst, float Factor);
 
+bool moveAndDanglePseudoProbes(BasicBlock *From, Instruction *To);
 } // end namespace llvm
 
 #endif // LLVM_IR_PSEUDOPROBE_H

llvm/lib/CodeGen/BranchFolding.cpp

bool BranchFolder::OptimizeBranches(MachineFunction &MF) {
   }
 
   return MadeChange;
 }
 
 // Blocks should be considered empty if they contain only debug info;
 // else the debug info would affect codegen.
 static bool IsEmptyBlock(MachineBasicBlock *MBB) {
-  return MBB->getFirstNonDebugInstr() == MBB->end();
+  return MBB->getFirstNonDebugInstr(true) == MBB->end();
 }
 
 // Blocks with only debug info and branches should be considered the same
 // as blocks with only branches.
 static bool IsBranchOnlyBlock(MachineBasicBlock *MBB) {
   MachineBasicBlock::iterator I = MBB->getFirstNonDebugInstr();
   assert(I != MBB->end() && "empty block!");
   return I->isBranch();
for (MachineBasicBlock *SuccBB : MBB.successors())
     if (SuccBB->pred_size() == 1)
       copyDebugInfoToSuccessor(TII, MBB, *SuccBB);
   // If this MBB is the only successor of a predecessor it is legal to copy the
   // DBG_VALUE instructions to the end of the predecessor (just before the
   // terminators, assuming that the terminator isn't affecting the DBG_VALUE).
   for (MachineBasicBlock *PredBB : MBB.predecessors())
     if (PredBB->succ_size() == 1)
       copyDebugInfoToPredecessor(TII, MBB, *PredBB);
+
+  // For AutoFDO, if the block is removed, we won't be able to sample it. To
+  // avoid assigning a zero weight for BB, move all its pseudo probes into once
+  // of its predecessors or successors and mark them dangling. This should allow
+  // the counts inference a chance to get a more reasonable weight for the
+  // block.
+  if (!MBB.pred_empty())
+    MBB.moveAndDanglePseudoProbes(*MBB.pred_begin());
+  else if (!MBB.succ_empty())
+    MBB.moveAndDanglePseudoProbes(*MBB.succ_begin());
 }
 
 bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
   bool MadeChange = false;
   MachineFunction &MF = *MBB->getParent();
 ReoptimizeBlock:
 
   MachineFunction::iterator FallThrough = MBB->getIterator();

llvm/lib/CodeGen/MachineBasicBlock.cpp

MachineBasicBlock::instr_iterator MachineBasicBlock::getFirstInstrTerminator() {
   instr_iterator B = instr_begin(), E = instr_end(), I = E;
   while (I != B && ((--I)->isTerminator() || I->isDebugInstr()))
     ; /*noop */
   while (I != E && !I->isTerminator())
     ++I;
   return I;
 }
 
-MachineBasicBlock::iterator MachineBasicBlock::getFirstNonDebugInstr() {
+MachineBasicBlock::iterator
+MachineBasicBlock::getFirstNonDebugInstr(bool SkipPseudoOp) {
   // Skip over begin-of-block dbg_value instructions.
-  return skipDebugInstructionsForward(begin(), end());
+  return skipDebugInstructionsForward(begin(), end(), SkipPseudoOp);
 }
 
-MachineBasicBlock::iterator MachineBasicBlock::getLastNonDebugInstr() {
+MachineBasicBlock::iterator
+MachineBasicBlock::getLastNonDebugInstr(bool SkipPseudoOp) {
   // Skip over end-of-block dbg_value instructions.
   instr_iterator B = instr_begin(), I = instr_end();
   while (I != B) {
     --I;
     // Return instruction that starts a bundle.
     if (I->isDebugInstr() || I->isInsideBundle())
       continue;
+    if (SkipPseudoOp && I->isPseudoProbe())
+      continue;
     return I;
   }
   // The block is all debug values.
   return end();
 }
 
 bool MachineBasicBlock::hasEHPadSuccessor() const {
   for (const_succ_iterator I = succ_begin(), E = succ_end(); I != E; ++I)
while (!FromMBB->succ_empty()) {
     FromMBB->removeSuccessor(Succ);
 
     // Fix up any PHI nodes in the successor.
     Succ->replacePhiUsesWith(FromMBB, this);
   }
   normalizeSuccProbs();
 }
 
+/// A block emptied (i.e., with all instructions moved out of it) won't be
+/// sampled at run time. In such cases, AutoFDO will be informed of zero samples
+/// collected for the block. This is not accurate and could lead to misleading
+/// weights assigned for the block. A way to mitigate that is to treat such
+/// block as having unknown counts in the AutoFDO profile loader and allow the
+/// counts inference tool a chance to calculate a relatively reasonable weight
+/// for it. This can be done by moving all pseudo probes in the emptied block
+/// i.e, /c this, to before /c ToMBB and tag them dangling. Note that this is
+/// not needed for dead blocks which really have a zero weight. It's per
+/// transforms to decide whether to call this function or not.
+void MachineBasicBlock::moveAndDanglePseudoProbes(MachineBasicBlock *ToMBB) {
+  SmallVector<MachineInstr *, 4> ToBeMoved;
+  for (MachineInstr &MI : instrs()) {
+    if (MI.isPseudoProbe()) {
+      MI.addPseudoProbeAttribute(PseudoProbeAttributes::Dangling);
+      ToBeMoved.push_back(&MI);
+    }
+  }
+
+  MachineBasicBlock::iterator I = ToMBB->getFirstTerminator();
+  for (MachineInstr *MI : ToBeMoved) {
+    MI->removeFromParent();
+    ToMBB->insert(I, MI);
+  }
+}
+
 bool MachineBasicBlock::isPredecessor(const MachineBasicBlock *MBB) const {
   return is_contained(predecessors(), MBB);
 }
 
 bool MachineBasicBlock::isSuccessor(const MachineBasicBlock *MBB) const {
   return is_contained(successors(), MBB);
 }
 

llvm/lib/CodeGen/TailDuplicator.cpp

 }
 
 /// True if this BB has only one unconditional jump.
 bool TailDuplicator::isSimpleBB(MachineBasicBlock *TailBB) {
   if (TailBB->succ_size() != 1)
     return false;
   if (TailBB->pred_empty())
     return false;
-  MachineBasicBlock::iterator I = TailBB->getFirstNonDebugInstr();
+  MachineBasicBlock::iterator I = TailBB->getFirstNonDebugInstr(true);
   if (I == TailBB->end())
     return true;
   return I->isUnconditionalBranch();
 }
 
 static bool bothUsedInPHI(const MachineBasicBlock &A,
                           const SmallPtrSet<MachineBasicBlock *, 8> &SuccsB) {
   for (MachineBasicBlock *BB : A.successors())
for (MachineBasicBlock *PredBB : Preds) {
 
     if (!PredBB->isSuccessor(NewTarget))
       PredBB->replaceSuccessor(TailBB, NewTarget);
     else {
       PredBB->removeSuccessor(TailBB, true);
       assert(PredBB->succ_size() <= 1);
     }
 
+    // For AutoFDO, since BB is going to be removed, we won't be able to sample
+    // it. To avoid assigning a zero weight for BB, move all its pseudo probes
+    // into Succ and mark them dangling. This should allow the counts inference
+    // a chance to get a more reasonable weight for BB.
+    TailBB->moveAndDanglePseudoProbes(PredBB);
+
     if (PredTBB)
       TII->insertBranch(*PredBB, PredTBB, PredFBB, PredCond, DL);
 
     TDBBs.push_back(PredBB);
   }
   return Changed;
 }
 

llvm/lib/IR/PseudoProbe.cpp

if (const DebugLoc &DLoc = Inst.getDebugLoc()) {
         uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
             Index, Type, Attr, IntFactor);
         DIL = DIL->cloneWithDiscriminator(V);
         Inst.setDebugLoc(DIL);
       }
     }
   }
 }
+
+void addPseudoProbeAttribute(PseudoProbeInst &Inst,
+                             PseudoProbeAttributes Attr) {
+  IRBuilder<> Builder(&Inst);
+  uint32_t OldAttr = Inst.getAttributes()->getZExtValue();
+  uint32_t NewAttr = OldAttr | (uint32_t)Attr;
+  if (OldAttr != NewAttr)
+    Inst.replaceUsesOfWith(Inst.getAttributes(), Builder.getInt32(NewAttr));
+}
+
+/// A block emptied (i.e., with all instructions moved out of it) won't be
+/// sampled at run time. In such cases, AutoFDO will be informed of zero samples
+/// collected for the block. This is not accurate and could lead to misleading
+/// weights assigned for the block. A way to mitigate that is to treat such
+/// block as having unknown counts in the AutoFDO profile loader and allow the
+/// counts inference tool a chance to calculate a relatively reasonable weight
+/// for it. This can be done by moving all pseudo probes in the emptied block
+/// i.e, /c From, to before /c To and tag them dangling. Note that this is
+/// not needed for dead blocks which really have a zero weight. It's per
+/// transforms to decide whether to call this function or not.
+bool moveAndDanglePseudoProbes(BasicBlock *From, Instruction *To) {
+  SmallVector<PseudoProbeInst *, 4> ToBeMoved;
+  for (auto &I : *From) {
+    if (auto *II = dyn_cast<PseudoProbeInst>(&I)) {
+      addPseudoProbeAttribute(*II, PseudoProbeAttributes::Dangling);
+      ToBeMoved.push_back(II);
+    }
+  }
+
+  for (auto *I : ToBeMoved)
+    I->moveBefore(To);
+
+  return !ToBeMoved.empty();
+}
 } // namespace llvm

llvm/lib/Transforms/IPO/SampleProfile.cpp

 
 ErrorOr<uint64_t> SampleProfileLoader::getProbeWeight(const Instruction &Inst) {
   assert(FunctionSamples::ProfileIsProbeBased &&
          "Profile is not pseudo probe based");
   Optional<PseudoProbe> Probe = extractProbe(Inst);
   if (!Probe)
     return std::error_code();
 
+  // Ignore danling probes since they are logically deleted and should do
+  // not consume any profile samples.
+  if (Probe->isDangling())
+    return std::error_code();
+
   const FunctionSamples *FS = findFunctionSamples(Inst);
   if (!FS)
     return std::error_code();
 
   // If a direct call/invoke instruction is inlined in profile
   // (findCalleeFunctionSamples returns non-empty result), but not inlined here,
   // it means that the inlined callsite has no sample, thus the call
   // instruction should have 0 count.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

return BFI.getBlockProfileCount(BB)
                ? BFI.getBlockProfileCount(BB).getValue()
                : 0;
   };
 
   // Collect the sum of execution weight for each probe.
   ProbeFactorMap ProbeFactors;
   for (auto &Block : F) {
     for (auto &I : Block) {
-      if (Optional<PseudoProbe> Probe = extractProbe(I))
+      if (Optional<PseudoProbe> Probe = extractProbe(I)) {
+        // Do not count dangling probes since they are logically deleted and the
+        // current block that a dangling probe resides in doesn't reflect the
+        // execution count of the probe. The original samples of the probe will
+        // be distributed among the rest probes if there are any, this is
+        // less-than-deal but at least we don't lose any samples.
+        if (!Probe->isDangling())
           ProbeFactors[Probe->Id] += BBProfileCount(&Block);
       }
     }
+  }
 
   // Fix up over-counted probes.
   for (auto &Block : F) {
     for (auto &I : Block) {
       if (Optional<PseudoProbe> Probe = extractProbe(I)) {
+        // Ignore danling probes since they are logically deleted and should do
+        // not consume any profile samples in the subsequent profile annotation.
+        if (!Probe->isDangling()) {
           float Sum = ProbeFactors[Probe->Id];
           if (Sum != 0)
             setProbeDistributionFactor(I, BBProfileCount(&Block) / Sum);
         }
       }
     }
   }
+}
 
 PreservedAnalyses PseudoProbeUpdatePass::run(Module &M,
                                              ModuleAnalysisManager &AM) {
   if (UpdatePseudoProbe) {
     for (auto &F : M) {
       if (F.isDeclaration())
         continue;
       FunctionAnalysisManager &FAM =
           AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
       runOnFunction(F, FAM);
     }
   }
   return PreservedAnalyses::none();
 }

llvm/lib/Transforms/Scalar/JumpThreading.cpp

for (auto &BB : F) {
 
       // processBlock doesn't thread BBs with unconditional TIs. However, if BB
       // is "almost empty", we attempt to merge BB with its sole successor.
       auto *BI = dyn_cast<BranchInst>(BB.getTerminator());
       if (BI && BI->isUnconditional()) {
         BasicBlock *Succ = BI->getSuccessor(0);
         if (
             // The terminator must be the only non-phi instruction in BB.
-            BB.getFirstNonPHIOrDbg()->isTerminator() &&
+            BB.getFirstNonPHIOrDbg(true)->isTerminator() &&
             // Don't alter Loop headers and latches to ensure another pass can
             // detect and transform nested loops later.
             !LoopHeaders.count(&BB) && !LoopHeaders.count(Succ) &&
             TryToSimplifyUncondBranchFromEmptyBlock(&BB, DTU)) {
           RemoveRedundantDbgInstrs(Succ);
           // BB is valid for cleanup here because we passed in DTU. F remains
           // BB's parent until a DTU->getDomTree() event.
           LVI->eraseBlock(&BB);

llvm/lib/Transforms/Utils/Local.cpp

 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
   // add the metadata to the branch instructions in the predecessors.
   unsigned LoopMDKind = BB->getContext().getMDKindID("llvm.loop");
   Instruction *TI = BB->getTerminator();
   if (TI)
     if (MDNode *LoopMD = TI->getMetadata(LoopMDKind))
       for (BasicBlock *Pred : predecessors(BB))
         Pred->getTerminator()->setMetadata(LoopMDKind, LoopMD);
 
+  // For AutoFDO, since BB is going to be removed, we won't be able to sample
+  // it. To avoid assigning a zero weight for BB, move all its pseudo probes
+  // into Succ and mark them dangling. This should allow the counts inference a
+  // chance to get a more reasonable weight for BB.
+  moveAndDanglePseudoProbes(BB, &*Succ->getFirstInsertionPt());
+
   // Everything that jumped to BB now goes to Succ.
   BB->replaceAllUsesWith(Succ);
   if (!Succ->hasName()) Succ->takeName(BB);
 
   // Clear the successor list of BB to match updates applying to DTU later.
   if (BB->getTerminator())
     BB->getInstList().pop_back();
   new UnreachableInst(BB->getContext(), BB);
void llvm::hoistAllInstructionsInto(BasicBlock *DomBlock, Instruction *InsertPt,
   // DILocation in either branch left after performing the transformation. We
   // can only insert a dbg.value after the two branches are joined again.
   //
   // See PR38762, PR39243 for more details.
   //
   // TODO: Extend llvm.dbg.value to take more than one SSA Value (PR39141) to
   // encode predicated DIExpressions that yield different results on different
   // code paths.
+
+  // A hoisted conditional probe should be treated as dangling so that it will
+  // not be over-counted when the samples collected on the non-conditional path
+  // are counted towards the conditional path. We leave it for the counts
+  // inference algorithm to figure out a proper count for a danglng probe.
+  moveAndDanglePseudoProbes(BB, InsertPt);
+
   for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
     Instruction *I = &*II;
     I->dropUnknownNonDebugMetadata();
     if (I->isUsedByMetadata())
       dropDebugUsers(*I);
     if (isa<DbgInfoIntrinsic>(I)) {
       // Remove DbgInfo Intrinsics.
       II = I->eraseFromParent();

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
if (isa<DbgInfoIntrinsic>(I)) {
       SpeculatedDbgIntrinsics.push_back(I);
       continue;
     }
 
     // Skip pseudo probes. The consequence is we lose track of the branch
     // probability for ThenBB, which is fine since the optimization here takes
     // place regardless of the branch probability.
     if (isa<PseudoProbeInst>(I)) {
-      SpeculatedDbgIntrinsics.push_back(I);
       continue;
     }
 
     // Only speculatively execute a single instruction (not counting the
     // terminator) for now.
     ++SpeculatedInstructions;
     if (SpeculatedInstructions > 1)
       return false;
bool SimplifyCFGOpt::SpeculativelyExecuteBB(BranchInst *BI, BasicBlock *ThenBB,
   // to avoid making it appear as if the condition is a constant, which would
   // be misleading while debugging.
   for (auto &I : *ThenBB) {
     if (!SpeculatedStoreValue || &I != SpeculatedStore)
       I.setDebugLoc(DebugLoc());
     I.dropUnknownNonDebugMetadata();
   }
 
+  // A hoisted conditional probe should be treated as dangling so that it will
+  // not be over-counted when the samples collected on the non-conditional path
+  // are counted towards the conditional path. We leave it for the counts
+  // inference algorithm to figure out a proper count for a danglng probe.
+  moveAndDanglePseudoProbes(ThenBB, BI);
+
   // Hoist the instructions.
   BB->getInstList().splice(BI->getIterator(), ThenBB->getInstList(),
                            ThenBB->begin(), std::prev(ThenBB->end()));
 
   // Insert selects and rewrite the PHI operands.
   IRBuilder<NoFolder> Builder(BI);
   for (PHINode &PN : EndBB->phis()) {
     unsigned OrigI = PN.getBasicBlockIndex(BB);
bool SimplifyCFGOpt::simplifyUncondBranch(BranchInst *BI,
   // vectorization to keep canonical loop forms for nested loops. These blocks
   // can be eliminated when the pass is invoked later in the back-end.)
   // Note that if BB has only one predecessor then we do not introduce new
   // backedge, so we can eliminate BB.
   bool NeedCanonicalLoop =
       Options.NeedCanonicalLoop &&
       (!LoopHeaders.empty() && BB->hasNPredecessorsOrMore(2) &&
        (is_contained(LoopHeaders, BB) || is_contained(LoopHeaders, Succ)));
-  BasicBlock::iterator I = BB->getFirstNonPHIOrDbg()->getIterator();
+  BasicBlock::iterator I = BB->getFirstNonPHIOrDbg(true)->getIterator();
   if (I->isTerminator() && BB != &BB->getParent()->getEntryBlock() &&
       !NeedCanonicalLoop && TryToSimplifyUncondBranchFromEmptyBlock(BB, DTU))
     return true;
 
   // If the only instruction in the block is a seteq/setne comparison against a
   // constant, try to simplify the block.
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(I))
     if (ICI->isEquality() && isa<ConstantInt>(ICI->getOperand(1))) {
if (isValueEqualityComparison(BI)) {
     // If we only have one predecessor, and if it is a branch on this value,
     // see if that predecessor totally determines the outcome of this
     // switch.
     if (BasicBlock *OnlyPred = BB->getSinglePredecessor())
       if (SimplifyEqualityComparisonWithOnlyPredecessor(BI, OnlyPred, Builder))
         return requestResimplify();
 
     // This block must be empty, except for the setcond inst, if it exists.
-    // Ignore dbg intrinsics.
-    auto I = BB->instructionsWithoutDebug().begin();
+    // Ignore dbg and pseudo intrinsics.
+    auto I = BB->instructionsWithoutDebug(true).begin();
     if (&*I == BI) {
       if (FoldValueComparisonIntoPredecessors(BI, Builder))
         return requestResimplify();
     } else if (&*I == cast<Instruction>(BI->getCondition())) {
       ++I;
       if (&*I == BI && FoldValueComparisonIntoPredecessors(BI, Builder))
         return requestResimplify();
     }

llvm/test/Transforms/SampleProfile/pseudo-probe-dangle.ll

+; REQUIRES: x86_64-linux
+; RUN: opt < %s -passes='pseudo-probe,jump-threading' -S -o %t
+; RUN: FileCheck %s < %t --check-prefix=JT
+; RUN: llc -pseudo-probe-for-profiling -function-sections <%t -filetype=asm | FileCheck %s --check-prefix=ASM
+; RUN: opt < %s -passes='pseudo-probe' -S -o %t1
+; RUN: llc -pseudo-probe-for-profiling -stop-after=tailduplication <%t1 | FileCheck %s --check-prefix=MIR-tail
+; RUN: opt < %s -passes='pseudo-probe,simplifycfg' -S | FileCheck %s --check-prefix=SC
+
+declare i32 @f1()
+
+define i32 @foo(i1 %cond) {
+; JT-LABEL: @foo(
+; JT: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 1, i32 0, i64 -1)
+; ASM: pseudoprobe  6699318081062747564 1 0 0
+  %call = call i32 @f1()
+  br i1 %cond, label %T, label %F
+T:
+  br label %Merge
+F:
+  br label %Merge
+Merge:
+;; Check branch T and F are gone, and their probes (probe 2 and 3) are dangling.
+; JT-LABEL-NO: T
+; JT-LABEL-NO: F
+; JT-LABEL: Merge
+; JT: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 3, i32 2, i64 -1)
+; JT: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 2, i32 2, i64 -1)
+; JT: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 4, i32 0, i64 -1)
+; ASM: .pseudoprobe 6699318081062747564 3 0 2
+; ASM: .pseudoprobe 6699318081062747564 2 0 2
+; ASM: .pseudoprobe 6699318081062747564 4 0 0
+  ret i32 %call
+}
+
+;; Check block T and F are gone, and their probes (probe 2 and 3) are dangling.
+; MIR-tail: bb.0
+; MIR-tail: PSEUDO_PROBE [[#GUID:]], 1, 0, 0
+; MIR-tail: PSEUDO_PROBE [[#GUID:]], 2, 0, 2
+; MIR-tail: PSEUDO_PROBE [[#GUID:]], 3, 0, 2
+; MIR-tail: PSEUDO_PROBE [[#GUID:]], 4, 0, 0
+
+
+define void @foo2() {
+bb:
+  %tmp = call i32 @f1()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5, label %bb8
+
+bb2:
+  %tmp4 = icmp ne i32 %tmp, 1
+  switch i1 %tmp4, label %bb2 [
+  i1 0, label %bb5
+  i1 1, label %bb8
+  ]
+
+bb5:
+;; Check the pseudo probe with id 3 only has one copy.
+; JT-COUNT-1: call void @llvm.pseudoprobe(i64 [[#GUID2:]], i64 3, i32 2, i64 -1)
+  %tmp6 = phi i1 [ %tmp1, %bb ], [ false, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7
+
+bb7:
+  br label %bb8
+
+bb8:
+  ret void
+}
+
+define i32 @test(i32 %a, i32 %b, i32 %c) {
+;; Check block bb1 and bb2 are gone, and their probes (probe 2 and 3) are dangling.
+; SC-LABEL: @test(
+; SC-LABEL-NO: bb1
+; SC-LABEL-NO: bb2
+; SC:    [[T1:%.*]] = icmp eq i32 [[B:%.*]], 0
+; SC-DAG:    call void @llvm.pseudoprobe(i64 [[#GUID3:]], i64 2, i32 2, i64 -1)
+; SC-DAG:    call void @llvm.pseudoprobe(i64 [[#GUID3]], i64 3, i32 2, i64 -1)
+; SC:    [[T2:%.*]] = icmp sgt i32 [[C:%.*]], 1
+; SC:    [[T3:%.*]] = add i32 [[A:%.*]], 1
+; SC:    [[SPEC_SELECT:%.*]] = select i1 [[T2]], i32 [[T3]], i32 [[A]]
+; SC:    [[T4:%.*]] = select i1 [[T1]], i32 [[SPEC_SELECT]], i32 [[B]]
+; SC:    [[T5:%.*]] = sub i32 [[T4]], 1
+; SC:    ret i32 [[T5]]
+
+entry:
+  %t1 = icmp eq i32 %b, 0
+  br i1 %t1, label %bb1, label %bb3
+
+bb1:
+  %t2 = icmp sgt i32 %c, 1
+  br i1 %t2, label %bb2, label %bb3
+
+bb2:
+  %t3 = add i32 %a, 1
+  br label %bb3
+
+bb3:
+  %t4 = phi i32 [ %b, %entry ], [ %a, %bb1 ], [ %t3, %bb2 ]
+  %t5 = sub i32 %t4, 1
+  ret i32 %t5
+}