[PGO] Avoid potential const_cast UB (NFC)

[Dinistro]

Summary

This commit removes potential UB in the PGO instrumentation passes that
was caused by casting away constness and then potentially modifying the
object.

[gysit]

LGTM!

[xur]

Thanks for working on this.
PGOEdge is shared by PGO instrumentation pass and PGO profile-use pass. We cast away const for PGO instrumentation. But for PGO profile-use, we don't change the BB. I would prefer to keep the const for profile-use passes. So I think a better fix is to make PGOEdge a template class, one with const and one without.

[Dinistro]

Added a comment to the critical UB causing part. Again, while there might be a way to keep constness, it requires extended changes. Fixing UB should be done first, afterwards, we can think about refactorings that allow us to keep constness where possible.

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

817

This full function is shared between the instrumentation and the use pass. Keeping const will always result in UB, as the SplitCriticalEdge function will be called with TI, which should be of type const Instruction *.
I fail to see how a templated edge class would help in resolving this issue, could you elaborate @xur ?

I believe that splitting the logic between instrument and usage might help to reduce the complexity and push const in some places, but that requires extended changes to the algorithms.

[gysit]

@xur how should we proceed with this. It seems like eliminating UB in favor of less constness seems like a good initial step?

llvm/include/llvm/Transforms/Instrumentation/CFGMST.h

public:
   }
 
   // Traverse the CFG using a stack. Find all the edges and assign the weight.
   // Edges with large weight will be put into MST first so they are less likely
   // to be instrumented.
   void buildEdges() {
     LLVM_DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
 
-    const BasicBlock *Entry = &(F.getEntryBlock());
+    BasicBlock *Entry = &(F.getEntryBlock());
     uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
     // If we want to instrument the entry count, lower the weight to 0.
     if (InstrumentFuncEntry)
       EntryWeight = 0;
     Edge *EntryIncoming = nullptr, *EntryOutgoing = nullptr,
          *ExitOutgoing = nullptr, *ExitIncoming = nullptr;
     uint64_t MaxEntryOutWeight = 0, MaxExitOutWeight = 0, MaxExitInWeight = 0;
 
OS << "  Number of Edges: " << AllEdges.size()
        << " (*: Instrument, C: CriticalEdge, -: Removed)\n";
     uint32_t Count = 0;
     for (auto &EI : AllEdges)
       OS << "  Edge " << Count++ << ": " << getBBInfo(EI->SrcBB).Index << "-->"
          << getBBInfo(EI->DestBB).Index << EI->infoString() << "\n";
   }
 
   // Add an edge to AllEdges with weight W.
-  Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) {
+  Edge &addEdge(BasicBlock *Src, BasicBlock *Dest, uint64_t W) {
     uint32_t Index = BBInfos.size();
     auto Iter = BBInfos.end();
     bool Inserted;
     std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Src, nullptr));
     if (Inserted) {
       // Newly inserted, update the real info.
       Iter->second = std::move(std::make_unique<BBInfo>(Index));
       Index++;

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

struct SelectInstVisitor : public InstVisitor<SelectInstVisitor> {
   unsigned getNumOfSelectInsts() const { return NSIs; }
 };
 
 /// This class implements the CFG edges for the Minimum Spanning Tree (MST)
 /// based instrumentation.
 /// Note that the CFG can be a multi-graph. So there might be multiple edges
 /// with the same SrcBB and DestBB.
 struct PGOEdge {
-  const BasicBlock *SrcBB;
-  const BasicBlock *DestBB;
+  BasicBlock *SrcBB;
+  BasicBlock *DestBB;
   uint64_t Weight;
   bool InMST = false;
   bool Removed = false;
   bool IsCritical = false;
 
-  PGOEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W = 1)
+  PGOEdge(BasicBlock *Src, BasicBlock *Dest, uint64_t W = 1)
       : SrcBB(Src), DestBB(Dest), Weight(W) {}
 
   /// Return the information string of an edge.
   std::string infoString() const {
     return (Twine(Removed ? "-" : " ") + (InMST ? " " : "*") +
             (IsCritical ? "c" : " ") + "  W=" + Twine(Weight))
         .str();
   }

 
 // Given a CFG E to be instrumented, find which BB to place the instrumented
 // code. The function will split the critical edge if necessary.
 template <class Edge, class BBInfo>
 BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   if (E->InMST || E->Removed)
     return nullptr;
 
-  BasicBlock *SrcBB = const_cast<BasicBlock *>(E->SrcBB);
-  BasicBlock *DestBB = const_cast<BasicBlock *>(E->DestBB);
+  BasicBlock *SrcBB = E->SrcBB;
+  BasicBlock *DestBB = E->DestBB;
   // For a fake edge, instrument the real BB.
   if (SrcBB == nullptr)
     return DestBB;
   if (DestBB == nullptr)
     return SrcBB;
 
   auto canInstrument = [](BasicBlock *BB) -> BasicBlock * {
     // There are basic blocks (such as catchswitch) cannot be instrumented.
if (TI->getNumSuccessors() <= 1)
     return canInstrument(SrcBB);
   if (!E->IsCritical)
     return canInstrument(DestBB);
 
   // Some IndirectBr critical edges cannot be split by the previous
   // SplitIndirectBrCriticalEdges call. Bail out.
   unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
   BasicBlock *InstrBB =
       isa<IndirectBrInst>(TI) ? nullptr : SplitCriticalEdge(TI, SuccNum);
   if (!InstrBB) {
     LLVM_DEBUG(
         dbgs() << "Fail to split critical edge: not instrument this edge.\n");
     return nullptr;
   }
   // For a critical edge, we have to split. Instrument the newly
   // created BB.
   IsCS ? NumOfCSPGOSplit++ : NumOfPGOSplit++;
for (uint32_t Kind = IPVK_First; Kind <= IPVK_Last; ++Kind) {
     }
   } // IPVK_First <= Kind <= IPVK_Last
 }
 
 namespace {
 
 // This class represents a CFG edge in profile use compilation.
 struct PGOUseEdge : public PGOEdge {
+  using PGOEdge::PGOEdge;
+
   bool CountValid = false;
   uint64_t CountValue = 0;
 
-  PGOUseEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W = 1)
-      : PGOEdge(Src, Dest, W) {}
-
   // Set edge count value
   void setEdgeCount(uint64_t Value) {
     CountValue = Value;
     CountValid = true;
   }
 
   // Return the information string for this object.
   std::string infoString() const {