[llvm] [CodeGen] Fold load into LZCNT/TZCNT/POPCNT after compare elimination in peephole optimizer (PR #194662)
via llvm-commits
llvm-commits at lists.llvm.org
Wed May 6 18:51:05 PDT 2026
https://github.com/Michael-Chen-NJU updated https://github.com/llvm/llvm-project/pull/194662
>From f1665985533cd57a78715b968ea10518837aaeb8 Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <221250121 at smail.nju.edu.cn>
Date: Tue, 28 Apr 2026 23:38:53 +0800
Subject: [PATCH 1/3] [CodeGen] Fold load into LZCNT/TZCNT/POPCNT after compare
elimination in peephole optimizer
---
llvm/lib/CodeGen/PeepholeOptimizer.cpp | 42 +++++++
llvm/test/CodeGen/X86/cttz.ll | 13 +--
llvm/test/CodeGen/X86/freeze-unary.ll | 3 +-
.../lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll | 104 ++++++++++++++++++
4 files changed, 153 insertions(+), 9 deletions(-)
create mode 100644 llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index 9365ea883eec9..2db6729add50b 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -1792,6 +1792,16 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
NAPhysToVirtMIs.clear();
}
+ // SrcReg must be captured before MI is erased inside optimizeCmpInstr.
+ Register CmpSrcReg;
+ if (MI->isCompare()) {
+ Register SrcReg2;
+ int64_t CmpMask, CmpValue;
+ if (!TII->analyzeCompare(*MI, CmpSrcReg, SrcReg2, CmpMask, CmpValue) ||
+ CmpSrcReg.isPhysical() || SrcReg2.isPhysical())
+ CmpSrcReg = Register();
+ }
+
if ((isUncoalescableCopy(*MI) &&
optimizeUncoalescableCopy(*MI, LocalMIs)) ||
(MI->isCompare() && optimizeCmpInstr(*MI)) ||
@@ -1799,6 +1809,38 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
// MI is deleted.
LocalMIs.erase(MI);
Changed = true;
+
+ // The eliminated compare may have been the extra use blocking load
+ // folding into the EFLAGS producer;
+ if (CmpSrcReg.isVirtual() && MRI->hasOneNonDBGUser(CmpSrcReg)) {
+ MachineInstr *FlagProducer =
+ MRI->use_nodbg_begin(CmpSrcReg)->getParent();
+ MachineInstr *LoadMI = MRI->getVRegDef(CmpSrcReg);
+ if (LocalMIs.count(FlagProducer) && LoadMI &&
+ LoadMI->canFoldAsLoad() && LoadMI->mayLoad() &&
+ LocalMIs.count(LoadMI)) {
+ Register FoldReg = CmpSrcReg;
+ MachineInstr *DefMI = nullptr;
+ MachineInstr *CopyMI = nullptr;
+ if (MachineInstr *FoldMI = TII->optimizeLoadInstr(
+ *FlagProducer, MRI, FoldReg, DefMI, CopyMI)) {
+ LLVM_DEBUG(dbgs() << "Replacing: " << *FlagProducer);
+ LLVM_DEBUG(dbgs() << " With: " << *FoldMI);
+ LocalMIs.erase(FlagProducer);
+ LocalMIs.erase(LoadMI);
+ LocalMIs.insert(FoldMI);
+ if (CopyMI)
+ LocalMIs.insert(CopyMI);
+ if (FlagProducer->shouldUpdateAdditionalCallInfo())
+ FlagProducer->getMF()->moveAdditionalCallInfo(FlagProducer,
+ FoldMI);
+ FlagProducer->eraseFromParent();
+ DefMI->eraseFromParent();
+ MRI->markUsesInDebugValueAsUndef(CmpSrcReg);
+ ++NumLoadFold;
+ }
+ }
+ }
continue;
}
diff --git a/llvm/test/CodeGen/X86/cttz.ll b/llvm/test/CodeGen/X86/cttz.ll
index a88fb96dd7c8c..9890450a1c195 100644
--- a/llvm/test/CodeGen/X86/cttz.ll
+++ b/llvm/test/CodeGen/X86/cttz.ll
@@ -379,13 +379,12 @@ define i64 @cttz_i64_zero_test(i64 %n) nounwind {
;
; X86-CMOV-LABEL: cttz_i64_zero_test:
; X86-CMOV: # %bb.0:
-; X86-CMOV-NEXT: movl {{[0-9]+}}(%esp), %eax
-; X86-CMOV-NEXT: bsfl {{[0-9]+}}(%esp), %ecx
-; X86-CMOV-NEXT: movl $32, %edx
-; X86-CMOV-NEXT: cmovnel %ecx, %edx
-; X86-CMOV-NEXT: addl $32, %edx
-; X86-CMOV-NEXT: bsfl %eax, %eax
-; X86-CMOV-NEXT: cmovel %edx, %eax
+; X86-CMOV-NEXT: bsfl {{[0-9]+}}(%esp), %eax
+; X86-CMOV-NEXT: movl $32, %ecx
+; X86-CMOV-NEXT: cmovnel %eax, %ecx
+; X86-CMOV-NEXT: addl $32, %ecx
+; X86-CMOV-NEXT: bsfl {{[0-9]+}}(%esp), %eax
+; X86-CMOV-NEXT: cmovel %ecx, %eax
; X86-CMOV-NEXT: xorl %edx, %edx
; X86-CMOV-NEXT: retl
;
diff --git a/llvm/test/CodeGen/X86/freeze-unary.ll b/llvm/test/CodeGen/X86/freeze-unary.ll
index bc9e29957c74a..4f4b2fbd77979 100644
--- a/llvm/test/CodeGen/X86/freeze-unary.ll
+++ b/llvm/test/CodeGen/X86/freeze-unary.ll
@@ -300,8 +300,7 @@ define i32 @freeze_cttz(i32 %a0) nounwind {
define i32 @freeze_cttz_undef(i32 %a0) nounwind {
; X86-LABEL: freeze_cttz_undef:
; X86: # %bb.0:
-; X86-NEXT: movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT: bsfl %eax, %ecx
+; X86-NEXT: bsfl {{[0-9]+}}(%esp), %ecx
; X86-NEXT: movl $32, %eax
; X86-NEXT: cmovnel %ecx, %eax
; X86-NEXT: retl
diff --git a/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll b/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
new file mode 100644
index 0000000000000..8a171859387a5
--- /dev/null
+++ b/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
@@ -0,0 +1,104 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+lzcnt,+bmi,+popcnt | FileCheck %s
+
+; Test that loads are folded into the memory form of LZCNT/TZCNT/POPCNT when
+; the loaded value is also used in an icmp-eq-zero/select sequence whose
+; condition is satisfied by the EFLAGS already produced by the bit-count
+; instruction (LZCNT/TZCNT set CF=1, POPCNT sets ZF=1 when input is zero).
+;
+; Before the fix, ISel chose the register form (rr) because the loaded value
+; had two uses: the bit-count instruction and the compare. The peephole
+; optimizer later removed the redundant TEST but missed folding the load into
+; the now single-use bit-count instruction.
+; See https://github.com/llvm/llvm-project/issues/164782
+
+; --- lzcnt: CF=1 when input is zero, cmovb takes the fallback ---
+
+define i32 @lzcnt32_load_cmov(ptr %p0, i32 %a1) {
+; CHECK-LABEL: lzcnt32_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: lzcntl (%rdi), %eax
+; CHECK-NEXT: cmovbl %esi, %eax
+; CHECK-NEXT: retq
+ %v = load i32, ptr %p0, align 4
+ %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
+ %iszero = icmp eq i32 %v, 0
+ %res = select i1 %iszero, i32 %a1, i32 %cnt
+ ret i32 %res
+}
+
+define i64 @lzcnt64_load_cmov(ptr %p0, i64 %a1) {
+; CHECK-LABEL: lzcnt64_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: lzcntq (%rdi), %rax
+; CHECK-NEXT: cmovbq %rsi, %rax
+; CHECK-NEXT: retq
+ %v = load i64, ptr %p0, align 8
+ %cnt = tail call range(i64 0, 65) i64 @llvm.ctlz.i64(i64 %v, i1 true)
+ %iszero = icmp eq i64 %v, 0
+ %res = select i1 %iszero, i64 %a1, i64 %cnt
+ ret i64 %res
+}
+
+; --- tzcnt: CF=1 when input is zero, cmovb takes the fallback ---
+
+define i32 @tzcnt32_load_cmov(ptr %p0, i32 %a1) {
+; CHECK-LABEL: tzcnt32_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: tzcntl (%rdi), %eax
+; CHECK-NEXT: cmovbl %esi, %eax
+; CHECK-NEXT: retq
+ %v = load i32, ptr %p0, align 4
+ %cnt = tail call range(i32 0, 33) i32 @llvm.cttz.i32(i32 %v, i1 true)
+ %iszero = icmp eq i32 %v, 0
+ %res = select i1 %iszero, i32 %a1, i32 %cnt
+ ret i32 %res
+}
+
+define i64 @tzcnt64_load_cmov(ptr %p0, i64 %a1) {
+; CHECK-LABEL: tzcnt64_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: tzcntq (%rdi), %rax
+; CHECK-NEXT: cmovbq %rsi, %rax
+; CHECK-NEXT: retq
+ %v = load i64, ptr %p0, align 8
+ %cnt = tail call range(i64 0, 65) i64 @llvm.cttz.i64(i64 %v, i1 true)
+ %iszero = icmp eq i64 %v, 0
+ %res = select i1 %iszero, i64 %a1, i64 %cnt
+ ret i64 %res
+}
+
+; --- popcnt: ZF=1 when input is zero, cmove takes the fallback ---
+
+define i32 @popcnt32_load_cmov(ptr %p0, i32 %a1) {
+; CHECK-LABEL: popcnt32_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: popcntl (%rdi), %eax
+; CHECK-NEXT: cmovel %esi, %eax
+; CHECK-NEXT: retq
+ %v = load i32, ptr %p0, align 4
+ %cnt = tail call i32 @llvm.ctpop.i32(i32 %v)
+ %iszero = icmp eq i32 %v, 0
+ %res = select i1 %iszero, i32 %a1, i32 %cnt
+ ret i32 %res
+}
+
+define i64 @popcnt64_load_cmov(ptr %p0, i64 %a1) {
+; CHECK-LABEL: popcnt64_load_cmov:
+; CHECK: # %bb.0:
+; CHECK-NEXT: popcntq (%rdi), %rax
+; CHECK-NEXT: cmoveq %rsi, %rax
+; CHECK-NEXT: retq
+ %v = load i64, ptr %p0, align 8
+ %cnt = tail call i64 @llvm.ctpop.i64(i64 %v)
+ %iszero = icmp eq i64 %v, 0
+ %res = select i1 %iszero, i64 %a1, i64 %cnt
+ ret i64 %res
+}
+
+declare i32 @llvm.ctlz.i32(i32, i1) nounwind readnone
+declare i64 @llvm.ctlz.i64(i64, i1) nounwind readnone
+declare i32 @llvm.cttz.i32(i32, i1) nounwind readnone
+declare i64 @llvm.cttz.i64(i64, i1) nounwind readnone
+declare i32 @llvm.ctpop.i32(i32) nounwind readnone
+declare i64 @llvm.ctpop.i64(i64) nounwind readnone
>From db7a7dfefdc503d6ac150e0f633ec209d5d93cba Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <221250121 at smail.nju.edu.cn>
Date: Wed, 29 Apr 2026 22:55:31 +0800
Subject: [PATCH 2/3] [CodeGen] Extract foldLoadInto helper to deduplicate
load-fold logic in PeepholeOptimizer
---
llvm/lib/CodeGen/PeepholeOptimizer.cpp | 82 +++++++++++---------------
1 file changed, 36 insertions(+), 46 deletions(-)
diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index 2db6729add50b..86560a235ecc9 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -494,6 +494,12 @@ class PeepholeOptimizer : private MachineFunction::Delegate {
bool isLoadFoldable(MachineInstr &MI,
SmallSet<Register, 16> &FoldAsLoadDefCandidates);
+ /// Try to fold the load defined by \p FoldReg into \p MI using
+ /// TII->optimizeLoadInstr. On success, updates LocalMIs, erases the old
+ /// instructions, and returns the replacement; returns nullptr otherwise.
+ MachineInstr *foldLoadInto(MachineInstr &MI, Register FoldReg,
+ SmallPtrSet<MachineInstr *, 16> &LocalMIs);
+
/// Check whether \p MI is understood by the register coalescer
/// but may require some rewriting.
static bool isCoalescableCopy(const MachineInstr &MI) {
@@ -1392,6 +1398,31 @@ bool PeepholeOptimizer::isLoadFoldable(
return false;
}
+MachineInstr *
+PeepholeOptimizer::foldLoadInto(MachineInstr &MI, Register FoldReg,
+ SmallPtrSet<MachineInstr *, 16> &LocalMIs) {
+ Register Reg = FoldReg;
+ MachineInstr *DefMI = nullptr;
+ MachineInstr *CopyMI = nullptr;
+ MachineInstr *FoldMI = TII->optimizeLoadInstr(MI, MRI, Reg, DefMI, CopyMI);
+ if (!FoldMI)
+ return nullptr;
+ LLVM_DEBUG(dbgs() << "Replacing: " << MI);
+ LLVM_DEBUG(dbgs() << " With: " << *FoldMI);
+ LocalMIs.erase(&MI);
+ LocalMIs.erase(DefMI);
+ LocalMIs.insert(FoldMI);
+ if (CopyMI)
+ LocalMIs.insert(CopyMI);
+ if (MI.shouldUpdateAdditionalCallInfo())
+ MI.getMF()->moveAdditionalCallInfo(&MI, FoldMI);
+ MI.eraseFromParent();
+ DefMI->eraseFromParent();
+ MRI->markUsesInDebugValueAsUndef(FoldReg);
+ ++NumLoadFold;
+ return FoldMI;
+}
+
bool PeepholeOptimizer::isMoveImmediate(
MachineInstr &MI, SmallSet<Register, 4> &ImmDefRegs,
DenseMap<Register, MachineInstr *> &ImmDefMIs) {
@@ -1811,35 +1842,15 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
Changed = true;
// The eliminated compare may have been the extra use blocking load
- // folding into the EFLAGS producer;
+ // folding into the EFLAGS producer.
if (CmpSrcReg.isVirtual() && MRI->hasOneNonDBGUser(CmpSrcReg)) {
MachineInstr *FlagProducer =
MRI->use_nodbg_begin(CmpSrcReg)->getParent();
MachineInstr *LoadMI = MRI->getVRegDef(CmpSrcReg);
if (LocalMIs.count(FlagProducer) && LoadMI &&
LoadMI->canFoldAsLoad() && LoadMI->mayLoad() &&
- LocalMIs.count(LoadMI)) {
- Register FoldReg = CmpSrcReg;
- MachineInstr *DefMI = nullptr;
- MachineInstr *CopyMI = nullptr;
- if (MachineInstr *FoldMI = TII->optimizeLoadInstr(
- *FlagProducer, MRI, FoldReg, DefMI, CopyMI)) {
- LLVM_DEBUG(dbgs() << "Replacing: " << *FlagProducer);
- LLVM_DEBUG(dbgs() << " With: " << *FoldMI);
- LocalMIs.erase(FlagProducer);
- LocalMIs.erase(LoadMI);
- LocalMIs.insert(FoldMI);
- if (CopyMI)
- LocalMIs.insert(CopyMI);
- if (FlagProducer->shouldUpdateAdditionalCallInfo())
- FlagProducer->getMF()->moveAdditionalCallInfo(FlagProducer,
- FoldMI);
- FlagProducer->eraseFromParent();
- DefMI->eraseFromParent();
- MRI->markUsesInDebugValueAsUndef(CmpSrcReg);
- ++NumLoadFold;
- }
- }
+ LocalMIs.count(LoadMI))
+ foldLoadInto(*FlagProducer, CmpSrcReg, LocalMIs);
}
continue;
}
@@ -1903,31 +1914,10 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
if (FoldAsLoadDefCandidates.count(FoldAsLoadDefReg)) {
// We need to fold load after optimizeCmpInstr, since
// optimizeCmpInstr can enable folding by converting SUB to CMP.
- // Save FoldAsLoadDefReg because optimizeLoadInstr() resets it and
- // we need it for markUsesInDebugValueAsUndef().
Register FoldedReg = FoldAsLoadDefReg;
- MachineInstr *DefMI = nullptr;
- MachineInstr *CopyMI = nullptr;
- if (MachineInstr *FoldMI = TII->optimizeLoadInstr(
- *MI, MRI, FoldAsLoadDefReg, DefMI, CopyMI)) {
- // Update LocalMIs since we replaced MI with FoldMI and deleted
- // DefMI.
- LLVM_DEBUG(dbgs() << "Replacing: " << *MI);
- LLVM_DEBUG(dbgs() << " With: " << *FoldMI);
- LocalMIs.erase(MI);
- LocalMIs.erase(DefMI);
- LocalMIs.insert(FoldMI);
- if (CopyMI)
- LocalMIs.insert(CopyMI);
- // Update the call info.
- if (MI->shouldUpdateAdditionalCallInfo())
- MI->getMF()->moveAdditionalCallInfo(MI, FoldMI);
- MI->eraseFromParent();
- DefMI->eraseFromParent();
- MRI->markUsesInDebugValueAsUndef(FoldedReg);
+ if (MachineInstr *FoldMI =
+ foldLoadInto(*MI, FoldAsLoadDefReg, LocalMIs)) {
FoldAsLoadDefCandidates.erase(FoldedReg);
- ++NumLoadFold;
-
// MI is replaced with FoldMI so we can continue trying to fold
Changed = true;
MI = FoldMI;
>From 47aff9e14e9544fde85199e4ce86fbccb604391f Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <221250121 at smail.nju.edu.cn>
Date: Thu, 7 May 2026 09:48:36 +0800
Subject: [PATCH 3/3] [CodeGen][test] Drop redundant declares and add negative
tests for lzcnt/tzcnt/popcnt load-fold
---
.../lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll | 37 ++++++++++++++++---
1 file changed, 31 insertions(+), 6 deletions(-)
diff --git a/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll b/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
index 8a171859387a5..41d87a75be7e0 100644
--- a/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
+++ b/llvm/test/CodeGen/X86/lzcnt-tzcnt-popcnt-load-fold-with-cmov.ll
@@ -96,9 +96,34 @@ define i64 @popcnt64_load_cmov(ptr %p0, i64 %a1) {
ret i64 %res
}
-declare i32 @llvm.ctlz.i32(i32, i1) nounwind readnone
-declare i64 @llvm.ctlz.i64(i64, i1) nounwind readnone
-declare i32 @llvm.cttz.i32(i32, i1) nounwind readnone
-declare i64 @llvm.cttz.i64(i64, i1) nounwind readnone
-declare i32 @llvm.ctpop.i32(i32) nounwind readnone
-declare i64 @llvm.ctpop.i64(i64) nounwind readnone
+; --- negative: icmp against non-zero constant; load must NOT be folded ---
+
+define i32 @lzcnt32_no_fold_nonzero_cmp(ptr %p0, i32 %a1) {
+; CHECK-LABEL: lzcnt32_no_fold_nonzero_cmp:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movl (%rdi), %ecx
+; CHECK-NEXT: lzcntl %ecx, %eax
+; CHECK-NEXT: cmpl $1, %ecx
+; CHECK-NEXT: cmovel %esi, %eax
+; CHECK-NEXT: retq
+ %v = load i32, ptr %p0, align 4
+ %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
+ %isone = icmp eq i32 %v, 1
+ %res = select i1 %isone, i32 %a1, i32 %cnt
+ ret i32 %res
+}
+
+define i32 @lzcnt32_no_fold_var_cmp(ptr %p0, i32 %a1, i32 %a2) {
+; CHECK-LABEL: lzcnt32_no_fold_var_cmp:
+; CHECK: # %bb.0:
+; CHECK-NEXT: movl (%rdi), %ecx
+; CHECK-NEXT: lzcntl %ecx, %eax
+; CHECK-NEXT: cmpl %edx, %ecx
+; CHECK-NEXT: cmovel %esi, %eax
+; CHECK-NEXT: retq
+ %v = load i32, ptr %p0, align 4
+ %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
+ %eq = icmp eq i32 %v, %a2
+ %res = select i1 %eq, i32 %a1, i32 %cnt
+ ret i32 %res
+}
More information about the llvm-commits
mailing list