1//===-- AArch64CondBrTuning.cpp --- Conditional branch tuning for AArch64 -===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8/// \file
9/// This file contains a pass that transforms CBZ/CBNZ/TBZ/TBNZ instructions
10/// into a conditional branch (B.cond), when the NZCV flags can be set for
11/// "free". This is preferred on targets that have more flexibility when
12/// scheduling B.cond instructions as compared to CBZ/CBNZ/TBZ/TBNZ (assuming
13/// all other variables are equal). This can also reduce register pressure.
14///
15/// A few examples:
16///
17/// 1) add w8, w0, w1 -> cmn w0, w1 ; CMN is an alias of ADDS.
18/// cbz w8, .LBB_2 -> b.eq .LBB0_2
19///
20/// 2) add w8, w0, w1 -> adds w8, w0, w1 ; w8 has multiple uses.
21/// cbz w8, .LBB1_2 -> b.eq .LBB1_2
22///
23/// 3) sub w8, w0, w1 -> subs w8, w0, w1 ; w8 has multiple uses.
24/// tbz w8, #31, .LBB6_2 -> b.pl .LBB6_2
25///
26//===----------------------------------------------------------------------===//
27
28#include "AArch64.h"
29#include "AArch64Subtarget.h"
30#include "llvm/CodeGen/MachineFunction.h"
31#include "llvm/CodeGen/MachineFunctionPass.h"
32#include "llvm/CodeGen/MachineInstrBuilder.h"
33#include "llvm/CodeGen/MachineRegisterInfo.h"
34#include "llvm/CodeGen/Passes.h"
35#include "llvm/CodeGen/RegisterClassInfo.h"
36#include "llvm/CodeGen/TargetInstrInfo.h"
37#include "llvm/CodeGen/TargetRegisterInfo.h"
38#include "llvm/CodeGen/TargetSubtargetInfo.h"
39#include "llvm/Support/Debug.h"
40#include "llvm/Support/raw_ostream.h"
41
42using namespace llvm;
43
44#define DEBUG_TYPE "aarch64-cond-br-tuning"
45#define AARCH64_CONDBR_TUNING_NAME "AArch64 Conditional Branch Tuning"
46
47namespace {
48class AArch64CondBrTuning : public MachineFunctionPass {
49 const AArch64InstrInfo *TII;
50 const TargetRegisterInfo *TRI;
51
52 MachineRegisterInfo *MRI;
53
54public:
55 static char ID;
56 AArch64CondBrTuning() : MachineFunctionPass(ID) {}
57 void getAnalysisUsage(AnalysisUsage &AU) const override;
58 bool runOnMachineFunction(MachineFunction &MF) override;
59 StringRef getPassName() const override { return AARCH64_CONDBR_TUNING_NAME; }
60
61private:
62 MachineInstr *getOperandDef(const MachineOperand &MO);
63 MachineInstr *tryConvertToFlagSetting(MachineInstr &MI, bool IsFlagSetting,
64 bool Is64Bit);
65 MachineInstr *convertToCondBr(MachineInstr &MI);
66 bool tryToTuneBranch(MachineInstr &MI, MachineInstr &DefMI);
67};
68} // end anonymous namespace
69
70char AArch64CondBrTuning::ID = 0;
71
72INITIALIZE_PASS(AArch64CondBrTuning, "aarch64-cond-br-tuning",
73 AARCH64_CONDBR_TUNING_NAME, false, false)
74
75void AArch64CondBrTuning::getAnalysisUsage(AnalysisUsage &AU) const {
76 AU.setPreservesCFG();
77 MachineFunctionPass::getAnalysisUsage(AU);
78}
79
80MachineInstr *AArch64CondBrTuning::getOperandDef(const MachineOperand &MO) {
81 if (!MO.getReg().isVirtual())
82 return nullptr;
83 return MRI->getUniqueVRegDef(Reg: MO.getReg());
84}
85
86MachineInstr *AArch64CondBrTuning::tryConvertToFlagSetting(MachineInstr &MI,
87 bool IsFlagSetting,
88 bool Is64Bit) {
89 // If the instruction has a frame index operand, we can't safely convert it
90 // to a flag-setting form, because it can be expanded later into multiple
91 // instructions, which don't all have flag-setting forms (e.g. ADDVL).
92 if (any_of(Range: MI.operands(), P: [](const MachineOperand &Op) { return Op.isFI(); }))
93 return nullptr;
94
95 // If this is already the flag setting version of the instruction (e.g., SUBS)
96 // just make sure the implicit-def of NZCV isn't marked dead.
97 if (IsFlagSetting) {
98 for (MachineOperand &MO : MI.implicit_operands())
99 if (MO.isReg() && MO.isDead() && MO.getReg() == AArch64::NZCV)
100 MO.setIsDead(false);
101 return &MI;
102 }
103 unsigned NewOpc = TII->convertToFlagSettingOpc(Opc: MI.getOpcode());
104 Register NewDestReg = MI.getOperand(i: 0).getReg();
105 if (MRI->hasOneNonDBGUse(RegNo: MI.getOperand(i: 0).getReg()))
106 NewDestReg = Is64Bit ? AArch64::XZR : AArch64::WZR;
107
108 MachineInstrBuilder MIB = BuildMI(BB&: *MI.getParent(), I&: MI, MIMD: MI.getDebugLoc(),
109 MCID: TII->get(Opcode: NewOpc), DestReg: NewDestReg);
110
111 // If the MI has a debug instruction number, preserve that in the new Machine
112 // Instruction that is created.
113 if (MI.peekDebugInstrNum() != 0)
114 MIB->setDebugInstrNum(MI.peekDebugInstrNum());
115
116 for (const MachineOperand &MO : llvm::drop_begin(RangeOrContainer: MI.operands()))
117 MIB.add(MO);
118
119 return MIB;
120}
121
122MachineInstr *AArch64CondBrTuning::convertToCondBr(MachineInstr &MI) {
123 AArch64CC::CondCode CC;
124 MachineBasicBlock *TargetMBB = TII->getBranchDestBlock(MI);
125 switch (MI.getOpcode()) {
126 default:
127 llvm_unreachable("Unexpected opcode!");
128
129 case AArch64::CBZW:
130 case AArch64::CBZX:
131 CC = AArch64CC::EQ;
132 break;
133 case AArch64::CBNZW:
134 case AArch64::CBNZX:
135 CC = AArch64CC::NE;
136 break;
137 case AArch64::TBZW:
138 case AArch64::TBZX:
139 CC = AArch64CC::PL;
140 break;
141 case AArch64::TBNZW:
142 case AArch64::TBNZX:
143 CC = AArch64CC::MI;
144 break;
145 }
146 return BuildMI(BB&: *MI.getParent(), I&: MI, MIMD: MI.getDebugLoc(), MCID: TII->get(Opcode: AArch64::Bcc))
147 .addImm(Val: CC)
148 .addMBB(MBB: TargetMBB);
149}
150
151bool AArch64CondBrTuning::tryToTuneBranch(MachineInstr &MI,
152 MachineInstr &DefMI) {
153 // We don't want NZCV bits live across blocks.
154 if (MI.getParent() != DefMI.getParent())
155 return false;
156
157 bool IsFlagSetting = true;
158 unsigned MIOpc = MI.getOpcode();
159 MachineInstr *NewCmp = nullptr, *NewBr = nullptr;
160 switch (DefMI.getOpcode()) {
161 default:
162 return false;
163 case AArch64::ADDWri:
164 case AArch64::ADDWrr:
165 case AArch64::ADDWrs:
166 case AArch64::ADDWrx:
167 case AArch64::ANDWri:
168 case AArch64::ANDWrr:
169 case AArch64::ANDWrs:
170 case AArch64::BICWrr:
171 case AArch64::BICWrs:
172 case AArch64::SUBWri:
173 case AArch64::SUBWrr:
174 case AArch64::SUBWrs:
175 case AArch64::SUBWrx:
176 IsFlagSetting = false;
177 [[fallthrough]];
178 case AArch64::ADDSWri:
179 case AArch64::ADDSWrr:
180 case AArch64::ADDSWrs:
181 case AArch64::ADDSWrx:
182 case AArch64::ANDSWri:
183 case AArch64::ANDSWrr:
184 case AArch64::ANDSWrs:
185 case AArch64::BICSWrr:
186 case AArch64::BICSWrs:
187 case AArch64::SUBSWri:
188 case AArch64::SUBSWrr:
189 case AArch64::SUBSWrs:
190 case AArch64::SUBSWrx:
191 switch (MIOpc) {
192 default:
193 llvm_unreachable("Unexpected opcode!");
194
195 case AArch64::CBZW:
196 case AArch64::CBNZW:
197 case AArch64::TBZW:
198 case AArch64::TBNZW:
199 // Check to see if the TBZ/TBNZ is checking the sign bit.
200 if ((MIOpc == AArch64::TBZW || MIOpc == AArch64::TBNZW) &&
201 MI.getOperand(i: 1).getImm() != 31)
202 return false;
203
204 // There must not be any instruction between DefMI and MI that clobbers or
205 // reads NZCV.
206 if (isNZCVTouchedInInstructionRange(DefMI, UseMI: MI, TRI))
207 return false;
208
209 NewCmp = tryConvertToFlagSetting(MI&: DefMI, IsFlagSetting, /*Is64Bit=*/false);
210 if (!NewCmp)
211 return false;
212
213 LLVM_DEBUG(dbgs() << " Replacing instructions:\n ");
214 LLVM_DEBUG(DefMI.print(dbgs()));
215 LLVM_DEBUG(dbgs() << " ");
216 LLVM_DEBUG(MI.print(dbgs()));
217
218 NewBr = convertToCondBr(MI);
219 break;
220 }
221 break;
222
223 case AArch64::ADDXri:
224 case AArch64::ADDXrr:
225 case AArch64::ADDXrs:
226 case AArch64::ADDXrx:
227 case AArch64::ANDXri:
228 case AArch64::ANDXrr:
229 case AArch64::ANDXrs:
230 case AArch64::BICXrr:
231 case AArch64::BICXrs:
232 case AArch64::SUBXri:
233 case AArch64::SUBXrr:
234 case AArch64::SUBXrs:
235 case AArch64::SUBXrx:
236 IsFlagSetting = false;
237 [[fallthrough]];
238 case AArch64::ADDSXri:
239 case AArch64::ADDSXrr:
240 case AArch64::ADDSXrs:
241 case AArch64::ADDSXrx:
242 case AArch64::ANDSXri:
243 case AArch64::ANDSXrr:
244 case AArch64::ANDSXrs:
245 case AArch64::BICSXrr:
246 case AArch64::BICSXrs:
247 case AArch64::SUBSXri:
248 case AArch64::SUBSXrr:
249 case AArch64::SUBSXrs:
250 case AArch64::SUBSXrx:
251 switch (MIOpc) {
252 default:
253 llvm_unreachable("Unexpected opcode!");
254
255 case AArch64::CBZX:
256 case AArch64::CBNZX:
257 case AArch64::TBZX:
258 case AArch64::TBNZX: {
259 // Check to see if the TBZ/TBNZ is checking the sign bit.
260 if ((MIOpc == AArch64::TBZX || MIOpc == AArch64::TBNZX) &&
261 MI.getOperand(i: 1).getImm() != 63)
262 return false;
263 // There must not be any instruction between DefMI and MI that clobbers or
264 // reads NZCV.
265 if (isNZCVTouchedInInstructionRange(DefMI, UseMI: MI, TRI))
266 return false;
267
268 NewCmp = tryConvertToFlagSetting(MI&: DefMI, IsFlagSetting, /*Is64Bit=*/true);
269 if (!NewCmp)
270 return false;
271
272 LLVM_DEBUG(dbgs() << " Replacing instructions:\n ");
273 LLVM_DEBUG(DefMI.print(dbgs()));
274 LLVM_DEBUG(dbgs() << " ");
275 LLVM_DEBUG(MI.print(dbgs()));
276
277 NewBr = convertToCondBr(MI);
278 break;
279 }
280 }
281 break;
282 }
283 (void)NewCmp; (void)NewBr;
284 assert(NewCmp && NewBr && "Expected new instructions.");
285
286 LLVM_DEBUG(dbgs() << " with instruction:\n ");
287 LLVM_DEBUG(NewCmp->print(dbgs()));
288 LLVM_DEBUG(dbgs() << " ");
289 LLVM_DEBUG(NewBr->print(dbgs()));
290
291 // If this was a flag setting version of the instruction, we use the original
292 // instruction by just clearing the dead marked on the implicit-def of NCZV.
293 // Therefore, we should not erase this instruction.
294 if (!IsFlagSetting)
295 DefMI.eraseFromParent();
296 MI.eraseFromParent();
297 return true;
298}
299
300bool AArch64CondBrTuning::runOnMachineFunction(MachineFunction &MF) {
301 if (skipFunction(F: MF.getFunction()))
302 return false;
303
304 LLVM_DEBUG(
305 dbgs() << "********** AArch64 Conditional Branch Tuning **********\n"
306 << "********** Function: " << MF.getName() << '\n');
307
308 TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
309 TRI = MF.getSubtarget().getRegisterInfo();
310 MRI = &MF.getRegInfo();
311
312 bool Changed = false;
313 for (MachineBasicBlock &MBB : MF) {
314 bool LocalChange = false;
315 for (MachineInstr &MI : MBB.terminators()) {
316 switch (MI.getOpcode()) {
317 default:
318 break;
319 case AArch64::CBZW:
320 case AArch64::CBZX:
321 case AArch64::CBNZW:
322 case AArch64::CBNZX:
323 case AArch64::TBZW:
324 case AArch64::TBZX:
325 case AArch64::TBNZW:
326 case AArch64::TBNZX:
327 MachineInstr *DefMI = getOperandDef(MO: MI.getOperand(i: 0));
328 LocalChange = (DefMI && tryToTuneBranch(MI, DefMI&: *DefMI));
329 break;
330 }
331 // If the optimization was successful, we can't optimize any other
332 // branches because doing so would clobber the NZCV flags.
333 if (LocalChange) {
334 Changed = true;
335 break;
336 }
337 }
338 }
339 return Changed;
340}
341
342FunctionPass *llvm::createAArch64CondBrTuning() {
343 return new AArch64CondBrTuning();
344}
345