[BasicBlockSections] Introduce the path cloning profile format to BasicBlockSectionsProfileReader. (#67214)

Following up on prior RFC
(https://lists.llvm.org/pipermail/llvm-dev/2020-September/145357.html)
we can now improve above our highly-optimized basic-block-sections
binary (e.g., 2% for clang) by applying path cloning. Cloning can
improve performance by reducing taken branches.

This patch prepares the profile format for applying cloning actions.

The basic block cloning profile format extends the basic block sections
profile in two ways.

1. Specifies the cloning paths with a 'p' specifier. For example, `p 1 4
5` specifies that blocks with BB ids 4 and 5 must be cloned along the
edge 1 --> 4.
2. For each cloned block, it will appear in the cluster info as
`<bb_id>.<clone_id>` where `clone_id` is the id associated with this
clone.

For example, the following profile specifies one cloned block (2) and
determines its cluster position as well.
```
f foo
p 1 2
c 0 1 2.1 3 2 5
```

This patch keeps backward-compatibility (retains the behavior for old
profile formats). This feature is only introduced for profile version >=
1.

1 files changed, 30 insertions, 43 deletions

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index de7c170..632fd68 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -168,31 +168,6 @@ updateBranches(MachineFunction &MF,
   }
 }
 
-// This function provides the BBCluster information associated with a function.
-// Returns true if a valid association exists and false otherwise.
-bool getBBClusterInfoForFunction(
-    const MachineFunction &MF,
-    BasicBlockSectionsProfileReader *BBSectionsProfileReader,
-    DenseMap<unsigned, BBClusterInfo> &V) {
-
-  // Find the assoicated cluster information.
-  std::pair<bool, SmallVector<BBClusterInfo, 4>> P =
-      BBSectionsProfileReader->getBBClusterInfoForFunction(MF.getName());
-  if (!P.first)
-    return false;
-
-  if (P.second.empty()) {
-    // This indicates that sections are desired for all basic blocks of this
-    // function. We clear the BBClusterInfo vector to denote this.
-    V.clear();
-    return true;
-  }
-
-  for (const BBClusterInfo &BBCI : P.second)
-    V[BBCI.BBID] = BBCI;
-  return true;
-}
-
 // This function sorts basic blocks according to the cluster's information.
 // All explicitly specified clusters of basic blocks will be ordered
 // accordingly. All non-specified BBs go into a separate "Cold" section.
@@ -200,12 +175,12 @@ bool getBBClusterInfoForFunction(
 // clusters, they are moved into a single "Exception" section. Eventually,
 // clusters are ordered in increasing order of their IDs, with the "Exception"
 // and "Cold" succeeding all other clusters.
-// FuncBBClusterInfo represent the cluster information for basic blocks. It
+// ClusterInfoByBBID represents the cluster information for basic blocks. It
 // maps from BBID of basic blocks to their cluster information. If this is
 // empty, it means unique sections for all basic blocks in the function.
-static void
-assignSections(MachineFunction &MF,
-               const DenseMap<unsigned, BBClusterInfo> &FuncBBClusterInfo) {
+static void assignSections(
+    MachineFunction &MF,
+    const DenseMap<unsigned, BBClusterInfo<unsigned>> &ClusterInfoByBBID) {
   assert(MF.hasBBSections() && "BB Sections is not set for function.");
   // This variable stores the section ID of the cluster containing eh_pads (if
   // all eh_pads are one cluster). If more than one cluster contain eh_pads, we
@@ -216,17 +191,17 @@ assignSections(MachineFunction &MF,
     // With the 'all' option, every basic block is placed in a unique section.
     // With the 'list' option, every basic block is placed in a section
     // associated with its cluster, unless we want individual unique sections
-    // for every basic block in this function (if FuncBBClusterInfo is empty).
+    // for every basic block in this function (if ClusterInfoByBBID is empty).
     if (MF.getTarget().getBBSectionsType() == llvm::BasicBlockSection::All ||
-        FuncBBClusterInfo.empty()) {
+        ClusterInfoByBBID.empty()) {
       // If unique sections are desired for all basic blocks of the function, we
       // set every basic block's section ID equal to its original position in
       // the layout (which is equal to its number). This ensures that basic
       // blocks are ordered canonically.
       MBB.setSectionID(MBB.getNumber());
     } else {
-      auto I = FuncBBClusterInfo.find(*MBB.getBBID());
-      if (I != FuncBBClusterInfo.end()) {
+      auto I = ClusterInfoByBBID.find(*MBB.getBBID());
+      if (I != ClusterInfoByBBID.end()) {
         MBB.setSectionID(I->second.ClusterID);
       } else {
         // BB goes into the special cold section if it is not specified in the
@@ -333,16 +308,28 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
     return true;
   }
 
-  BBSectionsProfileReader = &getAnalysis<BasicBlockSectionsProfileReader>();
+  DenseMap<unsigned, BBClusterInfo<unsigned>> ClusterInfoByBBID;
+  if (BBSectionsType == BasicBlockSection::List) {
+    auto [HasProfile, PathAndClusterInfo] =
+        getAnalysis<BasicBlockSectionsProfileReader>()
+            .getPathAndClusterInfoForFunction(MF.getName());
+    if (!HasProfile)
+      return true;
+    for (const BBClusterInfo<ProfileBBID> &BBP :
+         PathAndClusterInfo.ClusterInfo) {
+      // TODO: Apply the path cloning profile.
+      assert(!BBP.BasicBlockID.CloneID && "Path cloning is not supported yet");
+      const auto [I, Inserted] = ClusterInfoByBBID.try_emplace(
+          BBP.BasicBlockID.BBID,
+          BBClusterInfo<unsigned>{BBP.BasicBlockID.BBID, BBP.ClusterID,
+                                  BBP.PositionInCluster});
+      (void)I;
+      assert(Inserted && "Duplicate BBID found in profile");
+    }
+  }
 
-  // Map from BBID of blocks to their cluster information.
-  DenseMap<unsigned, BBClusterInfo> FuncBBClusterInfo;
-  if (BBSectionsType == BasicBlockSection::List &&
-      !getBBClusterInfoForFunction(MF, BBSectionsProfileReader,
-                                   FuncBBClusterInfo))
-    return true;
   MF.setBBSectionsType(BBSectionsType);
-  assignSections(MF, FuncBBClusterInfo);
+  assignSections(MF, ClusterInfoByBBID);
 
   // We make sure that the cluster including the entry basic block precedes all
   // other clusters.
@@ -376,8 +363,8 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
     // If the two basic block are in the same section, the order is decided by
     // their position within the section.
     if (XSectionID.Type == MBBSectionID::SectionType::Default)
-      return FuncBBClusterInfo.lookup(*X.getBBID()).PositionInCluster <
-             FuncBBClusterInfo.lookup(*Y.getBBID()).PositionInCluster;
+      return ClusterInfoByBBID.lookup(*X.getBBID()).PositionInCluster <
+             ClusterInfoByBBID.lookup(*Y.getBBID()).PositionInCluster;
     return X.getNumber() < Y.getNumber();
   };