Skip to content

Commit a058ee5

Browse files
michaelmaitlandmichaelmaitland
committed
[BranchFolding][RISCV] Optimize branches NE/EQ to zero or one
There are two changes here. The first that we teach analyzeBranch how to evaluate a conditional branch followed by a unconditional branch such that we can evaluate the conditional branch statically. Often, we will see comparison to one or zero since SelectionDAG often uses i1 for the conditional comparison. As a result, we handle this specific case. We handle only EQ and NEQ for now, but this can be expanded in the future. We can also expand on handling arbitrary constants in the future. The second change is that we pass AllowModify=false to analyzeBranch in the tail merging code. The reason we do this is because this code is doing some clever tricks to the branch code that it will restore later. Now that we are actually optimizing branches in analyzeBranch, we have to be careful not to mess up this canonical form that the tail merging code expects.
1 parent b77baa4 commit a058ee5

File tree

7 files changed

+169
-78
lines changed

7 files changed

+169
-78
lines changed

‎llvm/lib/CodeGen/BranchFolding.cpp

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,7 @@ static void FixTail(MachineBasicBlock *CurMBB, MachineBasicBlock *SuccBB,
492492
DebugLoc dl = CurMBB->findBranchDebugLoc();
493493
if (!dl)
494494
dl = BranchDL;
495-
if (I != MF->end() && !TII->analyzeBranch(*CurMBB, TBB, FBB, Cond, true)) {
495+
if (I != MF->end() && !TII->analyzeBranch(*CurMBB, TBB, FBB, Cond)) {
496496
MachineBasicBlock *NextBB = &*I;
497497
if (TBB == NextBB && !Cond.empty() && !FBB) {
498498
if (!TII->reverseBranchCondition(Cond)) {
@@ -1132,7 +1132,7 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
11321132

11331133
MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
11341134
SmallVector<MachineOperand, 4> Cond;
1135-
if (!TII->analyzeBranch(*PBB, TBB, FBB, Cond, true)) {
1135+
if (!TII->analyzeBranch(*PBB, TBB, FBB, Cond)) {
11361136
// Failing case: IBB is the target of a cbr, and we cannot reverse the
11371137
// branch.
11381138
SmallVector<MachineOperand, 4> NewCond(Cond);
@@ -1589,7 +1589,8 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
15891589
// Loop: xxx; jcc Out; jmp Loop
15901590
// we want:
15911591
// Loop: xxx; jncc Loop; jmp Out
1592-
if (CurTBB && CurFBB && CurFBB == MBB && CurTBB != MBB) {
1592+
if (CurTBB && CurFBB && CurFBB == MBB && CurTBB != MBB &&
1593+
!CurCond.empty()) {
15931594
SmallVector<MachineOperand, 4> NewCond(CurCond);
15941595
if (!TII->reverseBranchCondition(NewCond)) {
15951596
DebugLoc Dl = MBB->findBranchDebugLoc();

‎llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1014,6 +1014,109 @@ RISCVCC::CondCode RISCVCC::getOppositeBranchCondition(RISCVCC::CondCode CC) {
10141014
}
10151015
}
10161016

1017+
bool RISCVInstrInfo::trySimplifyCondBr(
1018+
MachineBasicBlock &MBB, MachineBasicBlock *TBB, MachineBasicBlock *FBB,
1019+
SmallVectorImpl<MachineOperand> &Cond) const {
1020+
1021+
if (!TBB || Cond.size() != 3)
1022+
return false;
1023+
1024+
RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(Cond[0].getImm());
1025+
auto LHS = Cond[1];
1026+
auto RHS = Cond[2];
1027+
1028+
// Return true if MO definitely contains the value one.
1029+
auto isOne = [](MachineOperand &MO) -> bool {
1030+
if (MO.isImm() && MO.getImm() == 1)
1031+
return true;
1032+
1033+
if (!MO.isReg() || !MO.getReg().isVirtual())
1034+
return false;
1035+
1036+
MachineRegisterInfo &MRI =
1037+
MO.getParent()->getParent()->getParent()->getRegInfo();
1038+
MachineInstr *DefMI = MRI.getUniqueVRegDef(MO.getReg());
1039+
if (!DefMI)
1040+
return false;
1041+
1042+
// For now, just check the canonical one value.
1043+
if (DefMI->getOpcode() == RISCV::ADDI &&
1044+
DefMI->getOperand(1).getReg() == RISCV::X0 &&
1045+
DefMI->getOperand(2).getImm() == 1)
1046+
return true;
1047+
1048+
return false;
1049+
};
1050+
1051+
// Return true if MO definitely contains the value zero.
1052+
auto isZero = [](MachineOperand &MO) -> bool {
1053+
if (MO.isImm() && MO.getImm() == 0)
1054+
return true;
1055+
if (MO.isReg() && MO.getReg() == RISCV::X0)
1056+
return true;
1057+
return false;
1058+
};
1059+
1060+
MachineBasicBlock *Folded = nullptr;
1061+
switch (CC) {
1062+
default:
1063+
// TODO: Implement for more CCs
1064+
return false;
1065+
case RISCVCC::COND_EQ: {
1066+
// We can statically evaluate that we take the first branch
1067+
if ((isZero(LHS) && isZero(RHS)) || (isOne(LHS) && isOne(RHS))) {
1068+
Folded = TBB;
1069+
break;
1070+
}
1071+
// We can statically evaluate that we take the second branch
1072+
if ((isZero(LHS) && isOne(RHS)) || (isOne(LHS) && isZero(RHS))) {
1073+
Folded = FBB;
1074+
break;
1075+
}
1076+
return false;
1077+
}
1078+
case RISCVCC::COND_NE: {
1079+
// We can statically evaluate that we take the first branch
1080+
if ((isOne(LHS) && isZero(RHS)) || (isZero(LHS) && isOne(RHS))) {
1081+
Folded = TBB;
1082+
break;
1083+
}
1084+
// We can statically evaluate that we take the second branch
1085+
if ((isZero(LHS) && isZero(RHS)) || (isOne(LHS) && isOne(RHS))) {
1086+
Folded = FBB;
1087+
break;
1088+
}
1089+
return false;
1090+
}
1091+
}
1092+
1093+
// At this point, its legal to optimize.
1094+
removeBranch(MBB);
1095+
Cond.clear();
1096+
1097+
// Only need to insert a branch if we're not falling through.
1098+
if (Folded) {
1099+
DebugLoc DL = MBB.findBranchDebugLoc();
1100+
insertBranch(MBB, Folded, nullptr, {}, DL);
1101+
}
1102+
1103+
// Update the successors. Remove them all and add back the correct one.
1104+
while (!MBB.succ_empty())
1105+
MBB.removeSuccessor(MBB.succ_end() - 1);
1106+
1107+
// If it's a fallthrough, we need to figure out where MBB is going.
1108+
if (!Folded) {
1109+
MachineFunction::iterator Fallthrough = ++MBB.getIterator();
1110+
if (Fallthrough != MBB.getParent()->end())
1111+
MBB.addSuccessor(&*Fallthrough);
1112+
} else
1113+
MBB.addSuccessor(Folded);
1114+
1115+
TBB = Folded;
1116+
FBB = nullptr;
1117+
return true;
1118+
}
1119+
10171120
bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
10181121
MachineBasicBlock *&TBB,
10191122
MachineBasicBlock *&FBB,
@@ -1071,6 +1174,9 @@ bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
10711174
// Handle a single conditional branch.
10721175
if (NumTerminators == 1 && I->getDesc().isConditionalBranch()) {
10731176
parseCondBranch(*I, TBB, Cond);
1177+
// Try to fold the branch of the conditional branch into a the fallthru.
1178+
if (AllowModify)
1179+
trySimplifyCondBr(MBB, TBB, FBB, Cond);
10741180
return false;
10751181
}
10761182

@@ -1079,6 +1185,10 @@ bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
10791185
I->getDesc().isUnconditionalBranch()) {
10801186
parseCondBranch(*std::prev(I), TBB, Cond);
10811187
FBB = getBranchDestBlock(*I);
1188+
// Try to fold the branch of the conditional branch into an unconditional
1189+
// branch.
1190+
if (AllowModify)
1191+
trySimplifyCondBr(MBB, TBB, FBB, Cond);
10821192
return false;
10831193
}
10841194

‎llvm/lib/Target/RISCV/RISCVInstrInfo.h

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,26 @@ class RISCVInstrInfo : public RISCVGenInstrInfo {
319319
const MachineInstr &MI2) const;
320320
bool hasReassociableVectorSibling(const MachineInstr &Inst,
321321
bool &Commuted) const;
322+
/// Return true if the branch represented by the conditional branch with
323+
/// components TBB, FBB, and CurCond was folded into an unconditional branch.
324+
///
325+
/// If FBB is nullptr, then the the input represents a conditional branch with
326+
/// a fallthrough.
327+
///
328+
/// For example:
329+
/// BRCOND EQ 0, 0, BB1
330+
/// BR BB2
331+
///
332+
/// can be simplified to BR BB1 since 0 == 0 statically. On the other hand,
333+
///
334+
///
335+
/// BRCOND EQ 0, 1, BB1
336+
/// BR BB2
337+
///
338+
/// can be simplified to BR BB2 because 0 != 1 statically.
339+
bool trySimplifyCondBr(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
340+
MachineBasicBlock *FBB,
341+
SmallVectorImpl<MachineOperand> &Cond) const;
322342
};
323343

324344
namespace RISCV {

‎llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -357,11 +357,6 @@ define i64 @ctpop_i64(i64 %a) nounwind {
357357
define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
358358
; RV32I-LABEL: ctpop_i64_ugt_two:
359359
; RV32I: # %bb.0:
360-
; RV32I-NEXT: beqz zero, .LBB6_2
361-
; RV32I-NEXT: # %bb.1:
362-
; RV32I-NEXT: sltiu a0, zero, 0
363-
; RV32I-NEXT: ret
364-
; RV32I-NEXT: .LBB6_2:
365360
; RV32I-NEXT: srli a2, a0, 1
366361
; RV32I-NEXT: lui a3, 349525
367362
; RV32I-NEXT: lui a4, 209715
@@ -404,11 +399,6 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
404399
;
405400
; RV32ZBB-LABEL: ctpop_i64_ugt_two:
406401
; RV32ZBB: # %bb.0:
407-
; RV32ZBB-NEXT: beqz zero, .LBB6_2
408-
; RV32ZBB-NEXT: # %bb.1:
409-
; RV32ZBB-NEXT: sltiu a0, zero, 0
410-
; RV32ZBB-NEXT: ret
411-
; RV32ZBB-NEXT: .LBB6_2:
412402
; RV32ZBB-NEXT: cpop a0, a0
413403
; RV32ZBB-NEXT: cpop a1, a1
414404
; RV32ZBB-NEXT: add a0, a1, a0
@@ -422,11 +412,6 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
422412
define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
423413
; RV32I-LABEL: ctpop_i64_ugt_one:
424414
; RV32I: # %bb.0:
425-
; RV32I-NEXT: beqz zero, .LBB7_2
426-
; RV32I-NEXT: # %bb.1:
427-
; RV32I-NEXT: snez a0, zero
428-
; RV32I-NEXT: ret
429-
; RV32I-NEXT: .LBB7_2:
430415
; RV32I-NEXT: srli a2, a0, 1
431416
; RV32I-NEXT: lui a3, 349525
432417
; RV32I-NEXT: lui a4, 209715
@@ -470,11 +455,6 @@ define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
470455
;
471456
; RV32ZBB-LABEL: ctpop_i64_ugt_one:
472457
; RV32ZBB: # %bb.0:
473-
; RV32ZBB-NEXT: beqz zero, .LBB7_2
474-
; RV32ZBB-NEXT: # %bb.1:
475-
; RV32ZBB-NEXT: snez a0, zero
476-
; RV32ZBB-NEXT: ret
477-
; RV32ZBB-NEXT: .LBB7_2:
478458
; RV32ZBB-NEXT: cpop a0, a0
479459
; RV32ZBB-NEXT: cpop a1, a1
480460
; RV32ZBB-NEXT: add a0, a1, a0

‎llvm/test/CodeGen/RISCV/branch_zero.ll

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,11 @@
55
define void @foo(i16 %finder_idx) {
66
; CHECK-LABEL: foo:
77
; CHECK: # %bb.0: # %entry
8-
; CHECK-NEXT: .LBB0_1: # %for.body
9-
; CHECK-NEXT: # =>This Inner Loop Header: Depth=1
8+
; CHECK-NEXT: # %bb.1: # %for.body
109
; CHECK-NEXT: slli a0, a0, 48
11-
; CHECK-NEXT: bltz a0, .LBB0_4
10+
; CHECK-NEXT: bltz a0, .LBB0_3
1211
; CHECK-NEXT: # %bb.2: # %while.cond.preheader.i
13-
; CHECK-NEXT: # in Loop: Header=BB0_1 Depth=1
14-
; CHECK-NEXT: li a0, 0
15-
; CHECK-NEXT: bnez zero, .LBB0_1
16-
; CHECK-NEXT: # %bb.3: # %while.body
17-
; CHECK-NEXT: .LBB0_4: # %while.cond1.preheader.i
12+
; CHECK-NEXT: .LBB0_3: # %while.cond1.preheader.i
1813
entry:
1914
br label %for.body
2015

@@ -46,16 +41,11 @@ if.then:
4641
define void @bar(i16 %finder_idx) {
4742
; CHECK-LABEL: bar:
4843
; CHECK: # %bb.0: # %entry
49-
; CHECK-NEXT: .LBB1_1: # %for.body
50-
; CHECK-NEXT: # =>This Inner Loop Header: Depth=1
44+
; CHECK-NEXT: # %bb.1: # %for.body
5145
; CHECK-NEXT: slli a0, a0, 48
52-
; CHECK-NEXT: bgez a0, .LBB1_4
46+
; CHECK-NEXT: bgez a0, .LBB1_3
5347
; CHECK-NEXT: # %bb.2: # %while.cond.preheader.i
54-
; CHECK-NEXT: # in Loop: Header=BB1_1 Depth=1
55-
; CHECK-NEXT: li a0, 0
56-
; CHECK-NEXT: bnez zero, .LBB1_1
57-
; CHECK-NEXT: # %bb.3: # %while.body
58-
; CHECK-NEXT: .LBB1_4: # %while.cond1.preheader.i
48+
; CHECK-NEXT: .LBB1_3: # %while.cond1.preheader.i
5949
entry:
6050
br label %for.body
6151

@@ -83,4 +73,3 @@ exit1:
8373
if.then:
8474
br label %for.body
8575
}
86-

0 commit comments

Comments
 (0)