[llvm] [AArch64] Extend intra-block optimizer to handle TBZ/TBNZ (PR #217808)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Aug 20 20:36:05 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Mugundan S (MGN-GIT)
<details>
<summary>Changes</summary>
PR Description:
* The patch extends the Condition Optimizer to transform CMP/CMN + Bcc sequences that compare against an immediate value of zero into TBZ/TBNZ instructions (e.g., if (a < 0) and if (a >= 0)).
* Added a pattern-matching helper to identify CMP/CMN + Bcc sequences that can be converted to TBZ/TBNZ.
* The patch supports both 32-bit and 64-bit registers for the GE and LT condition codes.
Suggestions are always welcome, please free to add your ideas
cc: @<!-- -->davemgreen & @<!-- -->hussam-alhassan
---
Full diff: https://github.com/llvm/llvm-project/pull/217808.diff
2 Files Affected:
- (modified) llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp (+81-2)
- (added) llvm/test/CodeGen/AArch64/aarch64-condopt-test-bit-zero-branch.mir (+185)
``````````diff
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;
}
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/217808
More information about the llvm-commits
mailing list