[llvm] [PeepholeOpt] Erase optimized compare from LocalMIs earlier (PR #217848)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/217848

We need to drop the compare instruction that was optimized away from LocalMIs before the LocalMIs-based load folding optimization a few lines below. Addresses a regression from https://github.com/llvm/llvm-project/pull/194662.

Fixes github.com/llvm/llvm-project/issues/208746.

This is a simplified version of https://github.com/llvm/llvm-project/pull/211479.

>From 1cef9bff099be94c79b54fedd0e4a9214198c755 Mon Sep 17 00:00:00 2001
From: woruyu <1214539920 at qq.com>
Date: Thu, 23 Jul 2026 15:03:34 +0800
Subject: [PATCH 1/2] Avoid stale LocalMIs entries after compare optimization

---
 llvm/lib/CodeGen/PeepholeOptimizer.cpp        | 10 ++++-
 .../X86/peephole-compare-load-fold-ext.mir    | 40 +++++++++++++++++++
 2 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/peephole-compare-load-fold-ext.mir

diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index e7b0150ae3a72..a9e1af54e359f 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -959,8 +959,15 @@ bool PeepholeOptimizer::optimizeCmpInstr(
 
   // Attempt to optimize the comparison instruction.
   LLVM_DEBUG(dbgs() << "Attempting to optimize compare: " << MI);
-  if (!TII->optimizeCompareInstr(MI, SrcReg, SrcReg2, CmpMask, CmpValue, MRI))
+  // Stop tracking MI before optimizeCompareInstr may erase it. Any instruction
+  // created below could otherwise reuse MI's address and be confused with the
+  // erased instruction in LocalMIs.
+  LocalMIs.erase(&MI);
+  if (!TII->optimizeCompareInstr(MI, SrcReg, SrcReg2, CmpMask, CmpValue, MRI)) {
+    // MI was not erased, and is still part of the already visited region.
+    LocalMIs.insert(&MI);
     return false;
+  }
 
   LLVM_DEBUG(dbgs() << "  -> Successfully optimized compare!\n");
   ++NumCmps;
@@ -1843,7 +1850,6 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
       }
 
       if (MI->isCompare() && optimizeCmpInstr(*MI, MF, LocalMIs)) {
-        LocalMIs.erase(MI);
         Changed = true;
         continue;
       }
diff --git a/llvm/test/CodeGen/X86/peephole-compare-load-fold-ext.mir b/llvm/test/CodeGen/X86/peephole-compare-load-fold-ext.mir
new file mode 100644
index 0000000000000..ca0ec429902c9
--- /dev/null
+++ b/llvm/test/CodeGen/X86/peephole-compare-load-fold-ext.mir
@@ -0,0 +1,40 @@
+# RUN: llc -mtriple=x86_64-- -run-pass=peephole-opt -verify-machineinstrs %s -o - | FileCheck %s
+
+# When removing the second compare makes a load foldable into the first one,
+# make sure the folded compare remains marked as already visited. Otherwise,
+# the extension optimization may incorrectly replace its source with a subreg
+# of a result defined later in the block.
+
+---
+name: compare_load_fold_before_ext
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $edi, $rsi
+
+    ; CHECK-LABEL: name: compare_load_fold_before_ext
+    ; CHECK: [[SRC:%[0-9]+]]:gr32 = COPY $edi
+    ; CHECK-NEXT: [[PTR:%[0-9]+]]:gr64 = COPY $rsi
+    ; CHECK-NEXT: CMP32mr [[PTR]], 1, $noreg, 0, $noreg, [[SRC]], implicit-def $eflags
+    ; CHECK-NEXT: [[SETCC:%[0-9]+]]:gr8 = SETCCr 14, implicit $eflags
+    ; CHECK-NEXT: [[EXT:%[0-9]+]]:gr64 = MOVSX64rr32 [[SRC]]
+    ; CHECK-NEXT: NOOP implicit [[EXT]]
+    ; CHECK-NEXT: JCC_1 %bb.1, 14, implicit $eflags
+
+    %0:gr32 = COPY $edi
+    %1:gr64 = COPY $rsi
+    %2:gr32 = MOV32rm %1, 1, $noreg, 0, $noreg :: (load (s32))
+    %3:gr32 = SUB32rr %2, %0, implicit-def $eflags
+    %4:gr8 = SETCCr 14, implicit $eflags
+    %5:gr32 = SUB32rr %2, %0, implicit-def $eflags
+    %6:gr64 = MOVSX64rr32 %0
+    NOOP implicit %6
+    JCC_1 %bb.1, 14, implicit $eflags
+    JMP_1 %bb.2
+
+  bb.1:
+    RET 0
+
+  bb.2:
+    RET 0
+...

>From 0cbe358a899260a188dce78c69a29f11789cd6d0 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Fri, 21 Aug 2026 10:38:19 +0200
Subject: [PATCH 2/2] Simplify code

---
 llvm/lib/CodeGen/PeepholeOptimizer.cpp | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index a9e1af54e359f..ceab2f317e056 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -959,17 +959,11 @@ bool PeepholeOptimizer::optimizeCmpInstr(
 
   // Attempt to optimize the comparison instruction.
   LLVM_DEBUG(dbgs() << "Attempting to optimize compare: " << MI);
-  // Stop tracking MI before optimizeCompareInstr may erase it. Any instruction
-  // created below could otherwise reuse MI's address and be confused with the
-  // erased instruction in LocalMIs.
-  LocalMIs.erase(&MI);
-  if (!TII->optimizeCompareInstr(MI, SrcReg, SrcReg2, CmpMask, CmpValue, MRI)) {
-    // MI was not erased, and is still part of the already visited region.
-    LocalMIs.insert(&MI);
+  if (!TII->optimizeCompareInstr(MI, SrcReg, SrcReg2, CmpMask, CmpValue, MRI))
     return false;
-  }
 
   LLVM_DEBUG(dbgs() << "  -> Successfully optimized compare!\n");
+  LocalMIs.erase(&MI);
   ++NumCmps;
 
   // The eliminated compare may have been the extra use preventing a