[llvm] [BranchFolding] Drop undef flag when hoisting common code from successors (PR #205135)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto updated https://github.com/llvm/llvm-project/pull/205135

>From 83174ce211a7bad9dcba401e2b609ed0a734c03c Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Mon, 22 Jun 2026 18:26:00 +0200
Subject: [PATCH] [BranchFolding] Drop undef flag when hoisting common code
 from successors

Similarly to what already done during tail merging (4040c0f4ec135c18e723c1807ec0d1dbbb4cf3fa),
make sure the intersection of undef flags is taken while
hoisting common code from successors.

Fixes: https://github.com/llvm/llvm-project/issues/204549.
---
 llvm/lib/CodeGen/BranchFolding.cpp            | 23 ++++++++++++-------
 .../MIR/X86/branch-folder-drop-undef.mir      |  3 +--
 .../branch-folder-drop-undef-end-to-end.ll    |  2 +-
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 55f2dd430d6cb..2fdd766102a7f 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -790,6 +790,15 @@ bool BranchFolder::CreateCommonTailOnlyBlock(MachineBasicBlock *&PredBB,
   return true;
 }
 
+/// Ensure undef flag is preserved only when it is present in both instructions.
+static void mergeUndefFlag(MachineInstr &Merged, const MachineInstr &Other) {
+  for (unsigned I = 0, E = Merged.getNumOperands(); I != E; ++I) {
+    MachineOperand &MO = Merged.getOperand(I);
+    if (MO.isReg() && MO.isUndef() && !Other.getOperand(I).isUndef())
+      MO.setIsUndef(false);
+  }
+}
+
 static void
 mergeOperations(MachineBasicBlock::iterator MBBIStartPos,
                 MachineBasicBlock &MBBCommon) {
@@ -825,15 +834,9 @@ mergeOperations(MachineBasicBlock::iterator MBBIStartPos,
     // Merge MMOs from memory operations in the common block.
     if (MBBICommon->mayLoadOrStore())
       MBBICommon->cloneMergedMemRefs(*MBB->getParent(), {&*MBBICommon, &*MBBI});
+
     // Drop undef flags if they aren't present in all merged instructions.
-    for (unsigned I = 0, E = MBBICommon->getNumOperands(); I != E; ++I) {
-      MachineOperand &MO = MBBICommon->getOperand(I);
-      if (MO.isReg() && MO.isUndef()) {
-        const MachineOperand &OtherMO = MBBI->getOperand(I);
-        if (!OtherMO.isUndef())
-          MO.setIsUndef(false);
-      }
-    }
+    mergeUndefFlag(*MBBICommon, *MBBI);
 
     ++MBBI;
     ++MBBICommon;
@@ -2159,6 +2162,10 @@ bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
       assert(TI->isIdenticalTo(*FI, MachineInstr::CheckDefs) &&
              "Expected non-debug lockstep");
 
+      // Drop undef flag on the hoisted instruction if it was not present in
+      // both of the original ones.
+      mergeUndefFlag(*TI, *FI);
+
       // Merge debug locs on hoisted instructions.
       TI->setDebugLoc(
           DILocation::getMergedLocation(TI->getDebugLoc(), FI->getDebugLoc()));
diff --git a/llvm/test/CodeGen/MIR/X86/branch-folder-drop-undef.mir b/llvm/test/CodeGen/MIR/X86/branch-folder-drop-undef.mir
index d87a90f80d7d5..9bec8f904f284 100644
--- a/llvm/test/CodeGen/MIR/X86/branch-folder-drop-undef.mir
+++ b/llvm/test/CodeGen/MIR/X86/branch-folder-drop-undef.mir
@@ -26,8 +26,7 @@ body:             |
   ; CHECK-NEXT:   successors: %bb.3(0x40000000), %bb.2(0x40000000)
   ; CHECK-NEXT:   liveins: $edi, $xmm0, $xmm1, $xmm2
   ; CHECK-NEXT: {{  $}}
-  ; FIXME: This is a miscompilation issue, should intersect flags when hoisting common code.
-  ; CHECK-NEXT:   renamable $xmm1 = PUNPCKLBWrr killed renamable $xmm1, undef renamable $xmm0
+  ; CHECK-NEXT:   renamable $xmm1 = PUNPCKLBWrr killed renamable $xmm1, renamable $xmm0
   ; CHECK-NEXT:   TEST8ri renamable $dil, 1, implicit-def $eflags
   ; CHECK-NEXT:   JCC_1 %bb.3, 5, implicit killed $eflags
   ; CHECK-NEXT: {{  $}}
diff --git a/llvm/test/CodeGen/X86/branch-folder-drop-undef-end-to-end.ll b/llvm/test/CodeGen/X86/branch-folder-drop-undef-end-to-end.ll
index 03144940ba2a0..0f9d9048c0259 100644
--- a/llvm/test/CodeGen/X86/branch-folder-drop-undef-end-to-end.ll
+++ b/llvm/test/CodeGen/X86/branch-folder-drop-undef-end-to-end.ll
@@ -19,7 +19,7 @@ define void @test_drop_undef_when_hoisting(<4 x i8> %v0, i1 %c) {
 ; CHECK-NEXT:    jne .LBB0_5
 ; CHECK-NEXT:  .LBB0_1: # %loop
 ; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    punpcklbw {{.*#+}} xmm1 = xmm1[0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7]
+; CHECK-NEXT:    punpcklbw {{.*#+}} xmm1 = xmm1[0],xmm0[0],xmm1[1],xmm0[1],xmm1[2],xmm0[2],xmm1[3],xmm0[3],xmm1[4],xmm0[4],xmm1[5],xmm0[5],xmm1[6],xmm0[6],xmm1[7],xmm0[7]
 ; CHECK-NEXT:    testb $1, %dil
 ; CHECK-NEXT:    jne .LBB0_3
 ; CHECK-NEXT:  # %bb.2: # in Loop: Header=BB0_1 Depth=1