[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 20:35:58 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/2] [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/2] 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
More information about the llvm-commits
mailing list