[llvm] [AArch64] Extend intra-block optimizer to handle TBZ/TBNZ (PR #217808)
Mugundan S via llvm-commits
llvm-commits at lists.llvm.org
Thu Aug 20 22:55:25 PDT 2026
https://github.com/MGN-GIT updated https://github.com/llvm/llvm-project/pull/217808
>From 2eb6ecacd2c7b034e6fa06aefc106afaab35a0a6 Mon Sep 17 00:00:00 2001
From: Mugundan S <137760120+MGN-GIT at users.noreply.github.com>
Date: Thu, 20 Aug 2026 17:08:59 +0530
Subject: [PATCH 1/6] [AArch64] Extend intra-block optimizer to handle TBZ/TBNZ
---
.../AArch64/AArch64ConditionOptimizer.cpp | 83 ++++++++++++++++++-
1 file changed, 81 insertions(+), 2 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
index e4164f091faf6..f9642e3d7144e 100644
--- a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
@@ -53,7 +53,6 @@
//
// See optimizeCrossBlock() and optimizeIntraBlock() for implementation details.
//
-// TODO: maybe handle TBNZ/TBZ the same way as CMP when used instead for "a < 0"
// TODO: For cross-block:
// - allow second branching to be anything if it doesn't require adjusting
//
@@ -144,6 +143,8 @@ class AArch64ConditionOptimizerImpl {
bool optimizeCrossBlock(MachineBasicBlock &HBB);
std::pair<MachineInstr *, AArch64CC::CondCode>
findCondConsumer(MachineBasicBlock *MBB);
+ std::tuple<MachineInstr *, MachineInstr *, AArch64CC::CondCode>
+ isCmpToTbzPattern(MachineBasicBlock &MBB);
};
class AArch64ConditionOptimizerLegacy : public MachineFunctionPass {
@@ -532,6 +533,33 @@ bool AArch64ConditionOptimizerImpl::commitPendingPair(
return Changed;
}
+std::tuple<MachineInstr *, MachineInstr *, AArch64CC::CondCode>
+AArch64ConditionOptimizerImpl::isCmpToTbzPattern(MachineBasicBlock &MBB) {
+ // Get the Bcc terminator for the given MBB.
+ MachineInstr *BrMI = getBccTerminator(&MBB);
+ if (!BrMI)
+ return {nullptr, nullptr, AArch64CC::Invalid};
+ // Get the CMP/CMN instruction consumed by the Bcc terminator.
+ MachineInstr *CmpMI = findAdjustableCmp(BrMI);
+ if (!CmpMI)
+ return {nullptr, nullptr, AArch64CC::Invalid};
+ // Check if the immediate operand is zero.
+ MachineOperand &Op = CmpMI->getOperand(2);
+ if (!Op.isImm() || Op.getImm() != 0)
+ return {nullptr, nullptr, AArch64CC::Invalid};
+ // Get the condition code used by Bcc.
+ int CCOpIdx =
+ AArch64InstrInfo::findCondCodeUseOperandIdxForBranchOrSelect(*BrMI);
+ if (CCOpIdx < 0)
+ return {nullptr, nullptr, AArch64CC::Invalid};
+ AArch64CC::CondCode CC =
+ (AArch64CC::CondCode)(int)BrMI->getOperand(CCOpIdx).getImm();
+ // Only GE and LT can be represented by a sign-bit test.
+ if (CC != AArch64CC::GE && CC != AArch64CC::LT)
+ return {nullptr, nullptr, AArch64CC::Invalid};
+ return {BrMI, CmpMI, CC};
+}
+
// This function transforms cmps and their consuming conditionals (CmpCondPairs)
// 1. Same direction: when both conditions are the same (e.g. GT/GT or LT/LT)
// and immediates differ by 1
@@ -607,7 +635,58 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
// consumer would be affected by any CMP adjustment we make.
if (!nzcvLivesOut(&MBB))
Changed |= commitPendingPair(PendingPair, PairsByReg);
-
+
+ // Check if there is any CMP/CMN + Bcc sequence which
+ // can be converted to TBZ/TBNZ.
+ auto [BrMI, CmpMI, CC] = isCmpToTbzPattern(MBB);
+
+ if (BrMI && CmpMI) {
+ Register Reg = CmpMI->getOperand(1).getReg();
+
+ MachineBasicBlock *TBB = nullptr;
+ MachineBasicBlock *FBB = nullptr;
+ SmallVector<MachineOperand, 4> Cond;
+
+ if (TII->analyzeBranch(MBB, TBB, FBB, Cond))
+ return false;
+ if (!TBB)
+ return false;
+
+ unsigned CmpOpc = CmpMI->getOpcode();
+ // Determine the new TBZ/TBNZ opcode and bit position.
+ unsigned TbzOpc;
+ unsigned Bit;
+
+ switch (CmpOpc) {
+ // 32-bit comparison
+ case AArch64::SUBSWri:
+ case AArch64::ADDSWri:
+ Bit = 31;
+ TbzOpc = CC == AArch64CC::GE ? AArch64::TBZW
+ : AArch64::TBNZW;
+ break;
+ // 64-bit comparison
+ case AArch64::SUBSXri:
+ case AArch64::ADDSXri:
+ Bit = 63;
+ TbzOpc = CC == AArch64CC::GE ? AArch64::TBZX
+ : AArch64::TBNZX;
+ break;
+ default:
+ return false;
+ }
+ // Build the new tbz/tbnz instruction
+ BuildMI(MBB, BrMI, BrMI->getDebugLoc(), TII->get(TbzOpc))
+ .addReg(Reg)
+ .addImm(Bit)
+ .addMBB(TBB);
+ // Remove the old Compare and Branch instructions
+ CmpMI->eraseFromParent();
+ BrMI->eraseFromParent();
+
+ Changed = true;
+ }
+
return Changed;
}
>From d2f8a0f974e1fa2ff09e6775b93d8e8fb1b573e5 Mon Sep 17 00:00:00 2001
From: Mugundan S <137760120+MGN-GIT at users.noreply.github.com>
Date: Fri, 21 Aug 2026 08:55:55 +0530
Subject: [PATCH 2/6] Add a Condition Optimizer test for TBZ/TBNZ branches
---
.../aarch64-condopt-test-bit-zero-branch.mir | 185 ++++++++++++++++++
1 file changed, 185 insertions(+)
create mode 100644 llvm/test/CodeGen/AArch64/aarch64-condopt-test-bit-zero-branch.mir
diff --git a/llvm/test/CodeGen/AArch64/aarch64-condopt-test-bit-zero-branch.mir b/llvm/test/CodeGen/AArch64/aarch64-condopt-test-bit-zero-branch.mir
new file mode 100644
index 0000000000000..23d36b0119865
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/aarch64-condopt-test-bit-zero-branch.mir
@@ -0,0 +1,185 @@
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass=aarch64-condopt -verify-machineinstrs %s -o - | FileCheck %s
+
+# Test validates CMP/CMN + Bcc -> TBZ/TBNZ transformations
+# when the immediate value being compared is zero and the
+# condition code is either GE or LT.
+
+---
+name: cmp_subs_zero_ge
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = SUBSWri $w0, 0, 0, implicit-def $nzcv
+ Bcc 10, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subs_zero_ge
+# CHECK: TBZW $w0, 31, %bb.1
+
+---
+name: cmp_subs_zero_lt
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = SUBSWri $w0, 0, 0, implicit-def $nzcv
+ Bcc 11, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subs_zero_lt
+# CHECK: TBNZW $w0, 31, %bb.1
+
+---
+name: cmp_subsx_zero_ge
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $x0
+
+ $x1 = SUBSXri $x0, 0, 0, implicit-def $nzcv
+ Bcc 10, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subsx_zero_ge
+# CHECK: TBZX $x0, 63, %bb.1
+
+---
+name: cmp_subsx_zero_lt
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $x0
+
+ $x1 = SUBSXri $x0, 0, 0, implicit-def $nzcv
+ Bcc 11, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subsx_zero_lt
+# CHECK: TBNZX $x0, 63, %bb.1
+
+---
+name: cmp_adds_zero_ge
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = ADDSWri $w0, 0, 0, implicit-def $nzcv
+ Bcc 10, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_adds_zero_ge
+# CHECK: TBZW $w0, 31, %bb.1
+
+---
+name: cmp_adds_zero_lt
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = ADDSWri $w0, 0, 0, implicit-def $nzcv
+ Bcc 11, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_adds_zero_lt
+# CHECK: TBNZW $w0, 31, %bb.1
+
+---
+name: cmp_addsx_zero_ge
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $x0
+
+ $x1 = ADDSXri $x0, 0, 0, implicit-def $nzcv
+ Bcc 10, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_addsx_zero_ge
+# CHECK: TBZX $x0, 63, %bb.1
+
+---
+name: cmp_addsx_zero_lt
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $x0
+
+ $x1 = ADDSXri $x0, 0, 0, implicit-def $nzcv
+ Bcc 11, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_addsx_zero_lt
+# CHECK: TBNZX $x0, 63, %bb.1
+
+---
+name: cmp_subs_nonzero
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = SUBSWri $w0, 1, 0, implicit-def $nzcv
+ Bcc 10, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subs_nonzero
+# CHECK: $w1 = SUBSWri $w0, 1, 0, implicit-def $nzcv
+# CHECK: Bcc 10, %bb.1, implicit $nzcv
+# CHECK-NOT: TBZW
+# CHECK-NOT: TBNZW
+# CHECK-NOT: TBZX
+# CHECK-NOT: TBNZX
+
+---
+name: cmp_subs_zero_eq
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w0
+
+ $w1 = SUBSWri $w0, 0, 0, implicit-def $nzcv
+ Bcc 0, %bb.1, implicit $nzcv
+
+ bb.1:
+ $w0 = MOVi32imm 0
+ RET_ReallyLR implicit $w0
+
+# CHECK-LABEL: name: cmp_subs_zero_eq
+# CHECK: $w1 = SUBSWri $w0, 0, 0, implicit-def $nzcv
+# CHECK: Bcc 0, %bb.1, implicit $nzcv
+# CHECK-NOT: TBZW
+# CHECK-NOT: TBNZW
+# CHECK-NOT: TBZX
+# CHECK-NOT: TBNZX
>From 69728e22cf3e0cafd9e944a8df7afd8465fa6a73 Mon Sep 17 00:00:00 2001
From: Mugundan S <137760120+MGN-GIT at users.noreply.github.com>
Date: Fri, 21 Aug 2026 09:13:27 +0530
Subject: [PATCH 3/6] Fix clang formatting
---
.../AArch64/AArch64ConditionOptimizer.cpp | 20 +++++++++----------
1 file changed, 9 insertions(+), 11 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
index f9642e3d7144e..5b4ad962963c5 100644
--- a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
@@ -635,18 +635,18 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
// consumer would be affected by any CMP adjustment we make.
if (!nzcvLivesOut(&MBB))
Changed |= commitPendingPair(PendingPair, PairsByReg);
-
+
// Check if there is any CMP/CMN + Bcc sequence which
// can be converted to TBZ/TBNZ.
auto [BrMI, CmpMI, CC] = isCmpToTbzPattern(MBB);
-
+
if (BrMI && CmpMI) {
Register Reg = CmpMI->getOperand(1).getReg();
-
+
MachineBasicBlock *TBB = nullptr;
MachineBasicBlock *FBB = nullptr;
SmallVector<MachineOperand, 4> Cond;
-
+
if (TII->analyzeBranch(MBB, TBB, FBB, Cond))
return false;
if (!TBB)
@@ -656,21 +656,19 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
// Determine the new TBZ/TBNZ opcode and bit position.
unsigned TbzOpc;
unsigned Bit;
-
+
switch (CmpOpc) {
// 32-bit comparison
case AArch64::SUBSWri:
case AArch64::ADDSWri:
Bit = 31;
- TbzOpc = CC == AArch64CC::GE ? AArch64::TBZW
- : AArch64::TBNZW;
+ TbzOpc = CC == AArch64CC::GE ? AArch64::TBZW : AArch64::TBNZW;
break;
// 64-bit comparison
case AArch64::SUBSXri:
case AArch64::ADDSXri:
Bit = 63;
- TbzOpc = CC == AArch64CC::GE ? AArch64::TBZX
- : AArch64::TBNZX;
+ bzOpc = CC == AArch64CC::GE ? AArch64::TBZX : AArch64::TBNZX;
break;
default:
return false;
@@ -683,10 +681,10 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
// Remove the old Compare and Branch instructions
CmpMI->eraseFromParent();
BrMI->eraseFromParent();
-
+
Changed = true;
}
-
+
return Changed;
}
>From 7c67ae73b68a131b36d38e5701036efdd8bc7965 Mon Sep 17 00:00:00 2001
From: Mugundan S <smugundan12a at gmail.com>
Date: Fri, 21 Aug 2026 09:16:29 +0530
Subject: [PATCH 4/6] [AArch64] Extend intra-block optimizer to handle
TBZ/TBNZ
---
llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
index 5b4ad962963c5..f6fc7923bbd36 100644
--- a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
@@ -651,7 +651,7 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
return false;
if (!TBB)
return false;
-
+
unsigned CmpOpc = CmpMI->getOpcode();
// Determine the new TBZ/TBNZ opcode and bit position.
unsigned TbzOpc;
>From 59baf30bdbbf6d2de20f06565923d654c453dfeb Mon Sep 17 00:00:00 2001
From: Mugundan S <smugundan12a at gmail.com>
Date: Fri, 21 Aug 2026 09:39:33 +0530
Subject: [PATCH 5/6] [AArch64] Extend intra-block optimizer to handle
TBZ/TBNZ
---
llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
index f6fc7923bbd36..ea40275922d75 100644
--- a/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp
@@ -668,7 +668,7 @@ bool AArch64ConditionOptimizerImpl::optimizeIntraBlock(MachineBasicBlock &MBB) {
case AArch64::SUBSXri:
case AArch64::ADDSXri:
Bit = 63;
- bzOpc = CC == AArch64CC::GE ? AArch64::TBZX : AArch64::TBNZX;
+ TbzOpc = CC == AArch64CC::GE ? AArch64::TBZX : AArch64::TBNZX;
break;
default:
return false;
>From 9d7fffde786342696594f54242526b49ba457e1f Mon Sep 17 00:00:00 2001
From: Mugundan S <smugundan12a at gmail.com>
Date: Fri, 21 Aug 2026 11:25:14 +0530
Subject: [PATCH 6/6] [AArch64] Update condopt test case to handle tbz/tbnz
optimization
---
llvm/test/CodeGen/AArch64/aarch64-condopt-unsigned.mir | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/llvm/test/CodeGen/AArch64/aarch64-condopt-unsigned.mir b/llvm/test/CodeGen/AArch64/aarch64-condopt-unsigned.mir
index 4269921e1b780..1985b3ed57dd9 100644
--- a/llvm/test/CodeGen/AArch64/aarch64-condopt-unsigned.mir
+++ b/llvm/test/CodeGen/AArch64/aarch64-condopt-unsigned.mir
@@ -385,8 +385,7 @@ body: |
; CHECK-NEXT: successors: %bb.2(0x80000000)
; CHECK-NEXT: liveins: $w0
; CHECK-NEXT: {{ $}}
- ; CHECK-NEXT: [[SUBSWri:%[0-9]+]]:gpr32 = SUBSWri [[COPY]], 0, 0, implicit-def $nzcv
- ; CHECK-NEXT: Bcc 10, %bb.2, implicit $nzcv
+ ; CHECK-NEXT: TBZW [[COPY]], 31, %bb.2
; CHECK-NEXT: B %bb.2
; CHECK-NEXT: {{ $}}
; CHECK-NEXT: bb.2:
More information about the llvm-commits
mailing list