[llvm-branch-commits] [llvm] [CodeGen][NPM] Update dominator tree and post dominator tree consistently (PR #173485)

[Vikram Hegde via llvm-branch-commits]

https://github.com/vikramRH updated https://github.com/llvm/llvm-project/pull/173485

>From 001bd7d7b614842759cdcd7d8abfdd3116b4d20e Mon Sep 17 00:00:00 2001
From: vikhegde <vikram.hegde at amd.com>
Date: Mon, 22 Dec 2025 16:13:48 +0530
Subject: [PATCH] [CodeGen][NPM] Update dominator tree and post dominator tree
 consistently

---
 llvm/lib/CodeGen/BranchFolding.cpp           | 12 +++++++++++-
 llvm/lib/CodeGen/BranchRelaxation.cpp        | 11 ++++++++++-
 llvm/lib/CodeGen/MachineBlockPlacement.cpp   | 11 +++++++----
 llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp | 13 +++++++++++--
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 0b212fb0beb20..d1408807e6bc5 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -28,6 +28,7 @@
 #include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
@@ -35,6 +36,7 @@
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/MachineSizeOpts.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
@@ -135,12 +137,20 @@ PreservedAnalyses BranchFolderPass::run(MachineFunction &MF,
         "ProfileSummaryAnalysis is required for BranchFoldingPass", false);
 
   auto &MBFI = MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  auto *MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
+  auto *MPDT = MFAM.getCachedResult<MachinePostDominatorTreeAnalysis>(MF);
   MBFIWrapper MBBFreqInfo(MBFI);
   BranchFolder Folder(EnableTailMerge, /*CommonHoist=*/true, MBBFreqInfo, MBPI,
                       PSI);
   if (!Folder.OptimizeFunction(MF, MF.getSubtarget().getInstrInfo(),
-                               MF.getSubtarget().getRegisterInfo()))
+                               MF.getSubtarget().getRegisterInfo())) {
+    if (MDT)
+      MDT->updateBlockNumbers();
+    if (MPDT)
+      MPDT->updateBlockNumbers();
     return PreservedAnalyses::all();
+  }
+
   return getMachineFunctionPassPreservedAnalyses();
 }
 
diff --git a/llvm/lib/CodeGen/BranchRelaxation.cpp b/llvm/lib/CodeGen/BranchRelaxation.cpp
index fae952e888b4b..02a550f2e311a 100644
--- a/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -11,9 +11,11 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
@@ -769,8 +771,15 @@ bool BranchRelaxation::relaxBranchInstructions() {
 PreservedAnalyses
 BranchRelaxationPass::run(MachineFunction &MF,
                           MachineFunctionAnalysisManager &MFAM) {
-  if (!BranchRelaxation().run(MF))
+  auto *MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
+  auto *MPDT = MFAM.getCachedResult<MachinePostDominatorTreeAnalysis>(MF);
+  if (!BranchRelaxation().run(MF)) {
+    if (MDT)
+      MDT->updateBlockNumbers();
+    if (MPDT)
+      MPDT->updateBlockNumbers();
     return PreservedAnalyses::all();
+  }
 
   return getMachineFunctionPassPreservedAnalyses();
 }
diff --git a/llvm/lib/CodeGen/MachineBlockPlacement.cpp b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
index e9c75f0753f89..beff9a111b32c 100644
--- a/llvm/lib/CodeGen/MachineBlockPlacement.cpp
+++ b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
@@ -3541,19 +3541,22 @@ MachineBlockPlacementPass::run(MachineFunction &MF,
   auto *MPDT = MachineBlockPlacement::allowTailDupPlacement(MF)
                    ? &MFAM.getResult<MachinePostDominatorTreeAnalysis>(MF)
                    : nullptr;
+  auto *MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
   auto *PSI = MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF)
                   .getCachedResult<ProfileSummaryAnalysis>(
                       *MF.getFunction().getParent());
   if (!PSI)
     report_fatal_error("MachineBlockPlacement requires ProfileSummaryAnalysis",
                        false);
-
   MachineBlockPlacement MBP(MBPI, MLI, PSI, std::move(MBFI), MPDT,
                             AllowTailMerge);
-
-  if (!MBP.run(MF))
+  if (!MBP.run(MF)) {
+    if (MDT)
+      MDT->updateBlockNumbers();
+    if (MPDT)
+      MPDT->updateBlockNumbers();
     return PreservedAnalyses::all();
-
+  }
   return getMachineFunctionPassPreservedAnalyses();
 }
 
diff --git a/llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp b/llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp
index 8785968569d92..856b367a1f8f7 100644
--- a/llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp
+++ b/llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp
@@ -22,10 +22,11 @@
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/Support/BranchProbability.h"
-
 using namespace llvm;
 
 #define DEBUG_TYPE "si-pre-emit-peephole"
@@ -705,8 +706,16 @@ MachineInstrBuilder SIPreEmitPeephole::createUnpackedMI(MachineInstr &I,
 PreservedAnalyses
 llvm::SIPreEmitPeepholePass::run(MachineFunction &MF,
                                  MachineFunctionAnalysisManager &MFAM) {
-  if (!SIPreEmitPeephole().run(MF))
+  auto *MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
+  auto *MPDT = MFAM.getCachedResult<MachinePostDominatorTreeAnalysis>(MF);
+
+  if (!SIPreEmitPeephole().run(MF)) {
+    if (MDT)
+      MDT->updateBlockNumbers();
+    if (MPDT)
+      MPDT->updateBlockNumbers();
     return PreservedAnalyses::all();
+  }
 
   return getMachineFunctionPassPreservedAnalyses();
 }