[llvm] [BasicBlockSections] Apply path cloning with -basic-block-sections. (PR #68860)

Rahman Lavaee via llvm-commits llvm-commits at lists.llvm.org
Thu Oct 19 14:49:36 PDT 2023


https://github.com/rlavaee updated https://github.com/llvm/llvm-project/pull/68860

>From f206269bb1e8c36277e8cfe99c36fa4e44af433b Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Thu, 12 Oct 2023 05:52:25 +0000
Subject: [PATCH 1/6] [BasicBlockSections] Apply path cloning with
 -basic-block-sections.

---
 .../CodeGen/BasicBlockSectionsProfileReader.h |   1 +
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 173 +++++++++++++++++-
 .../BasicBlockSectionsProfileReader.cpp       |   4 +
 .../X86/basic-block-sections-cloning.ll       | 169 +++++++++++++++++
 4 files changed, 337 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/basic-block-sections-cloning.ll

diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index 6e01dfd11ee6dad..e4cb7e5b6643d0a 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -25,6 +25,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/MemoryBuffer.h"
+using namespace llvm;
 
 namespace llvm {
 
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 632fd68d88b5c64..7913f069b0f1d72 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -77,8 +77,11 @@
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/WithColor.h"
 #include "llvm/Target/TargetMachine.h"
 #include <optional>
+#include <sstream>
 
 using namespace llvm;
 
@@ -99,6 +102,33 @@ static cl::opt<bool> BBSectionsDetectSourceDrift(
 
 namespace {
 
+MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock *MBB) {
+  auto &MF = *MBB->getParent();
+  auto TII = MF.getSubtarget().getInstrInfo();
+
+  auto CloneBB = MF.CreateMachineBasicBlock(MBB->getBasicBlock());
+  MF.push_back(CloneBB);
+  // Copy the instructions.
+  for (auto &I : MBB->instrs())
+    CloneBB->push_back(MF.CloneMachineInstr(&I));
+
+  // Add the successors of the original block as the new block's successors.
+  for (auto SI = MBB->succ_begin(), SE = MBB->succ_end(); SI != SE; ++SI)
+    CloneBB->copySuccessor(MBB, SI);
+
+  if (auto FT = MBB->getFallThrough(/*JumpToFallThrough=*/false)) {
+    // The original block has an implicit fall through.
+    // Insert an explicit unconditional jump from the cloned block to the
+    // fallthrough block.
+    TII->insertUnconditionalBranch(*CloneBB, FT, CloneBB->findBranchDebugLoc());
+  }
+
+  for (auto &LiveIn : MBB->liveins())
+    CloneBB->addLiveIn(LiveIn);
+
+  return CloneBB;
+}
+
 class BasicBlockSections : public MachineFunctionPass {
 public:
   static char ID;
@@ -285,6 +315,128 @@ static bool hasInstrProfHashMismatch(MachineFunction &MF) {
   return false;
 }
 
+// Returns `Error::success` if we can legally apply all clonings specified in
+// `ClonePaths`, and `Error` otherwise.
+static Error
+CheckValidCloning(const MachineFunction &MF,
+                  const SmallVector<SmallVector<unsigned>> &ClonePaths) {
+  // Map from the final BB IDs to the `MachineBasicBlock`s.
+  DenseMap<unsigned, const MachineBasicBlock *> BBIDToBlock;
+  for (auto &BB : MF)
+    BBIDToBlock.try_emplace(*BB.getBBID(), &BB);
+
+  for (auto &ClonePath : ClonePaths) {
+    const MachineBasicBlock *PrevBB = nullptr;
+    for (size_t I = 0; I < ClonePath.size(); ++I) {
+      unsigned BBID = ClonePath[I];
+      const MachineBasicBlock *PathBB = BBIDToBlock.lookup(BBID);
+      if (!PathBB) {
+        return make_error<StringError>(Twine("no block with id ") +
+                                           Twine(BBID) + " in function " +
+                                           MF.getName(),
+                                       inconvertibleErrorCode());
+      }
+
+      if (PrevBB && !PrevBB->isSuccessor(PathBB)) {
+        return make_error<StringError>(
+            Twine("block #") + Twine(BBID) + " is not a successor of block #" +
+                Twine(*PrevBB->getBBID()) + " in function " + MF.getName(),
+            inconvertibleErrorCode());
+      }
+
+      if (I != ClonePath.size() - 1 && !PathBB->empty() &&
+          PathBB->back().isIndirectBranch()) {
+        return make_error<StringError>(
+            Twine("block #") + Twine(BBID) +
+                " has indirect branch can only appear as the last block of the "
+                "path, in function " +
+                MF.getName(),
+            inconvertibleErrorCode());
+      }
+      PrevBB = PathBB;
+    }
+  }
+  return Error::success();
+}
+
+// Applies all clonings specified in `ClonePaths` to `MF` and returns a map
+// from `ProfileBBID`s of all clone blocks to their BBIDs (assigned by
+// `MachineFunction`).
+static DenseMap<ProfileBBID, unsigned>
+ApplyCloning(MachineFunction &MF,
+             const SmallVector<SmallVector<unsigned>> &ClonePaths) {
+  DenseMap<ProfileBBID, unsigned> CloneBBIDMap;
+  // Map from the final BB IDs to the `MachineBasicBlock`s.
+  DenseMap<unsigned, MachineBasicBlock *> BBIDToBlock;
+  for (auto &BB : MF)
+    BBIDToBlock.try_emplace(*BB.getBBID(), &BB);
+
+  DenseMap<unsigned, unsigned> NClonesForBBID;
+  auto TII = MF.getSubtarget().getInstrInfo();
+  for (const auto &ClonePath : ClonePaths) {
+
+    MachineBasicBlock *PrevBB = nullptr;
+    for (unsigned BBID : ClonePath) {
+      MachineBasicBlock *OrigBB = BBIDToBlock.at(BBID);
+      if (PrevBB == nullptr) {
+        if (auto FT = OrigBB->getFallThrough(/*JumpToFallThrough=*/false)) {
+          // Make fallthroughs explicit since we may change the layout.
+          TII->insertUnconditionalBranch(*OrigBB, FT,
+                                         OrigBB->findBranchDebugLoc());
+        }
+        PrevBB = OrigBB;
+        continue;
+      }
+
+      MachineBasicBlock *CloneBB = CloneMachineBasicBlock(OrigBB);
+
+      CloneBBIDMap.try_emplace({BBID, ++NClonesForBBID[BBID]},
+                               *CloneBB->getBBID());
+      // Set up the previous block in the path to jump to the clone.
+      PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
+      PrevBB = CloneBB;
+    }
+  }
+  return CloneBBIDMap;
+}
+
+// Processes the raw input profile `PathAndClusterInfo` for the given function
+// `MF` and returns the final profile (cluster information). Returns error if
+// `PathAndClusterInfo` is invalid and cannot be applied (e.g., invalid cloning
+// paths).
+//
+// Steps:
+//    1. Clones basic blocks based on `PathAndClusterInfo.ClonePaths` (if not
+//    empty) and
+//       updates the CFG accordingly.
+//    2. Creates and returns the cluster profile for the final blocks (original
+//       and cloned) based on `PathAndClusterInfo.ClusterInfo`.
+static Expected<DenseMap<unsigned, BBClusterInfo<unsigned>>>
+ProcessProfile(MachineFunction &MF,
+               const FunctionPathAndClusterInfo &PathAndClusterInfo) {
+  if (auto Err = CheckValidCloning(MF, PathAndClusterInfo.ClonePaths))
+    return std::move(Err);
+
+  // Apply the clonings and obtain the map from the input block ID of cloned
+  // blocks to their final BB IDs.
+  DenseMap<ProfileBBID, unsigned> CloneBBIDMap =
+      ApplyCloning(MF, PathAndClusterInfo.ClonePaths);
+
+  // Map from final BB IDs to their profile information.
+  DenseMap<unsigned, BBClusterInfo<unsigned>> ClusterInfoByBBID;
+  // This step creates all the necessary clones. It does not adjust the
+  // branches.
+  for (const BBClusterInfo<ProfileBBID> &P : PathAndClusterInfo.ClusterInfo) {
+    unsigned FinalBBID = P.BasicBlockID.CloneID == 0
+                             ? P.BasicBlockID.BBID
+                             : CloneBBIDMap.at(P.BasicBlockID);
+    ClusterInfoByBBID.try_emplace(
+        FinalBBID,
+        BBClusterInfo<unsigned>{FinalBBID, P.ClusterID, P.PositionInCluster});
+  }
+  return ClusterInfoByBBID;
+}
+
 bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   auto BBSectionsType = MF.getTarget().getBBSectionsType();
   assert(BBSectionsType != BasicBlockSection::None &&
@@ -315,17 +467,18 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
             .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");
+    auto ClusterInfoOrErr = ProcessProfile(MF, PathAndClusterInfo);
+    if (!ClusterInfoOrErr) {
+      auto Err =
+          handleErrors(ClusterInfoOrErr.takeError(), [&](const StringError &E) {
+            WithColor::warning()
+                << "Rejecting the BBSections profile for function "
+                << MF.getName() << " : " << E.getMessage() << "\n";
+          });
+      assert(!Err);
+      return true;
     }
+    ClusterInfoByBBID = *std::move(ClusterInfoOrErr);
   }
 
   MF.setBBSectionsType(BBSectionsType);
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index 6bb412a6c7534a6..650c2770ae3bee1 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -217,6 +217,10 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       CurrentCluster++;
       continue;
     case 'p': { // Basic block cloning path specifier.
+      // Skip the profile when we the profile iterator (FI) refers to the
+      // past-the-end element.
+      if (FI == ProgramPathAndClusterInfo.end())
+        continue;
       SmallSet<unsigned, 5> BBsInPath;
       FI->second.ClonePaths.push_back({});
       for (size_t I = 0; I < Values.size(); ++I) {
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
new file mode 100644
index 000000000000000..0cca18b5aa8ad3a
--- /dev/null
+++ b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
@@ -0,0 +1,169 @@
+;; Tests for path cloning with -basic-block-sections.
+
+declare void @effect(i32 zeroext)
+
+;; Test valid application of path cloning.
+; RUN: echo 'v1' > %t1
+; RUN: echo 'f foo' >> %t1
+; RUN: echo 'p 0 3 5' >> %t1
+; RUN: echo 'c 0 3.1 5.1 1 2 3 4 5' >> %t1
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS
+; RUN: echo 'v1' > %t2
+; RUN: echo 'f foo' >> %t2
+; RUN: echo 'p 0 3 5' >> %t2
+; RUN: echo 'p 1 3 4 5' >> %t2
+; RUN: echo 'c 0 3.1 5.1' >> %t2
+; RUN: echo 'c 1 3.2 4.1 5.2 2 3 4 5' >> %t2
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS
+
+;; Test failed application of path cloning.
+; RUN: echo 'v1' > %t3
+; RUN: echo 'f foo' >> %t3
+; RUN: echo 'p 0 2 3' >> %t3
+; RUN: echo 'c 0 2.1 3.1' >> %t3
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3 2> %t3.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: FileCheck %s --check-prefixes=PATH3-WARN < %t3.err
+; RUN: echo 'v1' > %t4
+; RUN: echo 'f foo' >> %t4
+; RUN: echo 'p 0 1 100' >> %t4
+; RUN: echo 'c 0' >> %t3
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t4 2> %t4.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: FileCheck %s --check-prefixes=PATH4-WARN < %t4.err
+
+define void @foo(i1 %a, i1 %b, i1 %c, i1 %d) {
+b0:
+  call void @effect(i32 0)
+  br i1 %a, label %b1, label %b3
+
+b1:                                           ; preds = %b0
+  call void @effect(i32 1)
+  br i1 %b, label %b2, label %b3
+
+b2:                                             ; preds = %b1
+  call void @effect(i32 2)
+  br label %b3
+
+b3:                                            ; preds = %b0, %b1, %b2
+  call void @effect(i32 3)
+  br i1 %c, label %b4, label %b5
+
+b4:                                             ; preds = %b3
+  call void @effect(i32 4)
+  br i1 %d, label %b5, label %cold
+
+b5:                                            ; preds = %b3, %b4
+  call void @effect(i32 5)
+  ret void
+cold:
+  call void @effect(i32 6)                     ; preds = %b4
+  ret void
+}
+
+; PATH1:         .section    .text.foo,"ax", at progbits
+; PATH1-LABEL: foo:
+; PATH1:       # %bb.0:        # %b0
+; PATH1:         jne .LBB0_1
+; PATH1-NEXT:  # %bb.7:        # %b3
+; PATH1:         jne .LBB0_4
+; PATH1-NEXT:  # %bb.8:        # %b5
+; PATH1:        retq
+; PATH1-NEXT:  .LBB0_1:        # %b1
+; PATH1:         je .LBB0_3
+; PATH1-NEXT:  # %bb.2:        # %b2
+; PATH1:         callq effect at PLT
+; PATH1-NEXT:  .LBB0_3:        # %b3
+; PATH1:         je .LBB0_5
+; PATH1-NEXT:  .LBB0_4:        # %b4
+; PATH1:         je foo.cold
+; PATH1-NEXT:  .LBB0_5:        # %b5
+; PATH1:         retq
+
+; PATH2:         .section    .text.foo,"ax", at progbits
+; PATH2-LABEL: foo:
+; PATH2:       # %bb.0:        # %b0
+; PATH2:         jne foo.__part.1
+; PATH2-NEXT:  # %bb.7:        # %b3
+; PATH2:         jne .LBB0_4
+; PATH2-NEXT:  # %bb.8:        # %b5
+; PATH2:        retq
+; PATH2:         .section    .text.foo,"ax", at progbits,unique,1
+; PATH2-NEXT:  foo.__part.1:   # %b1
+; PATH2:         jne .LBB0_2
+; PATH2-NEXT:  # %bb.9:        # %b3
+; PATH2:         je .LBB0_5
+; PATH2-NEXT:  # %bb.10:       # %b4
+; PATH2:         je foo.cold
+; PATH2-NEXT:  # %bb.11:       # %b5
+; PATH2:         retq
+; PATH2-NEXT:  .LBB0_2:        # %b2
+; PATH2:         callq	effect at PLT
+; PATH2-NEXT:  # %bb.3:        # %b3
+; PATH2:         je .LBB0_5
+; PATH2-NEXT:  .LBB0_4:        # %b4
+; PATH2:          je foo.cold
+; PATH2-NEXT:  .LBB0_5:       # %b5
+; PATH2:          retq
+; FOOSECTIONS: foo.cold
+
+; PATH3-WARN: warning: Rejecting the BBSections profile for function foo : block #2 is not a successor of block #0 in function foo
+; PATH4-WARN: warning: Rejecting the BBSections profile for function foo : no block with id 100 in function foo
+; NOSECTIONS-NOT: foo.cold
+
+
+;; Test valid application of cloning for paths with indirect branches.
+; RUN: echo 'v1' > %t5
+; RUN: echo 'f bar' >> %t5
+; RUN: echo 'p 0 1' >> %t5
+; RUN: echo 'c 0 1.1 2 1' >> %t5
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t5 | FileCheck %s --check-prefix=PATH5
+; RUN: echo 'v1' > %t6
+; RUN: echo 'f bar' >> %t6
+; RUN: echo 'p 0 1 2' >> %t6
+; RUN: echo 'c 0 1.1 2.1 1' >> %t6
+
+;; Test failed application of path cloning for paths with indirect branches.
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t6 2> %t6.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t6.err
+; RUN: echo 'v1' > %t7
+; RUN: echo 'f bar' >> %t7
+; RUN: echo 'p 1 2' >> %t7
+; RUN: echo 'c 0 1 2.1' >> %t7
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t7 2> %t7.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t7.err
+
+
+define void @bar(i1 %a, i1 %b) {
+b0:
+  call void @effect(i32 0)
+  br i1 %a, label %b1, label %b2
+b1:                                              ; preds = %b0
+  call void @effect(i32 1)
+  %0 = select i1 %b,                           ; <ptr> [#uses=1]
+              ptr blockaddress(@bar, %b2),
+              ptr blockaddress(@bar, %b3)
+  indirectbr ptr %0, [label %b2, label %b3]
+b2:                                              ; preds = %b0, %b1
+  call void @effect(i32 2)
+  ret void
+b3:
+  call void @effect(i32 3)                       ; preds = %b1
+  ret void
+}
+
+
+; PATH5:         .section    .text.bar,"ax", at progbits
+; PATH5-LABEL: bar:
+; PATH5:       # %bb.0:        # %b0
+; PATH5:         je .LBB1_2
+; PATH5-NEXT:  # %bb.4:        # %b1
+; PATH5:         jmpq *%rax
+; PATH5-NEXT:  .Ltmp0:         # Block address taken
+; PATH5-NEXT:  .LBB1_2:        # %b2
+; PATH5:         retq
+; PATH5-NEXT:  # %bb.1:        # %b1
+; PATH5:         jmpq *%rax
+; PATH5:         .section    .text.split.bar,"ax", at progbits
+; PATH5:       bar.cold:       # %b3
+; NOSECTIONS-NOT: bar.cold
+
+; PATH-INDIR-WARN: warning: Rejecting the BBSections profile for function bar : block #1 has indirect branch can only appear as the last block of the path, in function bar

>From 5a8dca654f632e17502af2367e19b88674820ec5 Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Mon, 16 Oct 2023 19:37:50 +0000
Subject: [PATCH 2/6] Apply BB sections even if cloning fails.

---
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 112 +++++++++---------
 .../X86/basic-block-sections-cloning.ll       |  92 ++++++++------
 2 files changed, 108 insertions(+), 96 deletions(-)

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 7913f069b0f1d72..248cfca6b94fd77 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -77,7 +77,6 @@
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
-#include "llvm/Support/Error.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Target/TargetMachine.h"
 #include <optional>
@@ -315,48 +314,41 @@ static bool hasInstrProfHashMismatch(MachineFunction &MF) {
   return false;
 }
 
-// Returns `Error::success` if we can legally apply all clonings specified in
-// `ClonePaths`, and `Error` otherwise.
-static Error
-CheckValidCloning(const MachineFunction &MF,
-                  const SmallVector<SmallVector<unsigned>> &ClonePaths) {
-  // Map from the final BB IDs to the `MachineBasicBlock`s.
-  DenseMap<unsigned, const MachineBasicBlock *> BBIDToBlock;
-  for (auto &BB : MF)
-    BBIDToBlock.try_emplace(*BB.getBBID(), &BB);
-
-  for (auto &ClonePath : ClonePaths) {
-    const MachineBasicBlock *PrevBB = nullptr;
-    for (size_t I = 0; I < ClonePath.size(); ++I) {
-      unsigned BBID = ClonePath[I];
-      const MachineBasicBlock *PathBB = BBIDToBlock.lookup(BBID);
-      if (!PathBB) {
-        return make_error<StringError>(Twine("no block with id ") +
-                                           Twine(BBID) + " in function " +
-                                           MF.getName(),
-                                       inconvertibleErrorCode());
-      }
+// Returns if we can legally apply all clonings specified in `ClonePaths`.
+static bool
+IsValidCloning(const MachineFunction &MF,
+               const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
+               const SmallVector<unsigned> &ClonePath) {
+  const MachineBasicBlock *PrevBB = nullptr;
+  for (size_t I = 0; I < ClonePath.size(); ++I) {
+    unsigned BBID = ClonePath[I];
+    const MachineBasicBlock *PathBB = BBIDToBlock.lookup(BBID);
+    if (!PathBB) {
+      WithColor::warning() << "no block with id " << BBID << " in function "
+                           << MF.getName() << "\n";
+      return false;
+    }
 
-      if (PrevBB && !PrevBB->isSuccessor(PathBB)) {
-        return make_error<StringError>(
-            Twine("block #") + Twine(BBID) + " is not a successor of block #" +
-                Twine(*PrevBB->getBBID()) + " in function " + MF.getName(),
-            inconvertibleErrorCode());
-      }
+    if (PrevBB && !PrevBB->isSuccessor(PathBB)) {
+      WithColor::warning() << "block #" << BBID
+                           << " is not a successor of block #"
+                           << *PrevBB->getBBID() << " in function "
+                           << MF.getName() << "\n";
+      return false;
+    }
 
-      if (I != ClonePath.size() - 1 && !PathBB->empty() &&
-          PathBB->back().isIndirectBranch()) {
-        return make_error<StringError>(
-            Twine("block #") + Twine(BBID) +
-                " has indirect branch can only appear as the last block of the "
-                "path, in function " +
-                MF.getName(),
-            inconvertibleErrorCode());
-      }
-      PrevBB = PathBB;
+    if (I != ClonePath.size() - 1 && !PathBB->empty() &&
+        PathBB->back().isIndirectBranch()) {
+      WithColor::warning()
+          << "block #" << BBID
+          << " has indirect branch and appears as the non-tail block of a "
+             "path in function "
+          << MF.getName() << "\n";
+      return false;
     }
+    PrevBB = PathBB;
   }
-  return Error::success();
+  return true;
 }
 
 // Applies all clonings specified in `ClonePaths` to `MF` and returns a map
@@ -374,7 +366,13 @@ ApplyCloning(MachineFunction &MF,
   DenseMap<unsigned, unsigned> NClonesForBBID;
   auto TII = MF.getSubtarget().getInstrInfo();
   for (const auto &ClonePath : ClonePaths) {
-
+    if (!IsValidCloning(MF, BBIDToBlock, ClonePath)) {
+      // We still need to increment the number of clones so we can map
+      // to the cluster info correctly.
+      for (unsigned BBID : ClonePath)
+        ++NClonesForBBID[BBID];
+      continue;
+    }
     MachineBasicBlock *PrevBB = nullptr;
     for (unsigned BBID : ClonePath) {
       MachineBasicBlock *OrigBB = BBIDToBlock.at(BBID);
@@ -411,12 +409,9 @@ ApplyCloning(MachineFunction &MF,
 //       updates the CFG accordingly.
 //    2. Creates and returns the cluster profile for the final blocks (original
 //       and cloned) based on `PathAndClusterInfo.ClusterInfo`.
-static Expected<DenseMap<unsigned, BBClusterInfo<unsigned>>>
+static DenseMap<unsigned, BBClusterInfo<unsigned>>
 ProcessProfile(MachineFunction &MF,
                const FunctionPathAndClusterInfo &PathAndClusterInfo) {
-  if (auto Err = CheckValidCloning(MF, PathAndClusterInfo.ClonePaths))
-    return std::move(Err);
-
   // Apply the clonings and obtain the map from the input block ID of cloned
   // blocks to their final BB IDs.
   DenseMap<ProfileBBID, unsigned> CloneBBIDMap =
@@ -427,9 +422,15 @@ ProcessProfile(MachineFunction &MF,
   // This step creates all the necessary clones. It does not adjust the
   // branches.
   for (const BBClusterInfo<ProfileBBID> &P : PathAndClusterInfo.ClusterInfo) {
-    unsigned FinalBBID = P.BasicBlockID.CloneID == 0
-                             ? P.BasicBlockID.BBID
-                             : CloneBBIDMap.at(P.BasicBlockID);
+    unsigned FinalBBID = P.BasicBlockID.BBID;
+    if (P.BasicBlockID.CloneID != 0) {
+      auto I = CloneBBIDMap.find(P.BasicBlockID);
+      // If there is no match, it means cloning was not applied for that block.
+      // We simply ignore these entries.
+      if (I == CloneBBIDMap.end())
+        continue;
+      FinalBBID = I->second;
+    }
     ClusterInfoByBBID.try_emplace(
         FinalBBID,
         BBClusterInfo<unsigned>{FinalBBID, P.ClusterID, P.PositionInCluster});
@@ -467,18 +468,13 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
             .getPathAndClusterInfoForFunction(MF.getName());
     if (!HasProfile)
       return true;
-    auto ClusterInfoOrErr = ProcessProfile(MF, PathAndClusterInfo);
-    if (!ClusterInfoOrErr) {
-      auto Err =
-          handleErrors(ClusterInfoOrErr.takeError(), [&](const StringError &E) {
-            WithColor::warning()
-                << "Rejecting the BBSections profile for function "
-                << MF.getName() << " : " << E.getMessage() << "\n";
-          });
-      assert(!Err);
-      return true;
+    // Empty ClusterInfo represents basic block sections for all blocks in this
+    // function.
+    if (!PathAndClusterInfo.ClusterInfo.empty()) {
+      ClusterInfoByBBID = ProcessProfile(MF, PathAndClusterInfo);
+      if (ClusterInfoByBBID.empty())
+        return true;
     }
-    ClusterInfoByBBID = *std::move(ClusterInfoOrErr);
   }
 
   MF.setBBSectionsType(BBSectionsType);
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
index 0cca18b5aa8ad3a..1c383771df77b1d 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
@@ -7,27 +7,35 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'f foo' >> %t1
 ; RUN: echo 'p 0 3 5' >> %t1
 ; RUN: echo 'c 0 3.1 5.1 1 2 3 4 5' >> %t1
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS,FOOCLONE
 ; RUN: echo 'v1' > %t2
 ; RUN: echo 'f foo' >> %t2
 ; RUN: echo 'p 0 3 5' >> %t2
 ; RUN: echo 'p 1 3 4 5' >> %t2
 ; RUN: echo 'c 0 3.1 5.1' >> %t2
 ; RUN: echo 'c 1 3.2 4.1 5.2 2 3 4 5' >> %t2
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS,FOOCLONE
 
 ;; Test failed application of path cloning.
 ; RUN: echo 'v1' > %t3
 ; RUN: echo 'f foo' >> %t3
 ; RUN: echo 'p 0 2 3' >> %t3
-; RUN: echo 'c 0 2.1 3.1' >> %t3
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3 2> %t3.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: echo 'c 0 2.1 3.1 1' >> %t3
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3 2> %t3.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH3-WARN < %t3.err
+;; Test that valid clonings are applied correctly, even if invalid clonings exist.
+; RUN: echo 'v1' > %t3_1
+; RUN: echo 'f foo' >> %t3_1
+; RUN: echo 'p 0 2 3' >> %t3_1
+; RUN: echo 'p 0 1 3' >> %t3_1
+; RUN: echo 'c 0 1.1 3.2 2.1 3.1 1' >> %t3_1
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3_1 2> %t3_1.err | FileCheck %s --check-prefixes=PATH3_1,FOONOCLONE,FOOSECTIONS
+; RUN: FileCheck %s --check-prefixes=PATH3-WARN < %t3_1.err
 ; RUN: echo 'v1' > %t4
 ; RUN: echo 'f foo' >> %t4
-; RUN: echo 'p 0 1 100' >> %t4
-; RUN: echo 'c 0' >> %t3
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t4 2> %t4.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: echo 'p 0 100' >> %t4
+; RUN: echo 'c 0 100.1 1' >> %t4
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t4 2> %t4.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH4-WARN < %t4.err
 
 define void @foo(i1 %a, i1 %b, i1 %c, i1 %d) {
@@ -59,9 +67,17 @@ cold:
   ret void
 }
 
-; PATH1:         .section    .text.foo,"ax", at progbits
-; PATH1-LABEL: foo:
-; PATH1:       # %bb.0:        # %b0
+; FOOSECTIONS:   .section    .text.foo,"ax", at progbits
+; FOOSECTIONS: foo:
+; FOOSECTIONS: # %bb.0:        # %b0
+
+; FOONOCLONE:    je .LBB0_3
+; PATH3_1:     # %bb.7:      # %b1
+; PATH3_1:     # %bb.8:      # %b3
+; PATH3_1:       jne .LBB0_4
+; FOONOCLONE:  # %bb.1:      # %b1
+; FOONOCLONE:    jne foo.cold
+
 ; PATH1:         jne .LBB0_1
 ; PATH1-NEXT:  # %bb.7:        # %b3
 ; PATH1:         jne .LBB0_4
@@ -78,9 +94,6 @@ cold:
 ; PATH1-NEXT:  .LBB0_5:        # %b5
 ; PATH1:         retq
 
-; PATH2:         .section    .text.foo,"ax", at progbits
-; PATH2-LABEL: foo:
-; PATH2:       # %bb.0:        # %b0
 ; PATH2:         jne foo.__part.1
 ; PATH2-NEXT:  # %bb.7:        # %b3
 ; PATH2:         jne .LBB0_4
@@ -103,11 +116,12 @@ cold:
 ; PATH2:          je foo.cold
 ; PATH2-NEXT:  .LBB0_5:       # %b5
 ; PATH2:          retq
-; FOOSECTIONS: foo.cold
+; FOOSECTIONS:        .section    .text.split.foo,"ax", at progbits
+; FOOCLONE:    foo.cold:      # %cold
+; FOONOCLONE:  foo.cold:      # %b2
 
-; PATH3-WARN: warning: Rejecting the BBSections profile for function foo : block #2 is not a successor of block #0 in function foo
-; PATH4-WARN: warning: Rejecting the BBSections profile for function foo : no block with id 100 in function foo
-; NOSECTIONS-NOT: foo.cold
+; PATH3-WARN: warning: block #2 is not a successor of block #0 in function foo
+; PATH4-WARN: warning: no block with id 100 in function foo
 
 
 ;; Test valid application of cloning for paths with indirect branches.
@@ -115,20 +129,20 @@ cold:
 ; RUN: echo 'f bar' >> %t5
 ; RUN: echo 'p 0 1' >> %t5
 ; RUN: echo 'c 0 1.1 2 1' >> %t5
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t5 | FileCheck %s --check-prefix=PATH5
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t5 | FileCheck %s --check-prefixes=PATH5,BARSECTIONS
+
+;; Test failed application of path cloning for paths with indirect branches.
 ; RUN: echo 'v1' > %t6
 ; RUN: echo 'f bar' >> %t6
 ; RUN: echo 'p 0 1 2' >> %t6
 ; RUN: echo 'c 0 1.1 2.1 1' >> %t6
-
-;; Test failed application of path cloning for paths with indirect branches.
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t6 2> %t6.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t6 2> %t6.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t6.err
 ; RUN: echo 'v1' > %t7
 ; RUN: echo 'f bar' >> %t7
 ; RUN: echo 'p 1 2' >> %t7
 ; RUN: echo 'c 0 1 2.1' >> %t7
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t7 2> %t7.err | FileCheck %s --check-prefix=NOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t7 2> %t7.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t7.err
 
 
@@ -151,19 +165,21 @@ b3:
 }
 
 
-; PATH5:         .section    .text.bar,"ax", at progbits
-; PATH5-LABEL: bar:
-; PATH5:       # %bb.0:        # %b0
-; PATH5:         je .LBB1_2
-; PATH5-NEXT:  # %bb.4:        # %b1
-; PATH5:         jmpq *%rax
-; PATH5-NEXT:  .Ltmp0:         # Block address taken
-; PATH5-NEXT:  .LBB1_2:        # %b2
-; PATH5:         retq
-; PATH5-NEXT:  # %bb.1:        # %b1
-; PATH5:         jmpq *%rax
-; PATH5:         .section    .text.split.bar,"ax", at progbits
-; PATH5:       bar.cold:       # %b3
-; NOSECTIONS-NOT: bar.cold
-
-; PATH-INDIR-WARN: warning: Rejecting the BBSections profile for function bar : block #1 has indirect branch can only appear as the last block of the path, in function bar
+; BARSECTIONS:       .section    .text.bar,"ax", at progbits
+; BARSECTIONS:     bar:
+; BARSECTIONS:     # %bb.0:        # %b0
+; PATH5:             je .LBB1_2
+; BARNOCLONE:        je bar.cold
+; BARNOCLONE-NEXT: # %bb.1:        # %b1
+; PATH5-NEXT:      # %bb.4:        # %b1
+; PATH5:             jmpq *%rax
+; PATH5-NEXT:      .Ltmp0:         # Block address taken
+; PATH5-NEXT:      .LBB1_2:        # %b2
+; PATH5:             retq
+; PATH5-NEXT:      # %bb.1:        # %b1
+; PATH5:             jmpq *%rax
+; BARSECTIONS:       .section    .text.split.bar,"ax", at progbits
+; PATH5:           bar.cold:       # %b3
+; BARNOCLONE:      bar.cold:       # %b2   
+
+; PATH-INDIR-WARN: warning: block #1 has indirect branch and appears as the non-tail block of a path in function bar

>From 9f55112572609294491740bd2cc2d216fd3b8bff Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Wed, 18 Oct 2023 22:40:03 +0000
Subject: [PATCH 3/6] Introduce a new pass for path cloning.

---
 .../llvm/CodeGen/BasicBlockSectionUtils.h     |   7 +
 .../CodeGen/BasicBlockSectionsProfileReader.h |  57 ++---
 llvm/include/llvm/CodeGen/MachineBasicBlock.h |  15 +-
 llvm/include/llvm/CodeGen/MachineFunction.h   |   5 +-
 llvm/include/llvm/CodeGen/Passes.h            |   2 +
 llvm/include/llvm/InitializePasses.h          |   1 +
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |   4 +-
 llvm/lib/CodeGen/BasicBlockPathCloning.cpp    | 213 ++++++++++++++++++
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 200 ++--------------
 .../BasicBlockSectionsProfileReader.cpp       |  48 ++--
 llvm/lib/CodeGen/CMakeLists.txt               |   1 +
 llvm/lib/CodeGen/CodeGen.cpp                  |   1 +
 llvm/lib/CodeGen/MIRParser/MIParser.cpp       |  18 +-
 llvm/lib/CodeGen/MachineBasicBlock.cpp        |   6 +-
 llvm/lib/CodeGen/MachineFunction.cpp          |   6 +-
 llvm/lib/CodeGen/TargetPassConfig.cpp         |   7 +
 .../X86/basic-block-sections-cloning.ll       |  16 +-
 17 files changed, 345 insertions(+), 262 deletions(-)
 create mode 100644 llvm/lib/CodeGen/BasicBlockPathCloning.cpp

diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionUtils.h b/llvm/include/llvm/CodeGen/BasicBlockSectionUtils.h
index d43f399b2c310a3..b3b6721c904c4ee 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionUtils.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionUtils.h
@@ -27,6 +27,13 @@ void sortBasicBlocksAndUpdateBranches(MachineFunction &MF,
 
 void avoidZeroOffsetLandingPad(MachineFunction &MF);
 
+// This checks if the source of this function has drifted since this binary was
+// profiled previously.  For now, we are piggy backing on what PGO does to
+// detect this with instrumented profiles.  PGO emits an hash of the IR and
+// checks if the hash has changed.  Advanced basic block layout is usually done
+// on top of PGO optimized binaries and hence this check works well in practice.
+bool hasInstrProfHashMismatch(MachineFunction &MF);
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_BASICBLOCKSECTIONUTILS_H
diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index e4cb7e5b6643d0a..60214db668a6d12 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -19,6 +19,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
@@ -29,24 +30,11 @@ using namespace llvm;
 
 namespace llvm {
 
-// This structure represents a unique ID for every block specified in the
-// input profile.
-struct ProfileBBID {
-  // Basic block id associated with `MachineBasicBlock::BBID`.
-  unsigned BBID;
-  // The clone id associated with the block. This is zero for the original
-  // block. For the cloned ones, it is equal to 1 + index of the associated
-  // path in `FunctionPathAndClusterInfo::ClonePaths`.
-  unsigned CloneID;
-};
-
 // This struct represents the cluster information for a machine basic block,
-// which is specifed by a unique ID. This templated struct is used for both the
-// raw input profile (as `BBClusterInfo<ProfileBBID>`) and the processed profile
-// after applying the clonings (as `BBClusterInfo<unsigned>`).
-template <typename BBIDType> struct BBClusterInfo {
+// which is specifed by a unique ID (`MachineBasicBlock::BBID`).
+struct BBClusterInfo {
   // Basic block ID.
-  BBIDType BasicBlockID;
+  UniqueBBID BBID;
   // Cluster ID this basic block belongs to.
   unsigned ClusterID;
   // Position of basic block within the cluster.
@@ -55,31 +43,31 @@ template <typename BBIDType> struct BBClusterInfo {
 
 // This represents the raw input profile for one function.
 struct FunctionPathAndClusterInfo {
-  // BB Cluster information specified by `ProfileBBID`s (before cloning).
-  SmallVector<BBClusterInfo<ProfileBBID>> ClusterInfo;
+  // BB Cluster information specified by `UniqueBBID`s.
+  SmallVector<BBClusterInfo> ClusterInfo;
   // Paths to clone. A path a -> b -> c -> d implies cloning b, c, and d along
   // the edge a -> b (a is not cloned). The index of the path in this vector
-  // determines the `ProfileBBID::CloneID` of the cloned blocks in that path.
+  // determines the `UniqueBBID::CloneID` of the cloned blocks in that path.
   SmallVector<SmallVector<unsigned>> ClonePaths;
 };
 
-// Provides DenseMapInfo for ProfileBBID.
-template <> struct DenseMapInfo<ProfileBBID> {
-  static inline ProfileBBID getEmptyKey() {
+// Provides DenseMapInfo for UniqueBBID.
+template <> struct DenseMapInfo<UniqueBBID> {
+  static inline UniqueBBID getEmptyKey() {
     unsigned EmptyKey = DenseMapInfo<unsigned>::getEmptyKey();
-    return ProfileBBID{EmptyKey, EmptyKey};
+    return UniqueBBID{EmptyKey, EmptyKey};
   }
-  static inline ProfileBBID getTombstoneKey() {
+  static inline UniqueBBID getTombstoneKey() {
     unsigned TombstoneKey = DenseMapInfo<unsigned>::getTombstoneKey();
-    return ProfileBBID{TombstoneKey, TombstoneKey};
+    return UniqueBBID{TombstoneKey, TombstoneKey};
   }
-  static unsigned getHashValue(const ProfileBBID &Val) {
+  static unsigned getHashValue(const UniqueBBID &Val) {
     std::pair<unsigned, unsigned> PairVal =
-        std::make_pair(Val.BBID, Val.CloneID);
+        std::make_pair(Val.BaseID, Val.CloneID);
     return DenseMapInfo<std::pair<unsigned, unsigned>>::getHashValue(PairVal);
   }
-  static bool isEqual(const ProfileBBID &LHS, const ProfileBBID &RHS) {
-    return DenseMapInfo<unsigned>::isEqual(LHS.BBID, RHS.BBID) &&
+  static bool isEqual(const UniqueBBID &LHS, const UniqueBBID &RHS) {
+    return DenseMapInfo<unsigned>::isEqual(LHS.BaseID, RHS.BaseID) &&
            DenseMapInfo<unsigned>::isEqual(LHS.CloneID, RHS.CloneID);
   }
 };
@@ -114,8 +102,11 @@ class BasicBlockSectionsProfileReader : public ImmutablePass {
   // function. If the first element is true and the second element is empty, it
   // means unique basic block sections are desired for all basic blocks of the
   // function.
-  std::pair<bool, FunctionPathAndClusterInfo>
-  getPathAndClusterInfoForFunction(StringRef FuncName) const;
+  std::pair<bool, SmallVector<BBClusterInfo>>
+    getClusterInfoForFunction(StringRef FuncName) const;
+ 
+  // Returns the path clonings for the given function.
+  SmallVector<SmallVector<unsigned>> getClonePathsForFunction(StringRef FuncName) const;
 
   // Initializes the FunctionNameToDIFilename map for the current module and
   // then reads the profile for the matching functions.
@@ -135,11 +126,11 @@ class BasicBlockSectionsProfileReader : public ImmutablePass {
         inconvertibleErrorCode());
   }
 
-  // Parses a `ProfileBBID` from `S`. `S` must be in the form "<bbid>"
+  // Parses a `UniqueBBID` from `S`. `S` must be in the form "<bbid>"
   // (representing an original block) or "<bbid>.<cloneid>" (representing a
   // cloned block) where bbid is a non-negative integer and cloneid is a
   // positive integer.
-  Expected<ProfileBBID> parseProfileBBID(StringRef S) const;
+  Expected<UniqueBBID> parseUniqueBBID(StringRef S) const;
 
   // Reads the basic block sections profile for functions in this module.
   Error ReadProfile();
diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
index 15c4fcd8399c181..4b5336fac33ea46 100644
--- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h
+++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -74,6 +74,13 @@ struct MBBSectionID {
   MBBSectionID(SectionType T) : Type(T), Number(0) {}
 };
 
+// This structure represents the information for a basic block.
+struct UniqueBBID {
+  unsigned BaseID;
+  // sections profile).
+  unsigned CloneID;
+};
+
 template <> struct ilist_traits<MachineInstr> {
 private:
   friend class MachineBasicBlock; // Set by the owning MachineBasicBlock.
@@ -180,7 +187,7 @@ class MachineBasicBlock
 
   /// Fixed unique ID assigned to this basic block upon creation. Used with
   /// basic block sections and basic block labels.
-  std::optional<unsigned> BBID;
+  std::optional<UniqueBBID> BBID;
 
   /// With basic block sections, this stores the Section ID of the basic block.
   MBBSectionID SectionID{0};
@@ -633,7 +640,7 @@ class MachineBasicBlock
 
   void setIsEndSection(bool V = true) { IsEndSection = V; }
 
-  std::optional<unsigned> getBBID() const { return BBID; }
+  std::optional<UniqueBBID> getBBID() const { return BBID; }
 
   /// Returns the section ID of this basic block.
   MBBSectionID getSectionID() const { return SectionID; }
@@ -645,7 +652,7 @@ class MachineBasicBlock
   }
 
   /// Sets the fixed BBID of this basic block.
-  void setBBID(unsigned V) {
+  void setBBID(const UniqueBBID &V) {
     assert(!BBID.has_value() && "Cannot change BBID.");
     BBID = V;
   }
@@ -753,7 +760,7 @@ class MachineBasicBlock
   ///
   /// This is useful when doing a partial clone of successors. Afterward, the
   /// probabilities may need to be normalized.
-  void copySuccessor(MachineBasicBlock *Orig, succ_iterator I);
+  void copySuccessor(const MachineBasicBlock *Orig, succ_iterator I);
 
   /// Split the old successor into old plus new and updates the probability
   /// info.
diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
index 8f1651c2958e591..5b2a9ec94168896 100644
--- a/llvm/include/llvm/CodeGen/MachineFunction.h
+++ b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -1005,8 +1005,9 @@ class LLVM_EXTERNAL_VISIBILITY MachineFunction {
   void deleteMachineInstr(MachineInstr *MI);
 
   /// CreateMachineBasicBlock - Allocate a new MachineBasicBlock. Use this
-  /// instead of `new MachineBasicBlock'.
-  MachineBasicBlock *CreateMachineBasicBlock(const BasicBlock *bb = nullptr);
+  /// instead of `new MachineBasicBlock'. Sets `MachineBasicBlock::BBID` if
+  /// basic-block-sections is enabled for the function.
+  MachineBasicBlock *CreateMachineBasicBlock(const BasicBlock *BB = nullptr, std::optional<UniqueBBID> BBID = std::nullopt);
 
   /// DeleteMachineBasicBlock - Delete the given MachineBasicBlock.
   void deleteMachineBasicBlock(MachineBasicBlock *MBB);
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index befa8a6eb9a27ce..adb28c5280d4224 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -64,6 +64,8 @@ namespace llvm {
   /// createBasicBlockSections Pass - This pass assigns sections to machine
   /// basic blocks and is enabled with -fbasic-block-sections.
   MachineFunctionPass *createBasicBlockSectionsPass();
+ 
+  MachineFunctionPass *createBasicBlockPathCloningPass();
 
   /// createMachineFunctionSplitterPass - This pass splits machine functions
   /// using profile information.
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index db653fff71ba95a..d59a706bc9c3869 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -55,6 +55,7 @@ void initializeAssignmentTrackingAnalysisPass(PassRegistry &);
 void initializeAssumeBuilderPassLegacyPassPass(PassRegistry &);
 void initializeAssumptionCacheTrackerPass(PassRegistry&);
 void initializeAtomicExpandPass(PassRegistry&);
+void initializeBasicBlockPathCloningPass(PassRegistry &);
 void initializeBasicBlockSectionsProfileReaderPass(PassRegistry &);
 void initializeBasicBlockSectionsPass(PassRegistry &);
 void initializeBarrierNoopPass(PassRegistry&);
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 97d2fe3426406f3..efee56543491c41 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1372,7 +1372,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
     if (BBAddrMapVersion > 1) {
       OutStreamer->AddComment("BB id");
       // Emit the BB ID for this basic block.
-      OutStreamer->emitULEB128IntValue(*MBB.getBBID());
+      OutStreamer->emitULEB128IntValue(MBB.getBBID()->BaseID);
     }
     // Emit the basic block offset relative to the end of the previous block.
     // This is zero unless the block is padded due to alignment.
@@ -1954,7 +1954,7 @@ void AsmPrinter::emitFunctionBody() {
       const double MBBRelFreq = MBFI.getBlockFreqRelativeToEntryBlock(&MBB);
       const double AbsMBBFreq = MBBRelFreq * EntryCount;
       *MBBProfileDumpFileOutput.get()
-          << MF->getName() << "," << MBB.getBBID() << "," << AbsMBBFreq << "\n";
+          << MF->getName() << "," << MBB.getBBID()->BaseID << "," << AbsMBBFreq << "\n";
     }
   }
 }
diff --git a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
new file mode 100644
index 000000000000000..9dcd8242e318562
--- /dev/null
+++ b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
@@ -0,0 +1,213 @@
+//===-- BasicBlockPathCloning.cpp ---=========-----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// BasicBlockPathCloning implementation.
+//
+// The purpose of this pass is to clone basic block paths based on information
+// provided by the -fbasic-block-sections=list option. 
+// Please refer to BasicBlockSectionsProfileReader.cpp to see a path cloning
+// example.
+//
+// ====================
+// This pass clones the machine basic blocks alongs the given paths and sets up
+// the CFG. It assigns BBIDs to the cloned blocks so that the `BasicBlockSections` pass can correctly map the cluster information to the blocks.
+// The cloned block's BBID will have the same BaseID as the original block, but will get a unique non-zero CloneID (original blocks all have zero CloneIDs).
+// This pass applies a path cloning if it satisfies the following conditions:
+//   1. All BBIDs in the path should be mapped to existing blocks.
+//   2. Each two consecutive BBIDs in the path must have a successor relationship in the CFG.
+//   3. The path should not include a block with indirect branches, except for the last block.
+// If a path does not satisfy all three conditions, it will be rejected, but the CloneIDs for its
+// (supposed to be cloned) blocks will be bypassed to make sure that the `BasicBlockSections` pass can map cluster info
+// correctly to the actually-cloned blocks.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/BasicBlockSectionUtils.h"
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Support/WithColor.h"
+#include "llvm/Target/TargetMachine.h"
+#include <optional>
+#include <sstream>
+
+using namespace llvm;
+
+namespace {
+
+// Clones the given block and assigns the given `CloneID` to its BBID. Copies the instructions into the new block and sets up its successors.
+MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock &OrigBB, unsigned CloneID) {
+  auto &MF = *OrigBB.getParent();
+  auto TII = MF.getSubtarget().getInstrInfo();
+  // Create the clone block and set its BBID based on the original block.
+  MachineBasicBlock *CloneBB = MF.CreateMachineBasicBlock(OrigBB.getBasicBlock(), UniqueBBID{OrigBB.getBBID()->BaseID, CloneID});
+  MF.push_back(CloneBB);
+ 
+  // Copy the instructions.
+  for (auto &I : OrigBB.instrs())
+    CloneBB->push_back(MF.CloneMachineInstr(&I));
+
+  // Add the successors of the original block as the new block's successors.
+  // We set the predecessor later.
+  for (auto SI = OrigBB.succ_begin(), SE = OrigBB.succ_end(); SI != SE; ++SI)
+    CloneBB->copySuccessor(&OrigBB, SI);
+
+  if (auto FT = OrigBB.getFallThrough(/*JumpToFallThrough=*/false)) {
+    // The original block has an implicit fall through.
+    // Insert an explicit unconditional jump from the cloned block to the
+    // fallthrough block.
+    TII->insertUnconditionalBranch(*CloneBB, FT, CloneBB->findBranchDebugLoc());
+  }
+  return CloneBB;
+}
+
+// Returns if we can legally apply all clonings specified in `ClonePaths`.
+bool IsValidCloning(const MachineFunction &MF,
+               const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
+               const SmallVector<unsigned> &ClonePath) {
+  const MachineBasicBlock *PrevBB = nullptr;
+  for (size_t I = 0; I < ClonePath.size(); ++I) {
+    unsigned BBID = ClonePath[I];
+    const MachineBasicBlock *PathBB = BBIDToBlock.lookup(BBID);
+    if (!PathBB) {
+      WithColor::warning() << "no block with id " << BBID << " in function "
+                           << MF.getName() << "\n";
+      return false;
+    }
+
+    if (PrevBB && !PrevBB->isSuccessor(PathBB)) {
+      WithColor::warning() << "block #" << BBID
+                           << " is not a successor of block #"
+                           << PrevBB->getBBID()->BaseID << " in function "
+                           << MF.getName() << "\n";
+      return false;
+    }
+
+    if (I != ClonePath.size() - 1 && !PathBB->empty() &&
+        PathBB->back().isIndirectBranch()) {
+      WithColor::warning()
+          << "block #" << BBID
+          << " has indirect branch and appears as the non-tail block of a "
+             "path in function "
+          << MF.getName() << "\n";
+      return false;
+    }
+    PrevBB = PathBB;
+  }
+  return true;
+}
+
+// Applies all clonings specified in `ClonePaths` to `MF`. Returns true
+// if any clonings have been applied.
+bool
+ApplyCloning(MachineFunction &MF,
+             const SmallVector<SmallVector<unsigned>> &ClonePaths) {
+  if (ClonePaths.empty()) return false;
+  bool AnyPathsCloned = false;
+  // Map from the final BB IDs to the `MachineBasicBlock`s.
+  DenseMap<unsigned, MachineBasicBlock *> BBIDToBlock;
+  for (auto &BB : MF)
+    BBIDToBlock.try_emplace(BB.getBBID()->BaseID, &BB);
+
+  DenseMap<unsigned, unsigned> NClonesForBBID;
+  auto TII = MF.getSubtarget().getInstrInfo();
+  for (const auto &ClonePath : ClonePaths) {
+    if (!IsValidCloning(MF, BBIDToBlock, ClonePath)) {
+      // We still need to increment the number of clones so we can map
+      // to the cluster info correctly.
+      for (unsigned BBID : ClonePath)
+        ++NClonesForBBID[BBID];
+      continue;
+    }
+    MachineBasicBlock *PrevBB = nullptr;
+    for (unsigned BBID : ClonePath) {
+      MachineBasicBlock *OrigBB = BBIDToBlock.at(BBID);
+      if (PrevBB == nullptr) {
+        if (auto FT = OrigBB->getFallThrough(/*JumpToFallThrough=*/false)) {
+          // Make fallthroughs explicit since we may change the layout.
+          TII->insertUnconditionalBranch(*OrigBB, FT,
+                                         OrigBB->findBranchDebugLoc());
+        }
+        PrevBB = OrigBB;
+        continue;
+      }
+      MachineBasicBlock *CloneBB = CloneMachineBasicBlock(*OrigBB, ++NClonesForBBID[BBID]);
+
+      for (auto &LiveOut : PrevBB->liveouts())
+       CloneBB->addLiveIn(LiveOut);
+
+      // Set up the previous block in the path to jump to the clone. This also transfers the successor/predecessor relationship between PrevBB and OrigBB to between PrevBB and CloneBB.
+      PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
+      PrevBB = CloneBB;
+    }
+    AnyPathsCloned = true;
+  }
+  return AnyPathsCloned;
+}
+} // end anonymous namespace
+
+
+namespace llvm {
+class BasicBlockPathCloning : public MachineFunctionPass {
+ public:
+  static char ID;
+
+  BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
+
+  BasicBlockPathCloning() : MachineFunctionPass(ID) {
+    initializeBasicBlockPathCloningPass(*PassRegistry::getPassRegistry());
+  }
+
+  StringRef getPassName() const override {
+    return "Basic Block Sections Analysis";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  /// Identify basic blocks that need separate sections and prepare to emit them
+  /// accordingly.
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+} // end anonymous namespace
+
+char BasicBlockPathCloning::ID = 0;
+INITIALIZE_PASS_BEGIN(
+    BasicBlockPathCloning, "bb-path-cloning",
+    "Prepares for basic block sections, by splitting functions "
+    "into clusters of basic blocks.",
+    false, false)
+INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReader)
+INITIALIZE_PASS_END(BasicBlockPathCloning, "bb-path-cloning",
+                    "Prepares for basic block sections, by splitting functions "
+                    "into clusters of basic blocks.",
+                    false, false)
+
+bool BasicBlockPathCloning::runOnMachineFunction(MachineFunction &MF) {
+  auto BBSectionsType = MF.getTarget().getBBSectionsType();
+  assert(BBSectionsType == BasicBlockSection::List &&
+         "BB Sections list not enabled!");
+  if (hasInstrProfHashMismatch(MF)) return false;
+
+  return ApplyCloning(MF, getAnalysis<BasicBlockSectionsProfileReader>()
+            .getClonePathsForFunction(MF.getName()));
+}
+
+void BasicBlockPathCloning::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.setPreservesAll();
+  AU.addRequired<BasicBlockSectionsProfileReader>();
+  MachineFunctionPass::getAnalysisUsage(AU);
+}
+
+MachineFunctionPass *llvm::createBasicBlockPathCloningPass() {
+  return new BasicBlockPathCloning();
+}
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 248cfca6b94fd77..4153fdcfbb25b8b 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -101,35 +101,8 @@ static cl::opt<bool> BBSectionsDetectSourceDrift(
 
 namespace {
 
-MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock *MBB) {
-  auto &MF = *MBB->getParent();
-  auto TII = MF.getSubtarget().getInstrInfo();
-
-  auto CloneBB = MF.CreateMachineBasicBlock(MBB->getBasicBlock());
-  MF.push_back(CloneBB);
-  // Copy the instructions.
-  for (auto &I : MBB->instrs())
-    CloneBB->push_back(MF.CloneMachineInstr(&I));
-
-  // Add the successors of the original block as the new block's successors.
-  for (auto SI = MBB->succ_begin(), SE = MBB->succ_end(); SI != SE; ++SI)
-    CloneBB->copySuccessor(MBB, SI);
-
-  if (auto FT = MBB->getFallThrough(/*JumpToFallThrough=*/false)) {
-    // The original block has an implicit fall through.
-    // Insert an explicit unconditional jump from the cloned block to the
-    // fallthrough block.
-    TII->insertUnconditionalBranch(*CloneBB, FT, CloneBB->findBranchDebugLoc());
-  }
-
-  for (auto &LiveIn : MBB->liveins())
-    CloneBB->addLiveIn(LiveIn);
-
-  return CloneBB;
-}
-
 class BasicBlockSections : public MachineFunctionPass {
-public:
+ public:
   static char ID;
 
   BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
@@ -204,12 +177,12 @@ updateBranches(MachineFunction &MF,
 // 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.
-// ClusterInfoByBBID represents the cluster information for basic blocks. It
+// FuncClusterInfo 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<unsigned>> &ClusterInfoByBBID) {
+    const DenseMap<UniqueBBID, BBClusterInfo> &FuncClusterInfo) {
   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
@@ -220,17 +193,17 @@ static void assignSections(
     // 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 ClusterInfoByBBID is empty).
+    // for every basic block in this function (if FuncClusterInfo is empty).
     if (MF.getTarget().getBBSectionsType() == llvm::BasicBlockSection::All ||
-        ClusterInfoByBBID.empty()) {
+        FuncClusterInfo.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 = ClusterInfoByBBID.find(*MBB.getBBID());
-      if (I != ClusterInfoByBBID.end()) {
+      auto I = FuncClusterInfo.find(*MBB.getBBID());
+      if (I != FuncClusterInfo.end()) {
         MBB.setSectionID(I->second.ClusterID);
       } else {
         // BB goes into the special cold section if it is not specified in the
@@ -293,12 +266,7 @@ void llvm::avoidZeroOffsetLandingPad(MachineFunction &MF) {
   }
 }
 
-// This checks if the source of this function has drifted since this binary was
-// profiled previously.  For now, we are piggy backing on what PGO does to
-// detect this with instrumented profiles.  PGO emits an hash of the IR and
-// checks if the hash has changed.  Advanced basic block layout is usually done
-// on top of PGO optimized binaries and hence this check works well in practice.
-static bool hasInstrProfHashMismatch(MachineFunction &MF) {
+bool llvm::hasInstrProfHashMismatch(MachineFunction &MF) {
   if (!BBSectionsDetectSourceDrift)
     return false;
 
@@ -314,136 +282,12 @@ static bool hasInstrProfHashMismatch(MachineFunction &MF) {
   return false;
 }
 
-// Returns if we can legally apply all clonings specified in `ClonePaths`.
-static bool
-IsValidCloning(const MachineFunction &MF,
-               const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
-               const SmallVector<unsigned> &ClonePath) {
-  const MachineBasicBlock *PrevBB = nullptr;
-  for (size_t I = 0; I < ClonePath.size(); ++I) {
-    unsigned BBID = ClonePath[I];
-    const MachineBasicBlock *PathBB = BBIDToBlock.lookup(BBID);
-    if (!PathBB) {
-      WithColor::warning() << "no block with id " << BBID << " in function "
-                           << MF.getName() << "\n";
-      return false;
-    }
-
-    if (PrevBB && !PrevBB->isSuccessor(PathBB)) {
-      WithColor::warning() << "block #" << BBID
-                           << " is not a successor of block #"
-                           << *PrevBB->getBBID() << " in function "
-                           << MF.getName() << "\n";
-      return false;
-    }
-
-    if (I != ClonePath.size() - 1 && !PathBB->empty() &&
-        PathBB->back().isIndirectBranch()) {
-      WithColor::warning()
-          << "block #" << BBID
-          << " has indirect branch and appears as the non-tail block of a "
-             "path in function "
-          << MF.getName() << "\n";
-      return false;
-    }
-    PrevBB = PathBB;
-  }
-  return true;
-}
-
-// Applies all clonings specified in `ClonePaths` to `MF` and returns a map
-// from `ProfileBBID`s of all clone blocks to their BBIDs (assigned by
-// `MachineFunction`).
-static DenseMap<ProfileBBID, unsigned>
-ApplyCloning(MachineFunction &MF,
-             const SmallVector<SmallVector<unsigned>> &ClonePaths) {
-  DenseMap<ProfileBBID, unsigned> CloneBBIDMap;
-  // Map from the final BB IDs to the `MachineBasicBlock`s.
-  DenseMap<unsigned, MachineBasicBlock *> BBIDToBlock;
-  for (auto &BB : MF)
-    BBIDToBlock.try_emplace(*BB.getBBID(), &BB);
-
-  DenseMap<unsigned, unsigned> NClonesForBBID;
-  auto TII = MF.getSubtarget().getInstrInfo();
-  for (const auto &ClonePath : ClonePaths) {
-    if (!IsValidCloning(MF, BBIDToBlock, ClonePath)) {
-      // We still need to increment the number of clones so we can map
-      // to the cluster info correctly.
-      for (unsigned BBID : ClonePath)
-        ++NClonesForBBID[BBID];
-      continue;
-    }
-    MachineBasicBlock *PrevBB = nullptr;
-    for (unsigned BBID : ClonePath) {
-      MachineBasicBlock *OrigBB = BBIDToBlock.at(BBID);
-      if (PrevBB == nullptr) {
-        if (auto FT = OrigBB->getFallThrough(/*JumpToFallThrough=*/false)) {
-          // Make fallthroughs explicit since we may change the layout.
-          TII->insertUnconditionalBranch(*OrigBB, FT,
-                                         OrigBB->findBranchDebugLoc());
-        }
-        PrevBB = OrigBB;
-        continue;
-      }
-
-      MachineBasicBlock *CloneBB = CloneMachineBasicBlock(OrigBB);
-
-      CloneBBIDMap.try_emplace({BBID, ++NClonesForBBID[BBID]},
-                               *CloneBB->getBBID());
-      // Set up the previous block in the path to jump to the clone.
-      PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
-      PrevBB = CloneBB;
-    }
-  }
-  return CloneBBIDMap;
-}
-
-// Processes the raw input profile `PathAndClusterInfo` for the given function
-// `MF` and returns the final profile (cluster information). Returns error if
-// `PathAndClusterInfo` is invalid and cannot be applied (e.g., invalid cloning
-// paths).
-//
-// Steps:
-//    1. Clones basic blocks based on `PathAndClusterInfo.ClonePaths` (if not
-//    empty) and
-//       updates the CFG accordingly.
-//    2. Creates and returns the cluster profile for the final blocks (original
-//       and cloned) based on `PathAndClusterInfo.ClusterInfo`.
-static DenseMap<unsigned, BBClusterInfo<unsigned>>
-ProcessProfile(MachineFunction &MF,
-               const FunctionPathAndClusterInfo &PathAndClusterInfo) {
-  // Apply the clonings and obtain the map from the input block ID of cloned
-  // blocks to their final BB IDs.
-  DenseMap<ProfileBBID, unsigned> CloneBBIDMap =
-      ApplyCloning(MF, PathAndClusterInfo.ClonePaths);
-
-  // Map from final BB IDs to their profile information.
-  DenseMap<unsigned, BBClusterInfo<unsigned>> ClusterInfoByBBID;
-  // This step creates all the necessary clones. It does not adjust the
-  // branches.
-  for (const BBClusterInfo<ProfileBBID> &P : PathAndClusterInfo.ClusterInfo) {
-    unsigned FinalBBID = P.BasicBlockID.BBID;
-    if (P.BasicBlockID.CloneID != 0) {
-      auto I = CloneBBIDMap.find(P.BasicBlockID);
-      // If there is no match, it means cloning was not applied for that block.
-      // We simply ignore these entries.
-      if (I == CloneBBIDMap.end())
-        continue;
-      FinalBBID = I->second;
-    }
-    ClusterInfoByBBID.try_emplace(
-        FinalBBID,
-        BBClusterInfo<unsigned>{FinalBBID, P.ClusterID, P.PositionInCluster});
-  }
-  return ClusterInfoByBBID;
-}
-
 bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   auto BBSectionsType = MF.getTarget().getBBSectionsType();
   assert(BBSectionsType != BasicBlockSection::None &&
          "BB Sections not enabled!");
 
-  // Check for source drift.  If the source has changed since the profiles
+  // Check for source drift. If the source has changed since the profiles
   // were obtained, optimizing basic blocks might be sub-optimal.
   // This only applies to BasicBlockSection::List as it creates
   // clusters of basic blocks using basic block ids. Source drift can
@@ -451,34 +295,30 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   // regards to performance.
   if (BBSectionsType == BasicBlockSection::List &&
       hasInstrProfHashMismatch(MF))
-    return true;
+    return false;
   // Renumber blocks before sorting them. This is useful for accessing the
   // original layout positions and finding the original fallthroughs.
   MF.RenumberBlocks();
 
   if (BBSectionsType == BasicBlockSection::Labels) {
     MF.setBBSectionsType(BBSectionsType);
-    return true;
+    return false;
   }
 
-  DenseMap<unsigned, BBClusterInfo<unsigned>> ClusterInfoByBBID;
+  DenseMap<UniqueBBID, BBClusterInfo> FuncClusterInfo;
   if (BBSectionsType == BasicBlockSection::List) {
-    auto [HasProfile, PathAndClusterInfo] =
+    auto [HasProfile, ClusterInfo] =
         getAnalysis<BasicBlockSectionsProfileReader>()
-            .getPathAndClusterInfoForFunction(MF.getName());
+            .getClusterInfoForFunction(MF.getName());
     if (!HasProfile)
-      return true;
-    // Empty ClusterInfo represents basic block sections for all blocks in this
-    // function.
-    if (!PathAndClusterInfo.ClusterInfo.empty()) {
-      ClusterInfoByBBID = ProcessProfile(MF, PathAndClusterInfo);
-      if (ClusterInfoByBBID.empty())
-        return true;
+      return false;
+    for (auto &BBClusterInfo: ClusterInfo) {
+      FuncClusterInfo.try_emplace(BBClusterInfo.BBID, BBClusterInfo);
     }
   }
 
   MF.setBBSectionsType(BBSectionsType);
-  assignSections(MF, ClusterInfoByBBID);
+  assignSections(MF, FuncClusterInfo);
 
   // We make sure that the cluster including the entry basic block precedes all
   // other clusters.
@@ -512,8 +352,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 ClusterInfoByBBID.lookup(*X.getBBID()).PositionInCluster <
-             ClusterInfoByBBID.lookup(*Y.getBBID()).PositionInCluster;
+      return FuncClusterInfo.lookup(*X.getBBID()).PositionInCluster <
+             FuncClusterInfo.lookup(*Y.getBBID()).PositionInCluster;
     return X.getNumber() < Y.getNumber();
   };
 
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index 650c2770ae3bee1..ea5e5e809fc7646 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -35,15 +35,15 @@ INITIALIZE_PASS(BasicBlockSectionsProfileReader, "bbsections-profile-reader",
                 "Reads and parses a basic block sections profile.", false,
                 false)
 
-Expected<ProfileBBID>
-BasicBlockSectionsProfileReader::parseProfileBBID(StringRef S) const {
+Expected<UniqueBBID>
+BasicBlockSectionsProfileReader::parseUniqueBBID(StringRef S) const {
   SmallVector<StringRef, 2> Parts;
   S.split(Parts, '.');
   if (Parts.size() > 2)
     return createProfileParseError(Twine("unable to parse basic block id: '") +
                                    S + "'");
-  unsigned long long BBID;
-  if (getAsUnsignedInteger(Parts[0], 10, BBID))
+  unsigned long long BaseBBID;
+  if (getAsUnsignedInteger(Parts[0], 10, BaseBBID))
     return createProfileParseError(
         Twine("unable to parse BB id: '" + Parts[0]) +
         "': unsigned integer expected");
@@ -51,21 +51,25 @@ BasicBlockSectionsProfileReader::parseProfileBBID(StringRef S) const {
   if (Parts.size() > 1 && getAsUnsignedInteger(Parts[1], 10, CloneID))
     return createProfileParseError(Twine("unable to parse clone id: '") +
                                    Parts[1] + "': unsigned integer expected");
-  return ProfileBBID{static_cast<unsigned>(BBID),
+  return UniqueBBID{static_cast<unsigned>(BaseBBID),
                      static_cast<unsigned>(CloneID)};
 }
 
 bool BasicBlockSectionsProfileReader::isFunctionHot(StringRef FuncName) const {
-  return getPathAndClusterInfoForFunction(FuncName).first;
+  return getClusterInfoForFunction(FuncName).first;
 }
 
-std::pair<bool, FunctionPathAndClusterInfo>
-BasicBlockSectionsProfileReader::getPathAndClusterInfoForFunction(
+std::pair<bool, SmallVector<BBClusterInfo>>
+BasicBlockSectionsProfileReader::getClusterInfoForFunction(
     StringRef FuncName) const {
   auto R = ProgramPathAndClusterInfo.find(getAliasName(FuncName));
   return R != ProgramPathAndClusterInfo.end()
-             ? std::pair(true, R->second)
-             : std::pair(false, FunctionPathAndClusterInfo());
+             ? std::pair(true, R->second.ClusterInfo)
+             : std::pair(false, SmallVector<BBClusterInfo>());
+}
+
+SmallVector<SmallVector<unsigned>> BasicBlockSectionsProfileReader::getClonePathsForFunction(StringRef FuncName) const {
+  return ProgramPathAndClusterInfo.lookup(getAliasName(FuncName)).ClonePaths;
 }
 
 // Reads the version 1 basic block sections profile. Profile for each function
@@ -133,7 +137,7 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
 
   // Temporary set to ensure every basic block ID appears once in the clusters
   // of a function.
-  DenseSet<ProfileBBID> FuncBBIDs;
+  DenseSet<UniqueBBID> FuncBBIDs;
 
   // Debug-info-based module filename for the current function. Empty string
   // means no filename.
@@ -199,7 +203,7 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       // Reset current cluster position.
       CurrentPosition = 0;
       for (auto BasicBlockIDStr : Values) {
-        auto BasicBlockID = parseProfileBBID(BasicBlockIDStr);
+        auto BasicBlockID = parseUniqueBBID(BasicBlockIDStr);
         if (!BasicBlockID)
           return BasicBlockID.takeError();
         if (!FuncBBIDs.insert(*BasicBlockID).second)
@@ -207,11 +211,11 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
               Twine("duplicate basic block id found '") + BasicBlockIDStr +
               "'");
 
-        if (!BasicBlockID->BBID && CurrentPosition)
+        if (!BasicBlockID->BaseID && CurrentPosition)
           return createProfileParseError(
               "entry BB (0) does not begin a cluster.");
 
-        FI->second.ClusterInfo.emplace_back(BBClusterInfo<ProfileBBID>{
+        FI->second.ClusterInfo.emplace_back(BBClusterInfo{
             *std::move(BasicBlockID), CurrentCluster, CurrentPosition++});
       }
       CurrentCluster++;
@@ -224,15 +228,15 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       SmallSet<unsigned, 5> BBsInPath;
       FI->second.ClonePaths.push_back({});
       for (size_t I = 0; I < Values.size(); ++I) {
-        auto BBIDStr = Values[I];
-        unsigned long long BBID = 0;
-        if (getAsUnsignedInteger(BBIDStr, 10, BBID))
+        auto BaseBBIDStr = Values[I];
+        unsigned long long BaseBBID = 0;
+        if (getAsUnsignedInteger(BaseBBIDStr, 10, BaseBBID))
           return createProfileParseError(Twine("unsigned integer expected: '") +
-                                         BBIDStr + "'");
-        if (I != 0 && !BBsInPath.insert(BBID).second)
+                                         BaseBBIDStr + "'");
+        if (I != 0 && !BBsInPath.insert(BaseBBID).second)
           return createProfileParseError(
-              Twine("duplicate cloned block in path: '") + BBIDStr + "'");
-        FI->second.ClonePaths.back().push_back(BBID);
+              Twine("duplicate cloned block in path: '") + BaseBBIDStr + "'");
+        FI->second.ClonePaths.back().push_back(BaseBBID);
       }
       continue;
     }
@@ -286,7 +290,7 @@ Error BasicBlockSectionsProfileReader::ReadV0Profile() {
               "entry BB (0) does not begin a cluster");
 
         FI->second.ClusterInfo.emplace_back(
-            BBClusterInfo<ProfileBBID>({{static_cast<unsigned>(BBID), 0},
+            BBClusterInfo({{static_cast<unsigned>(BBID), 0},
                                         CurrentCluster,
                                         CurrentPosition++}));
       }
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 389c70d04f17ba3..df2d1831ee5fdbf 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -46,6 +46,7 @@ add_llvm_component_library(LLVMCodeGen
   BranchRelaxation.cpp
   BreakFalseDeps.cpp
   BasicBlockSections.cpp
+  BasicBlockPathCloning.cpp
   BasicBlockSectionsProfileReader.cpp
   CalcSpillWeights.cpp
   CallBrPrepare.cpp
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index 6272b654b329539..79a95ee0d747a1c 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -20,6 +20,7 @@ using namespace llvm;
 void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeAssignmentTrackingAnalysisPass(Registry);
   initializeAtomicExpandPass(Registry);
+  initializeBasicBlockPathCloningPass(Registry);
   initializeBasicBlockSectionsPass(Registry);
   initializeBranchFolderPassPass(Registry);
   initializeBranchRelaxationPass(Registry);
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index 65280c65b68781e..bc3d5ec212faea6 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -500,7 +500,7 @@ class MIParser {
   bool parseAlignment(uint64_t &Alignment);
   bool parseAddrspace(unsigned &Addrspace);
   bool parseSectionID(std::optional<MBBSectionID> &SID);
-  bool parseBBID(std::optional<unsigned> &BBID);
+  bool parseBBID(std::optional<UniqueBBID> &BBID);
   bool parseCallFrameSize(unsigned &CallFrameSize);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
@@ -666,14 +666,20 @@ bool MIParser::parseSectionID(std::optional<MBBSectionID> &SID) {
 }
 
 // Parse Machine Basic Block ID.
-bool MIParser::parseBBID(std::optional<unsigned> &BBID) {
+bool MIParser::parseBBID(std::optional<UniqueBBID> &BBID) {
   assert(Token.is(MIToken::kw_bb_id));
   lex();
-  unsigned Value = 0;
-  if (getUnsigned(Value))
+  unsigned BaseID = 0;
+  unsigned CloneID = 0;
+  if (getUnsigned(BaseID))
     return error("Unknown BB ID");
-  BBID = Value;
   lex();
+  if (Token.is(MIToken::dot)) {
+    if (getUnsigned(CloneID))
+      return error("Uknown clone ID");
+    lex();
+  }
+  BBID = {BaseID, CloneID};
   return false;
 }
 
@@ -705,7 +711,7 @@ bool MIParser::parseBasicBlockDefinition(
   bool IsEHFuncletEntry = false;
   std::optional<MBBSectionID> SectionID;
   uint64_t Alignment = 0;
-  std::optional<unsigned> BBID;
+  std::optional<UniqueBBID> BBID;
   unsigned CallFrameSize = 0;
   BasicBlock *BB = nullptr;
   if (consumeIfPresent(MIToken::lparen)) {
diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
index 7d3d8b6fba1b7df..e605fbefbd75719 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -567,7 +567,9 @@ void MachineBasicBlock::printName(raw_ostream &os, unsigned printNameFlags,
     }
     if (getBBID().has_value()) {
       os << (hasAttributes ? ", " : " (");
-      os << "bb_id " << *getBBID();
+      os << "bb_id " << getBBID()->BaseID;
+      if (getBBID()->CloneID != 0)
+        os << "." << getBBID()->CloneID;
       hasAttributes = true;
     }
     if (CallFrameSize != 0) {
@@ -886,7 +888,7 @@ void MachineBasicBlock::replaceSuccessor(MachineBasicBlock *Old,
   removeSuccessor(OldI);
 }
 
-void MachineBasicBlock::copySuccessor(MachineBasicBlock *Orig,
+void MachineBasicBlock::copySuccessor(const MachineBasicBlock *Orig,
                                       succ_iterator I) {
   if (!Orig->Probs.empty())
     addSuccessor(*I, Orig->getSuccProbability(I));
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
index 1f14546a25b1c5d..9aac5a6507c34d0 100644
--- a/llvm/lib/CodeGen/MachineFunction.cpp
+++ b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -451,16 +451,16 @@ void MachineFunction::deleteMachineInstr(MachineInstr *MI) {
 /// Allocate a new MachineBasicBlock. Use this instead of
 /// `new MachineBasicBlock'.
 MachineBasicBlock *
-MachineFunction::CreateMachineBasicBlock(const BasicBlock *bb) {
+MachineFunction::CreateMachineBasicBlock(const BasicBlock *BB, std::optional<UniqueBBID> BBID) {
   MachineBasicBlock *MBB =
       new (BasicBlockRecycler.Allocate<MachineBasicBlock>(Allocator))
-          MachineBasicBlock(*this, bb);
+          MachineBasicBlock(*this, BB);
   // Set BBID for `-basic-block=sections=labels` and
   // `-basic-block-sections=list` to allow robust mapping of profiles to basic
   // blocks.
   if (Target.getBBSectionsType() == BasicBlockSection::Labels ||
       Target.getBBSectionsType() == BasicBlockSection::List)
-    MBB->setBBID(NextBBID++);
+    MBB->setBBID(BBID.has_value() ? *BBID : UniqueBBID{NextBBID++, 0});
   return MBB;
 }
 
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index e6ecbc9b03f7149..86fbc00d128da7a 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -255,6 +255,11 @@ static cl::opt<bool>
     GCEmptyBlocks("gc-empty-basic-blocks", cl::init(false), cl::Hidden,
                   cl::desc("Enable garbage-collecting empty basic blocks"));
 
+/// Enable basic block path cloning with basic-block-sections.
+static cl::opt<bool> EnableBasicBlockPathCloning(
+		"enable-basic-block-path-cloning", cl::Hidden,
+		cl::desc("Applies basic cloning clonings specified in the basic block sections profile."));
+
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
 /// i.e. -disable-mypass=false has no effect.
@@ -1267,6 +1272,8 @@ void TargetPassConfig::addMachinePasses() {
     if (TM->getBBSectionsType() == llvm::BasicBlockSection::List) {
       addPass(llvm::createBasicBlockSectionsProfileReaderPass(
           TM->getBBSectionsFuncListBuf()));
+      if (EnableBasicBlockPathCloning)
+        addPass(llvm::createBasicBlockPathCloningPass());
     }
     addPass(llvm::createBasicBlockSectionsPass());
   } else if (TM->Options.EnableMachineFunctionSplitter ||
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
index 1c383771df77b1d..2defbc1eadcc1fb 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
@@ -7,21 +7,21 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'f foo' >> %t1
 ; RUN: echo 'p 0 3 5' >> %t1
 ; RUN: echo 'c 0 3.1 5.1 1 2 3 4 5' >> %t1
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS,FOOCLONE
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS,FOOCLONE
 ; RUN: echo 'v1' > %t2
 ; RUN: echo 'f foo' >> %t2
 ; RUN: echo 'p 0 3 5' >> %t2
 ; RUN: echo 'p 1 3 4 5' >> %t2
 ; RUN: echo 'c 0 3.1 5.1' >> %t2
 ; RUN: echo 'c 1 3.2 4.1 5.2 2 3 4 5' >> %t2
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS,FOOCLONE
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS,FOOCLONE
 
 ;; Test failed application of path cloning.
 ; RUN: echo 'v1' > %t3
 ; RUN: echo 'f foo' >> %t3
 ; RUN: echo 'p 0 2 3' >> %t3
 ; RUN: echo 'c 0 2.1 3.1 1' >> %t3
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3 2> %t3.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t3 2> %t3.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH3-WARN < %t3.err
 ;; Test that valid clonings are applied correctly, even if invalid clonings exist.
 ; RUN: echo 'v1' > %t3_1
@@ -29,13 +29,13 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'p 0 2 3' >> %t3_1
 ; RUN: echo 'p 0 1 3' >> %t3_1
 ; RUN: echo 'c 0 1.1 3.2 2.1 3.1 1' >> %t3_1
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3_1 2> %t3_1.err | FileCheck %s --check-prefixes=PATH3_1,FOONOCLONE,FOOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t3_1 2> %t3_1.err | FileCheck %s --check-prefixes=PATH3_1,FOONOCLONE,FOOSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH3-WARN < %t3_1.err
 ; RUN: echo 'v1' > %t4
 ; RUN: echo 'f foo' >> %t4
 ; RUN: echo 'p 0 100' >> %t4
 ; RUN: echo 'c 0 100.1 1' >> %t4
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t4 2> %t4.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t4 2> %t4.err | FileCheck %s --check-prefixes=FOONOCLONE,FOOSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH4-WARN < %t4.err
 
 define void @foo(i1 %a, i1 %b, i1 %c, i1 %d) {
@@ -129,20 +129,20 @@ cold:
 ; RUN: echo 'f bar' >> %t5
 ; RUN: echo 'p 0 1' >> %t5
 ; RUN: echo 'c 0 1.1 2 1' >> %t5
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t5 | FileCheck %s --check-prefixes=PATH5,BARSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t5 | FileCheck %s --check-prefixes=PATH5,BARSECTIONS
 
 ;; Test failed application of path cloning for paths with indirect branches.
 ; RUN: echo 'v1' > %t6
 ; RUN: echo 'f bar' >> %t6
 ; RUN: echo 'p 0 1 2' >> %t6
 ; RUN: echo 'c 0 1.1 2.1 1' >> %t6
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t6 2> %t6.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t6 2> %t6.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t6.err
 ; RUN: echo 'v1' > %t7
 ; RUN: echo 'f bar' >> %t7
 ; RUN: echo 'p 1 2' >> %t7
 ; RUN: echo 'c 0 1 2.1' >> %t7
-; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t7 2> %t7.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t7 2> %t7.err | FileCheck %s --check-prefixes=BARNOCLONE,BARSECTIONS
 ; RUN: FileCheck %s --check-prefixes=PATH-INDIR-WARN < %t7.err
 
 

>From cbfc2f76886c4311960ccfb70dfba99c27850ae1 Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Wed, 18 Oct 2023 22:40:31 +0000
Subject: [PATCH 4/6] clang-format.

---
 .../CodeGen/BasicBlockSectionsProfileReader.h |  7 +-
 llvm/include/llvm/CodeGen/MachineFunction.h   |  4 +-
 llvm/include/llvm/CodeGen/Passes.h            |  2 +-
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |  3 +-
 llvm/lib/CodeGen/BasicBlockPathCloning.cpp    | 74 +++++++++++--------
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 10 +--
 .../BasicBlockSectionsProfileReader.cpp       | 10 ++-
 llvm/lib/CodeGen/MachineFunction.cpp          |  3 +-
 llvm/lib/CodeGen/TargetPassConfig.cpp         |  5 +-
 9 files changed, 68 insertions(+), 50 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index 60214db668a6d12..dfb8d5d9f2f5d33 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -103,10 +103,11 @@ class BasicBlockSectionsProfileReader : public ImmutablePass {
   // means unique basic block sections are desired for all basic blocks of the
   // function.
   std::pair<bool, SmallVector<BBClusterInfo>>
-    getClusterInfoForFunction(StringRef FuncName) const;
- 
+  getClusterInfoForFunction(StringRef FuncName) const;
+
   // Returns the path clonings for the given function.
-  SmallVector<SmallVector<unsigned>> getClonePathsForFunction(StringRef FuncName) const;
+  SmallVector<SmallVector<unsigned>>
+  getClonePathsForFunction(StringRef FuncName) const;
 
   // Initializes the FunctionNameToDIFilename map for the current module and
   // then reads the profile for the matching functions.
diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
index 5b2a9ec94168896..76a4f9d7eafa468 100644
--- a/llvm/include/llvm/CodeGen/MachineFunction.h
+++ b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -1007,7 +1007,9 @@ class LLVM_EXTERNAL_VISIBILITY MachineFunction {
   /// CreateMachineBasicBlock - Allocate a new MachineBasicBlock. Use this
   /// instead of `new MachineBasicBlock'. Sets `MachineBasicBlock::BBID` if
   /// basic-block-sections is enabled for the function.
-  MachineBasicBlock *CreateMachineBasicBlock(const BasicBlock *BB = nullptr, std::optional<UniqueBBID> BBID = std::nullopt);
+  MachineBasicBlock *
+  CreateMachineBasicBlock(const BasicBlock *BB = nullptr,
+                          std::optional<UniqueBBID> BBID = std::nullopt);
 
   /// DeleteMachineBasicBlock - Delete the given MachineBasicBlock.
   void deleteMachineBasicBlock(MachineBasicBlock *MBB);
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index adb28c5280d4224..7a0d0ea7b09e731 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -64,7 +64,7 @@ namespace llvm {
   /// createBasicBlockSections Pass - This pass assigns sections to machine
   /// basic blocks and is enabled with -fbasic-block-sections.
   MachineFunctionPass *createBasicBlockSectionsPass();
- 
+
   MachineFunctionPass *createBasicBlockPathCloningPass();
 
   /// createMachineFunctionSplitterPass - This pass splits machine functions
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index efee56543491c41..762ac52f2e87dbe 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1954,7 +1954,8 @@ void AsmPrinter::emitFunctionBody() {
       const double MBBRelFreq = MBFI.getBlockFreqRelativeToEntryBlock(&MBB);
       const double AbsMBBFreq = MBBRelFreq * EntryCount;
       *MBBProfileDumpFileOutput.get()
-          << MF->getName() << "," << MBB.getBBID()->BaseID << "," << AbsMBBFreq << "\n";
+          << MF->getName() << "," << MBB.getBBID()->BaseID << "," << AbsMBBFreq
+          << "\n";
     }
   }
 }
diff --git a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
index 9dcd8242e318562..8b146929c8be95d 100644
--- a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
+++ b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
@@ -9,21 +9,27 @@
 // BasicBlockPathCloning implementation.
 //
 // The purpose of this pass is to clone basic block paths based on information
-// provided by the -fbasic-block-sections=list option. 
+// provided by the -fbasic-block-sections=list option.
 // Please refer to BasicBlockSectionsProfileReader.cpp to see a path cloning
 // example.
 //
 // ====================
 // This pass clones the machine basic blocks alongs the given paths and sets up
-// the CFG. It assigns BBIDs to the cloned blocks so that the `BasicBlockSections` pass can correctly map the cluster information to the blocks.
-// The cloned block's BBID will have the same BaseID as the original block, but will get a unique non-zero CloneID (original blocks all have zero CloneIDs).
-// This pass applies a path cloning if it satisfies the following conditions:
+// the CFG. It assigns BBIDs to the cloned blocks so that the
+// `BasicBlockSections` pass can correctly map the cluster information to the
+// blocks. The cloned block's BBID will have the same BaseID as the original
+// block, but will get a unique non-zero CloneID (original blocks all have zero
+// CloneIDs). This pass applies a path cloning if it satisfies the following
+// conditions:
 //   1. All BBIDs in the path should be mapped to existing blocks.
-//   2. Each two consecutive BBIDs in the path must have a successor relationship in the CFG.
-//   3. The path should not include a block with indirect branches, except for the last block.
-// If a path does not satisfy all three conditions, it will be rejected, but the CloneIDs for its
-// (supposed to be cloned) blocks will be bypassed to make sure that the `BasicBlockSections` pass can map cluster info
-// correctly to the actually-cloned blocks.
+//   2. Each two consecutive BBIDs in the path must have a successor
+//   relationship in the CFG.
+//   3. The path should not include a block with indirect branches, except for
+//   the last block.
+// If a path does not satisfy all three conditions, it will be rejected, but the
+// CloneIDs for its (supposed to be cloned) blocks will be bypassed to make sure
+// that the `BasicBlockSections` pass can map cluster info correctly to the
+// actually-cloned blocks.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallVector.h"
@@ -44,14 +50,17 @@ using namespace llvm;
 
 namespace {
 
-// Clones the given block and assigns the given `CloneID` to its BBID. Copies the instructions into the new block and sets up its successors.
-MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock &OrigBB, unsigned CloneID) {
+// Clones the given block and assigns the given `CloneID` to its BBID. Copies
+// the instructions into the new block and sets up its successors.
+MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock &OrigBB,
+                                          unsigned CloneID) {
   auto &MF = *OrigBB.getParent();
   auto TII = MF.getSubtarget().getInstrInfo();
   // Create the clone block and set its BBID based on the original block.
-  MachineBasicBlock *CloneBB = MF.CreateMachineBasicBlock(OrigBB.getBasicBlock(), UniqueBBID{OrigBB.getBBID()->BaseID, CloneID});
+  MachineBasicBlock *CloneBB = MF.CreateMachineBasicBlock(
+      OrigBB.getBasicBlock(), UniqueBBID{OrigBB.getBBID()->BaseID, CloneID});
   MF.push_back(CloneBB);
- 
+
   // Copy the instructions.
   for (auto &I : OrigBB.instrs())
     CloneBB->push_back(MF.CloneMachineInstr(&I));
@@ -72,8 +81,8 @@ MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock &OrigBB, unsigned Cl
 
 // Returns if we can legally apply all clonings specified in `ClonePaths`.
 bool IsValidCloning(const MachineFunction &MF,
-               const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
-               const SmallVector<unsigned> &ClonePath) {
+                    const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
+                    const SmallVector<unsigned> &ClonePath) {
   const MachineBasicBlock *PrevBB = nullptr;
   for (size_t I = 0; I < ClonePath.size(); ++I) {
     unsigned BBID = ClonePath[I];
@@ -108,10 +117,10 @@ bool IsValidCloning(const MachineFunction &MF,
 
 // Applies all clonings specified in `ClonePaths` to `MF`. Returns true
 // if any clonings have been applied.
-bool
-ApplyCloning(MachineFunction &MF,
-             const SmallVector<SmallVector<unsigned>> &ClonePaths) {
-  if (ClonePaths.empty()) return false;
+bool ApplyCloning(MachineFunction &MF,
+                  const SmallVector<SmallVector<unsigned>> &ClonePaths) {
+  if (ClonePaths.empty())
+    return false;
   bool AnyPathsCloned = false;
   // Map from the final BB IDs to the `MachineBasicBlock`s.
   DenseMap<unsigned, MachineBasicBlock *> BBIDToBlock;
@@ -140,12 +149,15 @@ ApplyCloning(MachineFunction &MF,
         PrevBB = OrigBB;
         continue;
       }
-      MachineBasicBlock *CloneBB = CloneMachineBasicBlock(*OrigBB, ++NClonesForBBID[BBID]);
+      MachineBasicBlock *CloneBB =
+          CloneMachineBasicBlock(*OrigBB, ++NClonesForBBID[BBID]);
 
       for (auto &LiveOut : PrevBB->liveouts())
-       CloneBB->addLiveIn(LiveOut);
+        CloneBB->addLiveIn(LiveOut);
 
-      // Set up the previous block in the path to jump to the clone. This also transfers the successor/predecessor relationship between PrevBB and OrigBB to between PrevBB and CloneBB.
+      // Set up the previous block in the path to jump to the clone. This also
+      // transfers the successor/predecessor relationship between PrevBB and
+      // OrigBB to between PrevBB and CloneBB.
       PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
       PrevBB = CloneBB;
     }
@@ -155,10 +167,9 @@ ApplyCloning(MachineFunction &MF,
 }
 } // end anonymous namespace
 
-
 namespace llvm {
 class BasicBlockPathCloning : public MachineFunctionPass {
- public:
+public:
   static char ID;
 
   BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
@@ -168,7 +179,7 @@ class BasicBlockPathCloning : public MachineFunctionPass {
   }
 
   StringRef getPassName() const override {
-    return "Basic Block Sections Analysis";
+    return "Basic Block Path Cloning";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
@@ -178,28 +189,27 @@ class BasicBlockPathCloning : public MachineFunctionPass {
   bool runOnMachineFunction(MachineFunction &MF) override;
 };
 
-} // end anonymous namespace
+} // namespace llvm
 
 char BasicBlockPathCloning::ID = 0;
 INITIALIZE_PASS_BEGIN(
     BasicBlockPathCloning, "bb-path-cloning",
-    "Prepares for basic block sections, by splitting functions "
-    "into clusters of basic blocks.",
+    "Applies path clonings for the -basic-block-sections=list option",
     false, false)
 INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReader)
 INITIALIZE_PASS_END(BasicBlockPathCloning, "bb-path-cloning",
-                    "Prepares for basic block sections, by splitting functions "
-                    "into clusters of basic blocks.",
+    "Applies path clonings for the -basic-block-sections=list option",
                     false, false)
 
 bool BasicBlockPathCloning::runOnMachineFunction(MachineFunction &MF) {
   auto BBSectionsType = MF.getTarget().getBBSectionsType();
   assert(BBSectionsType == BasicBlockSection::List &&
          "BB Sections list not enabled!");
-  if (hasInstrProfHashMismatch(MF)) return false;
+  if (hasInstrProfHashMismatch(MF))
+    return false;
 
   return ApplyCloning(MF, getAnalysis<BasicBlockSectionsProfileReader>()
-            .getClonePathsForFunction(MF.getName()));
+                              .getClonePathsForFunction(MF.getName()));
 }
 
 void BasicBlockPathCloning::getAnalysisUsage(AnalysisUsage &AU) const {
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 4153fdcfbb25b8b..ebc7b3da0d97d48 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -102,7 +102,7 @@ static cl::opt<bool> BBSectionsDetectSourceDrift(
 namespace {
 
 class BasicBlockSections : public MachineFunctionPass {
- public:
+public:
   static char ID;
 
   BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
@@ -180,9 +180,9 @@ updateBranches(MachineFunction &MF,
 // FuncClusterInfo 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<UniqueBBID, BBClusterInfo> &FuncClusterInfo) {
+static void
+assignSections(MachineFunction &MF,
+               const DenseMap<UniqueBBID, BBClusterInfo> &FuncClusterInfo) {
   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
@@ -312,7 +312,7 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
             .getClusterInfoForFunction(MF.getName());
     if (!HasProfile)
       return false;
-    for (auto &BBClusterInfo: ClusterInfo) {
+    for (auto &BBClusterInfo : ClusterInfo) {
       FuncClusterInfo.try_emplace(BBClusterInfo.BBID, BBClusterInfo);
     }
   }
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index ea5e5e809fc7646..96662378a869316 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -52,7 +52,7 @@ BasicBlockSectionsProfileReader::parseUniqueBBID(StringRef S) const {
     return createProfileParseError(Twine("unable to parse clone id: '") +
                                    Parts[1] + "': unsigned integer expected");
   return UniqueBBID{static_cast<unsigned>(BaseBBID),
-                     static_cast<unsigned>(CloneID)};
+                    static_cast<unsigned>(CloneID)};
 }
 
 bool BasicBlockSectionsProfileReader::isFunctionHot(StringRef FuncName) const {
@@ -68,7 +68,9 @@ BasicBlockSectionsProfileReader::getClusterInfoForFunction(
              : std::pair(false, SmallVector<BBClusterInfo>());
 }
 
-SmallVector<SmallVector<unsigned>> BasicBlockSectionsProfileReader::getClonePathsForFunction(StringRef FuncName) const {
+SmallVector<SmallVector<unsigned>>
+BasicBlockSectionsProfileReader::getClonePathsForFunction(
+    StringRef FuncName) const {
   return ProgramPathAndClusterInfo.lookup(getAliasName(FuncName)).ClonePaths;
 }
 
@@ -291,8 +293,8 @@ Error BasicBlockSectionsProfileReader::ReadV0Profile() {
 
         FI->second.ClusterInfo.emplace_back(
             BBClusterInfo({{static_cast<unsigned>(BBID), 0},
-                                        CurrentCluster,
-                                        CurrentPosition++}));
+                           CurrentCluster,
+                           CurrentPosition++}));
       }
       CurrentCluster++;
     } else {
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
index 9aac5a6507c34d0..857eedebcf72d36 100644
--- a/llvm/lib/CodeGen/MachineFunction.cpp
+++ b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -451,7 +451,8 @@ void MachineFunction::deleteMachineInstr(MachineInstr *MI) {
 /// Allocate a new MachineBasicBlock. Use this instead of
 /// `new MachineBasicBlock'.
 MachineBasicBlock *
-MachineFunction::CreateMachineBasicBlock(const BasicBlock *BB, std::optional<UniqueBBID> BBID) {
+MachineFunction::CreateMachineBasicBlock(const BasicBlock *BB,
+                                         std::optional<UniqueBBID> BBID) {
   MachineBasicBlock *MBB =
       new (BasicBlockRecycler.Allocate<MachineBasicBlock>(Allocator))
           MachineBasicBlock(*this, BB);
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index 86fbc00d128da7a..704ece3bd653d05 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -257,8 +257,9 @@ static cl::opt<bool>
 
 /// Enable basic block path cloning with basic-block-sections.
 static cl::opt<bool> EnableBasicBlockPathCloning(
-		"enable-basic-block-path-cloning", cl::Hidden,
-		cl::desc("Applies basic cloning clonings specified in the basic block sections profile."));
+    "enable-basic-block-path-cloning", cl::Hidden,
+    cl::desc("Applies basic cloning clonings specified in the basic block "
+             "sections profile."));
 
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.

>From 9e2928701b7f61147af92b8f8f548e06583d132d Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Thu, 19 Oct 2023 21:19:36 +0000
Subject: [PATCH 5/6] Add clone id to MIR serialization.

---
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp        |  4 ++++
 llvm/lib/CodeGen/BasicBlockPathCloning.cpp        | 15 +++++++--------
 llvm/lib/CodeGen/MIRParser/MIParser.cpp           |  4 ++--
 llvm/lib/CodeGen/MachineBasicBlock.cpp            |  2 +-
 llvm/test/CodeGen/Generic/bb-profile-dump.ll      |  2 +-
 .../CodeGen/X86/basic-block-labels-mir-parse.mir  |  2 +-
 .../CodeGen/X86/basic-block-sections-cloning.ll   | 10 ++++++++++
 7 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 762ac52f2e87dbe..349f82fc612b2b0 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1372,6 +1372,10 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
     if (BBAddrMapVersion > 1) {
       OutStreamer->AddComment("BB id");
       // Emit the BB ID for this basic block.
+      // We only emit BaseID since CloneID is unset for
+      // basic-block-sections=labels.
+      // TODO: Emit the full BBID when labels and sections can be mixed
+      // together.
       OutStreamer->emitULEB128IntValue(MBB.getBBID()->BaseID);
     }
     // Emit the basic block offset relative to the end of the previous block.
diff --git a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
index 8b146929c8be95d..51a7caded3205aa 100644
--- a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
+++ b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
@@ -178,9 +178,7 @@ class BasicBlockPathCloning : public MachineFunctionPass {
     initializeBasicBlockPathCloningPass(*PassRegistry::getPassRegistry());
   }
 
-  StringRef getPassName() const override {
-    return "Basic Block Path Cloning";
-  }
+  StringRef getPassName() const override { return "Basic Block Path Cloning"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
@@ -194,12 +192,13 @@ class BasicBlockPathCloning : public MachineFunctionPass {
 char BasicBlockPathCloning::ID = 0;
 INITIALIZE_PASS_BEGIN(
     BasicBlockPathCloning, "bb-path-cloning",
-    "Applies path clonings for the -basic-block-sections=list option",
-    false, false)
+    "Applies path clonings for the -basic-block-sections=list option", false,
+    false)
 INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReader)
-INITIALIZE_PASS_END(BasicBlockPathCloning, "bb-path-cloning",
-    "Applies path clonings for the -basic-block-sections=list option",
-                    false, false)
+INITIALIZE_PASS_END(
+    BasicBlockPathCloning, "bb-path-cloning",
+    "Applies path clonings for the -basic-block-sections=list option", false,
+    false)
 
 bool BasicBlockPathCloning::runOnMachineFunction(MachineFunction &MF) {
   auto BBSectionsType = MF.getTarget().getBBSectionsType();
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index bc3d5ec212faea6..c01b34d6f490b0e 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -674,9 +674,9 @@ bool MIParser::parseBBID(std::optional<UniqueBBID> &BBID) {
   if (getUnsigned(BaseID))
     return error("Unknown BB ID");
   lex();
-  if (Token.is(MIToken::dot)) {
+  if (Token.is(MIToken::IntegerLiteral)) {
     if (getUnsigned(CloneID))
-      return error("Uknown clone ID");
+      return error("Unknown Clone ID");
     lex();
   }
   BBID = {BaseID, CloneID};
diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
index e605fbefbd75719..e6dfe3800e0f120 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -569,7 +569,7 @@ void MachineBasicBlock::printName(raw_ostream &os, unsigned printNameFlags,
       os << (hasAttributes ? ", " : " (");
       os << "bb_id " << getBBID()->BaseID;
       if (getBBID()->CloneID != 0)
-        os << "." << getBBID()->CloneID;
+        os << " " << getBBID()->CloneID;
       hasAttributes = true;
     }
     if (CallFrameSize != 0) {
diff --git a/llvm/test/CodeGen/Generic/bb-profile-dump.ll b/llvm/test/CodeGen/Generic/bb-profile-dump.ll
index 7391a6ee6f9128d..cae33635382a4b7 100644
--- a/llvm/test/CodeGen/Generic/bb-profile-dump.ll
+++ b/llvm/test/CodeGen/Generic/bb-profile-dump.ll
@@ -53,7 +53,7 @@ exit:
 ; Check that if we pass -mbb-profile-dump but don't set -basic-block-sections,
 ; we get an appropriate error message
 
-; RUN: not llc -mtriple=x86_64-linux-unknown -o /dev/null -mbb-profile-dump=- %s 2>&1 | FileCheck --check-prefix=NO-SECTIONS %s
+; RUN: not --crash llc -mtriple=x86_64-linux-unknown -o /dev/null -mbb-profile-dump=- %s 2>&1 | FileCheck --check-prefix=NO-SECTIONS %s
 
 ; NO-SECTIONS: <unknown>:0: error: Unable to find BB labels for MBB profile dump. -mbb-profile-dump must be called with -basic-block-sections=labels
 
diff --git a/llvm/test/CodeGen/X86/basic-block-labels-mir-parse.mir b/llvm/test/CodeGen/X86/basic-block-labels-mir-parse.mir
index 74a7bcf3ae82f1a..f11707c719895da 100644
--- a/llvm/test/CodeGen/X86/basic-block-labels-mir-parse.mir
+++ b/llvm/test/CodeGen/X86/basic-block-labels-mir-parse.mir
@@ -136,7 +136,7 @@ body:             |
   
     MOV32mi $rbp, 1, $noreg, -8, $noreg, 0 :: (store (s32) into %ir.2)
   
-  bb.3 (%ir-block.9, bb_id 3):
+  bb.3 (%ir-block.9, bb_id 3 2):
     renamable $eax = MOV32rm $rbp, 1, $noreg, -8, $noreg :: (load (s32) from %ir.2)
     $rbp = frame-destroy POP64r implicit-def $rsp, implicit $rsp
     frame-destroy CFI_INSTRUCTION def_cfa $rsp, 8
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
index 2defbc1eadcc1fb..b05961f5860664a 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-cloning.ll
@@ -8,6 +8,7 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'p 0 3 5' >> %t1
 ; RUN: echo 'c 0 3.1 5.1 1 2 3 4 5' >> %t1
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t1 | FileCheck %s --check-prefixes=PATH1,FOOSECTIONS,FOOCLONE
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t1 -stop-after=bb-path-cloning | FileCheck %s --check-prefix=PATH1MIR
 ; RUN: echo 'v1' > %t2
 ; RUN: echo 'f foo' >> %t2
 ; RUN: echo 'p 0 3 5' >> %t2
@@ -15,6 +16,7 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'c 0 3.1 5.1' >> %t2
 ; RUN: echo 'c 1 3.2 4.1 5.2 2 3 4 5' >> %t2
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t2 | FileCheck %s --check-prefixes=PATH2,FOOSECTIONS,FOOCLONE
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -enable-basic-block-path-cloning -basic-block-sections=%t2 -stop-after=bb-path-cloning | FileCheck %s --check-prefix=PATH2MIR
 
 ;; Test failed application of path cloning.
 ; RUN: echo 'v1' > %t3
@@ -67,6 +69,14 @@ cold:
   ret void
 }
 
+; PATH1MIR:    bb.7.b3 (bb_id 3 1):
+; PATH1MIR     bb.8.b5 (bb_id 5 1):
+; PATH2MIR:    bb.7.b3 (bb_id 3 1):
+; PATH2MIR:    bb.8.b5 (bb_id 5 1):
+; PATH2MIR:    bb.9.b3 (bb_id 3 2):
+; PATH2MIR:    bb.10.b4 (bb_id 4 1):
+; PATH2MIR:    bb.11.b5 (bb_id 5 2):
+
 ; FOOSECTIONS:   .section    .text.foo,"ax", at progbits
 ; FOOSECTIONS: foo:
 ; FOOSECTIONS: # %bb.0:        # %b0

>From 956a4eceb0914248d331e9b4d1e3101435f60a03 Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Thu, 19 Oct 2023 21:36:42 +0000
Subject: [PATCH 6/6] Address remaining comments.

---
 llvm/lib/CodeGen/BasicBlockPathCloning.cpp | 26 +++++++++++++---------
 llvm/lib/CodeGen/BasicBlockSections.cpp    |  2 --
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
index 51a7caded3205aa..add2b2e20455567 100644
--- a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
+++ b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
@@ -43,8 +43,6 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Target/TargetMachine.h"
-#include <optional>
-#include <sstream>
 
 using namespace llvm;
 
@@ -66,20 +64,23 @@ MachineBasicBlock *CloneMachineBasicBlock(MachineBasicBlock &OrigBB,
     CloneBB->push_back(MF.CloneMachineInstr(&I));
 
   // Add the successors of the original block as the new block's successors.
-  // We set the predecessor later.
+  // We set the predecessor after returning from this call.
   for (auto SI = OrigBB.succ_begin(), SE = OrigBB.succ_end(); SI != SE; ++SI)
     CloneBB->copySuccessor(&OrigBB, SI);
 
   if (auto FT = OrigBB.getFallThrough(/*JumpToFallThrough=*/false)) {
     // The original block has an implicit fall through.
     // Insert an explicit unconditional jump from the cloned block to the
-    // fallthrough block.
+    // fallthrough block. Technically, this is only needed for the last block
+    // of the path, but we do it for all clones for consistency.
     TII->insertUnconditionalBranch(*CloneBB, FT, CloneBB->findBranchDebugLoc());
   }
   return CloneBB;
 }
 
-// Returns if we can legally apply all clonings specified in `ClonePaths`.
+// Returns if we can legally apply the cloning represented by `ClonePath`.
+// `BBIDToBlock` contains the original basic blocks in function `MF` keyed by
+// their `BBID::BaseID`.
 bool IsValidCloning(const MachineFunction &MF,
                     const DenseMap<unsigned, MachineBasicBlock *> &BBIDToBlock,
                     const SmallVector<unsigned> &ClonePath) {
@@ -141,8 +142,10 @@ bool ApplyCloning(MachineFunction &MF,
     for (unsigned BBID : ClonePath) {
       MachineBasicBlock *OrigBB = BBIDToBlock.at(BBID);
       if (PrevBB == nullptr) {
+        // The first block in the path is not cloned. We only need to make it
+        // branch to the next cloned block in the path. Here, we make its
+        // fallthrough explicit so we can change it later.
         if (auto FT = OrigBB->getFallThrough(/*JumpToFallThrough=*/false)) {
-          // Make fallthroughs explicit since we may change the layout.
           TII->insertUnconditionalBranch(*OrigBB, FT,
                                          OrigBB->findBranchDebugLoc());
         }
@@ -152,13 +155,16 @@ bool ApplyCloning(MachineFunction &MF,
       MachineBasicBlock *CloneBB =
           CloneMachineBasicBlock(*OrigBB, ++NClonesForBBID[BBID]);
 
+      // Set up the previous block in the path to jump to the clone. This also
+      // transfers the successor/predecessor relationship of PrevBB and OrigBB
+      // to that of PrevBB and CloneBB.
+      PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
+
+      // CloneBB has a single predecessor. Therefore, its livein is the liveout
+      // of the predecessor block.
       for (auto &LiveOut : PrevBB->liveouts())
         CloneBB->addLiveIn(LiveOut);
 
-      // Set up the previous block in the path to jump to the clone. This also
-      // transfers the successor/predecessor relationship between PrevBB and
-      // OrigBB to between PrevBB and CloneBB.
-      PrevBB->ReplaceUsesOfBlockWith(OrigBB, CloneBB);
       PrevBB = CloneBB;
     }
     AnyPathsCloned = true;
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index ebc7b3da0d97d48..42997d2287d61d7 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -77,10 +77,8 @@
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
-#include "llvm/Support/WithColor.h"
 #include "llvm/Target/TargetMachine.h"
 #include <optional>
-#include <sstream>
 
 using namespace llvm;
 



More information about the llvm-commits mailing list