[AArch64][GlobalISel] Optimize conditional branches followed by unconditional branches

[aemerson]

Summary

If we have an icmp->brcond->br sequence where the brcond just branches to the next block jumping over the br, while the br takes the false edge, then we can modify the conditional branch to jump to the br's target while inverting the condition of the incoming icmp. This means we can eliminate the br as an unconditional branch to the fallthrough block.

[arsenm]

Why is this done in the target selector instead of in a generic combine?

[arsenm]

In D64354#1573842, @arsenm wrote:

Why is this done in the target selector instead of in a generic combine?

I'm pretty sure SelectionDAGBuilder does this today

[aditya_nandakumar]

I haven't read the patch yet, but just from reading the description I remembered that the DAG does this while building up the DAG. Maybe this needs to be done in the IRTranslator to keep this behavior? (This helps while comparing codegen with sdag/gisel).

[qcolombet]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1171 ↗

(On Diff #208448)

Taking advantage of fall-through sounds like something beneficial for all targets.
Could we turn this into a combiner helper?

(Then you can call this specific helper in here.)

[arsenm]

FWIW I would like some control over what boolean sources are acceptable to invert. SelectionDAG inserting a condition invert between a control flow intrinsic and the branch required work to undo

[aemerson]

This seemed a little presumptive to do for all targets for me. I can try to move it earlier into a Combiner helper, since I don't want to put optimizations into the IRTranslator unless we can't do it in other places.

In D64354#1573876, @arsenm wrote:

FWIW I would like some control over what boolean sources are acceptable to invert. SelectionDAG inserting a condition invert between a control flow intrinsic and the branch required work to undo

In this case I'm just inverting the condition code, we shouldn't need to insert extra instructions.

[paquette]

Combiner helper makes sense here IMO.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1186 ↗

(On Diff #208448)

Would BrIt->isTerminator() work here?

1197 ↗

(On Diff #208448)

Is trunc the only thing we should look through here? Can we also handle copies, for example?

[aemerson]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1186 ↗

(On Diff #208448)

Branch is a terminator by nature, so it doesn't check the same thing, although the usefulness of this assert is debatable.

1197 ↗

(On Diff #208448)

For arm64 it seems to be the only thing in between compares and branches in practice due to our legalization rules. This code will be removed anyway since I'm going to put this before the legalizer so we shouldn't see any artifacts in between the ops,

[aemerson]

New patch to do the optimization in the CombinerHelper and add it to the prelegalizer combiner for AArch64.

This doesn't get all the code size improvements of the last patch, I'll make that a separate patch. On it's own, this change gives a 2.5% geomean improvement on CTMark -O0.

[aemerson]

Add missing test.

[aditya_nandakumar]

Unless others have comments, this looks good to go in to me.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

392

Orthogonal to this patch - I've found that MachineOperand::getPredicate returning an unsigned and users having to explicitly cast to Predicate is a little annoying. What do you think of getPredicate returning Predicate type (in a separate patch)? In my use cases I've almost never used the unsigned value without casting to Predicate.

[aemerson]

Marking as accepted.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

392

That sounds reasonable to me.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

public:
   void applyCombineCopy(MachineInstr &MI);
 
   /// If \p MI is extend that consumes the result of a load, try to combine it.
   /// Returns true if MI changed.
   bool tryCombineExtendingLoads(MachineInstr &MI);
   bool matchCombineExtendingLoads(MachineInstr &MI, PreferredTuple &MatchInfo);
   void applyCombineExtendingLoads(MachineInstr &MI, PreferredTuple &MatchInfo);
 
+  bool matchCombineBr(MachineInstr &MI);
+  bool tryCombineBr(MachineInstr &MI);
+
   /// Try to transform \p MI by using all of the above
   /// combine functions. Returns true if changed.
   bool tryCombine(MachineInstr &MI);
 };
 } // namespace llvm
 
 #endif

llvm/include/llvm/CodeGen/MachineOperand.h

public:
   /// operand does not take ownership of the memory referenced by Mask, it must
   /// remain valid for the lifetime of the operand. See CreateRegMask().
   /// Any physreg with a 0 bit in the mask is clobbered by the instruction.
   void setRegMask(const uint32_t *RegMaskPtr) {
     assert(isRegMask() && "Wrong MachineOperand mutator");
     Contents.RegMask = RegMaskPtr;
   }
 
+  void setPredicate(unsigned Predicate) {
+    assert(isPredicate() && "Wrong MachineOperand mutator");
+    Contents.Pred = Predicate;
+  }
+
   //===--------------------------------------------------------------------===//
   // Other methods.
   //===--------------------------------------------------------------------===//
 
   /// Returns true if this operand is identical to the specified operand except
   /// for liveness related flags (isKill, isUndef and isDead). Note that this
   /// should stay in sync with the hash_value overload below.
   bool isIdenticalTo(const MachineOperand &Other) const;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

for (auto *UseMO : Uses) {
     // This is free on many targets.
     InsertInsnsWithoutSideEffectsBeforeUse(Builder, MI, *UseMO, InsertTruncAt);
   }
 
   MI.getOperand(0).setReg(ChosenDstReg);
   Observer.changedInstr(MI);
 }
 
+bool CombinerHelper::matchCombineBr(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_BR && "Expected a G_BR");
+  // Try to match the following:
+  // bb1:
+  //   %c(s32) = G_ICMP pred, %a, %b
+  //   %c1(s1) = G_TRUNC %c(s32)
+  //   G_BRCOND %c1, %bb2
+  //   G_BR %bb3
+  // bb2:
+  // ...
+  // bb3:
+
+  // The above pattern does not have a fall through to the successor bb2, always
+  // resulting in a branch no matter which path is taken. Here we try to find
+  // and replace that pattern with conditional branch to bb3 and otherwise
+  // fallthrough to bb2.
+
+  MachineBasicBlock *MBB = MI.getParent();
+  MachineBasicBlock::iterator BrIt(MI);
+  if (BrIt == MBB->begin())
+    return false;
+  assert(std::next(BrIt) == MBB->end() && "expected G_BR to be a terminator");
+
+  MachineInstr *BrCond = &*std::prev(BrIt);
+  if (BrCond->getOpcode() != TargetOpcode::G_BRCOND)
+    return false;
+
+  // Check that the next block is the conditional branch target.
+  if (!MBB->isLayoutSuccessor(BrCond->getOperand(1).getMBB()))
+    return false;
+
+  MachineInstr *CmpMI = MRI.getVRegDef(BrCond->getOperand(0).getReg());
+  if (!CmpMI || CmpMI->getOpcode() != TargetOpcode::G_ICMP ||
+      !MRI.hasOneUse(CmpMI->getOperand(0).getReg()))
+    return false;
+  return true;
+}
+
+bool CombinerHelper::tryCombineBr(MachineInstr &MI) {
+  if (!matchCombineBr(MI))
+    return false;
+  MachineBasicBlock *BrTarget = MI.getOperand(0).getMBB();
+  MachineBasicBlock::iterator BrIt(MI);
+  MachineInstr *BrCond = &*std::prev(BrIt);
+  MachineInstr *CmpMI = MRI.getVRegDef(BrCond->getOperand(0).getReg());
+
+  CmpInst::Predicate InversePred = CmpInst::getInversePredicate(
+      (CmpInst::Predicate)CmpMI->getOperand(1).getPredicate());
+
+  // Invert the G_ICMP condition.
+  Observer.changingInstr(*CmpMI);
+  CmpMI->getOperand(1).setPredicate(InversePred);
+  Observer.changedInstr(*CmpMI);
+
+  // Change the conditional branch target.
+  Observer.changingInstr(*BrCond);
+  BrCond->getOperand(1).setMBB(BrTarget);
+  Observer.changedInstr(*BrCond);
+  MI.eraseFromParent();
+  return true;
+}
+
 bool CombinerHelper::tryCombine(MachineInstr &MI) {
   if (tryCombineCopy(MI))
     return true;
   return tryCombineExtendingLoads(MI);
 }

llvm/lib/Target/AArch64/AArch64PreLegalizerCombiner.cpp

bool AArch64PreLegalizerCombinerInfo::combine(GISelChangeObserver &Observer,
                                               MachineIRBuilder &B) const {
   CombinerHelper Helper(Observer, B);
 
   switch (MI.getOpcode()) {
   default:
     return false;
   case TargetOpcode::COPY:
     return Helper.tryCombineCopy(MI);
+  case TargetOpcode::G_BR:
+    return Helper.tryCombineBr(MI);
   case TargetOpcode::G_LOAD:
   case TargetOpcode::G_SEXTLOAD:
   case TargetOpcode::G_ZEXTLOAD:
     return Helper.tryCombineExtendingLoads(MI);
   }
 
   return false;
 }

llvm/test/CodeGen/AArch64/GlobalISel/localizer-arm64-tti.ll

define i32 @foo() {
   ; CHECK:   [[GV:%[0-9]+]]:gpr(p0) = G_GLOBAL_VALUE @var1
   ; CHECK:   [[C:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 1
   ; CHECK:   [[C1:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 2
   ; CHECK:   [[GV1:%[0-9]+]]:gpr(p0) = G_GLOBAL_VALUE @var2
   ; CHECK:   [[C2:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 3
   ; CHECK:   [[GV2:%[0-9]+]]:gpr(p0) = G_GLOBAL_VALUE @var3
   ; CHECK:   [[C3:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 0
   ; CHECK:   [[LOAD:%[0-9]+]]:gpr(s32) = G_LOAD [[GV]](p0) :: (load 4 from @var1)
-  ; CHECK:   [[ICMP:%[0-9]+]]:gpr(s32) = G_ICMP intpred(eq), [[LOAD]](s32), [[C]]
+  ; CHECK:   [[ICMP:%[0-9]+]]:gpr(s32) = G_ICMP intpred(ne), [[LOAD]](s32), [[C]]
   ; CHECK:   [[TRUNC:%[0-9]+]]:gpr(s1) = G_TRUNC [[ICMP]](s32)
-  ; CHECK:   G_BRCOND [[TRUNC]](s1), %bb.2
-  ; CHECK:   G_BR %bb.3
+  ; CHECK:   G_BRCOND [[TRUNC]](s1), %bb.3
   ; CHECK: bb.2.if.then:
   ; CHECK:   successors: %bb.3(0x80000000)
   ; CHECK:   [[GV3:%[0-9]+]]:gpr(p0) = G_GLOBAL_VALUE @var2
   ; CHECK:   [[C4:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 2
   ; CHECK:   G_STORE [[C4]](s32), [[GV3]](p0) :: (store 4 into @var2)
   ; CHECK:   [[C5:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 3
   ; CHECK:   G_STORE [[C5]](s32), [[GV]](p0) :: (store 4 into @var1)
   ; CHECK:   [[GV4:%[0-9]+]]:gpr(p0) = G_GLOBAL_VALUE @var3

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-br.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -O0 -run-pass=aarch64-prelegalizer-combiner -global-isel -verify-machineinstrs %s -o - | FileCheck %s
+--- |
+  target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+  target triple = "arm64-apple-ios5.0.0"
+
+  define i32 @foo(i32 %a, i32 %b) {
+  entry:
+    %cmp = icmp sgt i32 %a, 0
+    br i1 %cmp, label %if.then, label %if.end
+
+  if.then:
+    %add = add nsw i32 %b, %a
+    %add1 = add nsw i32 %a, %b
+    br label %return
+
+  if.end:
+    %mul = mul nsw i32 %b, %b
+    %add2 = add nuw nsw i32 %mul, 2
+    br label %return
+
+  return:
+    %retval.0 = phi i32 [ %add1, %if.then ], [ %add2, %if.end ]
+    ret i32 %retval.0
+  }
+
+
+...
+---
+name:            foo
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: foo
+  ; CHECK: bb.0.entry:
+  ; CHECK:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK:   liveins: $w0, $w1
+  ; CHECK:   [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+  ; CHECK:   [[COPY1:%[0-9]+]]:_(s32) = COPY $w1
+  ; CHECK:   [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+  ; CHECK:   [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 2
+  ; CHECK:   [[ICMP:%[0-9]+]]:_(s1) = G_ICMP intpred(sle), [[COPY]](s32), [[C]]
+  ; CHECK:   G_BRCOND [[ICMP]](s1), %bb.2
+  ; CHECK: bb.1.if.then:
+  ; CHECK:   successors: %bb.3(0x80000000)
+  ; CHECK:   %5:_(s32) = nsw G_ADD [[COPY1]], [[COPY]]
+  ; CHECK:   %6:_(s32) = nsw G_ADD %5, [[COPY1]]
+  ; CHECK:   G_BR %bb.3
+  ; CHECK: bb.2.if.end:
+  ; CHECK:   successors: %bb.3(0x80000000)
+  ; CHECK:   %7:_(s32) = nsw G_MUL [[COPY1]], [[COPY1]]
+  ; CHECK:   %8:_(s32) = nuw nsw G_ADD %7, [[C1]]
+  ; CHECK: bb.3.return:
+  ; CHECK:   [[PHI:%[0-9]+]]:_(s32) = G_PHI %6(s32), %bb.1, %8(s32), %bb.2
+  ; CHECK:   $w0 = COPY [[PHI]](s32)
+  ; CHECK:   RET_ReallyLR implicit $w0
+  bb.1.entry:
+    liveins: $w0, $w1
+
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = COPY $w1
+    %2:_(s32) = G_CONSTANT i32 0
+    %5:_(s32) = G_CONSTANT i32 2
+    %3:_(s1) = G_ICMP intpred(sgt), %0(s32), %2
+    G_BRCOND %3(s1), %bb.2
+    G_BR %bb.3
+
+  bb.2.if.then:
+    %7:_(s32) = nsw G_ADD %1, %0
+    %8:_(s32) = nsw G_ADD %7, %1
+    G_BR %bb.4
+
+  bb.3.if.end:
+    %4:_(s32) = nsw G_MUL %1, %1
+    %6:_(s32) = nuw nsw G_ADD %4, %5
+
+  bb.4.return:
+    %10:_(s32) = G_PHI %8(s32), %bb.2, %6(s32), %bb.3
+    $w0 = COPY %10(s32)
+    RET_ReallyLR implicit $w0
+
+...

llvm/test/CodeGen/AArch64/speculation-hardening.ll

 ; NOSLH-NOT:  mov sp, [[TMPREG]]
   %call = tail call i32 @tail_callee(i32 %i)
 ; SLH:  cmp sp, #0
 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
   %cmp = icmp slt i32 %call, %N
   br i1 %cmp, label %if.then, label %return
-; GlobalISel lowers the branch to a b.ne sometimes instead of b.ge as expected..
-; CHECK: b.[[COND:(ge)|(lt)|(ne)]]
+; CHECK: b.[[COND:(ge)|(lt)|(ne)|(eq)]]
 
 if.then:                                          ; preds = %entry
-; NOSLH-NOT: csel x16, x16, xzr, {{(lt)|(ge)|(eq)}}
-; SLH-DAG: csel x16, x16, xzr, {{(lt)|(ge)|(eq)}}
+; NOSLH-NOT: csel x16, x16, xzr, {{(lt)|(ge)|(eq)|(ne)}}
+; SLH-DAG: csel x16, x16, xzr, {{(lt)|(ge)|(eq)|(ne)}}
   %idxprom = sext i32 %i to i64
   %arrayidx = getelementptr inbounds i8, i8* %p, i64 %idxprom
   %0 = load i8, i8* %arrayidx, align 1
 ; CHECK-DAG:      ldrb [[LOADED:w[0-9]+]],
   %conv = zext i8 %0 to i32
   br label %return
 
 ; SLH-DAG: csel x16, x16, xzr, [[COND]]

 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
   %l5 = extractvalue { i8*, i32 } %l4, 0
   %l6 = tail call i8* @__cxa_begin_catch(i8* %l5)
   %l7 = icmp sgt i32 %l0, %l1
   br i1 %l7, label %then, label %else
 ; GlobalISel lowers the branch to a b.ne sometimes instead of b.ge as expected..
-; CHECK: b.[[COND:(le)|(gt)|(ne)]]
+; CHECK: b.[[COND:(le)|(gt)|(ne)|(eq)]]
 
 then:
 ; SLH-DAG: csel x16, x16, xzr, [[COND]]
   %l9 = sdiv i32 %l0, %l1
   br label %postif
 
 else:
-; SLH-DAG: csel x16, x16, xzr, {{(gt)|(le)|(eq)}}
+; SLH-DAG: csel x16, x16, xzr, {{(gt)|(le)|(eq)|(ne)}}
   %l11 = sdiv i32 %l1, %l0
   br label %postif
 
 postif:
   %l13 = phi i32 [ %l9, %then ], [ %l11, %else ]
   tail call void @__cxa_end_catch()
   br label %exit