AArch64FastISel: Use cbz/cbnz to branch on i1

[MatzeB]

Summary

In the case of a conditional branch without a preceding cmp or similar we
used to emit a "and; cmp; b.eq/b.ne" sequence, use tbz/tbnz instead.

rdar://18783875

[t.p.northover]

I don't think this is quite right. The i1 content of a register is only the low bit, isn't it? This is most easily seen outside Darwin (we're the only ones that extend incoming paramters in the caller), but I think this gets miscompiled even for us:

define i32 @foo(i32 %in) {
entry:
  %val = trunc i32 %in to i1
  br label %do_test

do_test:
  %tst = phi i1 [%val, %entry]
  br i1 %tst, label %true, label %false

true:

  ret i32 0

false:
  ret i32 1
}

So I think we really do need the TBZ for correctness.

[MatzeB]

In D15122#299704, @t.p.northover wrote:

I don't think this is quite right. The i1 content of a register is only the low bit, isn't it? This is most easily seen outside Darwin (we're the only ones that extend incoming paramters in the caller), but I think this gets miscompiled even for us:

define i32 @foo(i32 %in) {
entry:
  %val = trunc i32 %in to i1
  br label %do_test

do_test:
  %tst = phi i1 [%val, %entry]
  br i1 %tst, label %true, label %false

true:

  ret i32 0

false:
  ret i32 1
}

So I think we really do need the TBZ for correctness.

I tried your example and even the non-fast-isel produce by "llc -march arm64 -O0" seems to use plain cbz, is this a bug?

I was looking at TargetLowering::getBooleanContents() comments: "For targets without i1 registers, this gives the nature of the high-bits of boolean values held in types wider than i1." for aarch64 this is "ZeroOrOneBooleanContent".
Is it wrong to rely on this in this context? I admit this warrants an assert(getBooleanContents() != UndefinedBooleanContents) to express my reasoning.

[MatzeB]

After discussing this with Tim, I realized that the "getBooleanContents() properties" are only guaranteed only for the usual boolean inputs of branches/selects/... and only after selection DAG legalization. So we shouldn't rely on that in the FastISel variants and I changed the code to use tbz/tbnz which should be as efficient anyway.

This new version also removes a BrCond(Trunc) special case which was creating more complicated code than the default case now.

(I also see no regressions on the test-suite with this patch).

[t.p.northover]

This looks fine to me. Thanks for updating the patch!

Tim.

lib/Target/AArch64/AArch64FastISel.cpp

if (BI->isUnconditional()) {
     MachineBasicBlock *MSucc = FuncInfo.MBBMap[BI->getSuccessor(0)];
     fastEmitBranch(MSucc, BI->getDebugLoc());
     return true;
   }
 
   MachineBasicBlock *TBB = FuncInfo.MBBMap[BI->getSuccessor(0)];
   MachineBasicBlock *FBB = FuncInfo.MBBMap[BI->getSuccessor(1)];
 
-  AArch64CC::CondCode CC = AArch64CC::NE;
   if (const CmpInst *CI = dyn_cast<CmpInst>(BI->getCondition())) {
     if (CI->hasOneUse() && isValueAvailable(CI)) {
       // Try to optimize or fold the cmp.
       CmpInst::Predicate Predicate = optimizeCmpPredicate(CI);
       switch (Predicate) {
       default:
         break;
       case CmpInst::FCMP_FALSE:
if (CI->hasOneUse() && isValueAvailable(CI)) {
       }
 
       // Emit the cmp.
       if (!emitCmp(CI->getOperand(0), CI->getOperand(1), CI->isUnsigned()))
         return false;
 
       // FCMP_UEQ and FCMP_ONE cannot be checked with a single branch
       // instruction.
-      CC = getCompareCC(Predicate);
+      AArch64CC::CondCode CC = getCompareCC(Predicate);
       AArch64CC::CondCode ExtraCC = AArch64CC::AL;
       switch (Predicate) {
       default:
         break;
       case CmpInst::FCMP_UEQ:
         ExtraCC = AArch64CC::EQ;
         CC = AArch64CC::VS;
         break;
if (CI->hasOneUse() && isValueAvailable(CI)) {
       // Emit the branch.
       BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(AArch64::Bcc))
           .addImm(CC)
           .addMBB(TBB);
 
       finishCondBranch(BI->getParent(), TBB, FBB);
       return true;
     }
-  } else if (TruncInst *TI = dyn_cast<TruncInst>(BI->getCondition())) {
-    MVT SrcVT;
-    if (TI->hasOneUse() && isValueAvailable(TI) &&
-        isTypeSupported(TI->getOperand(0)->getType(), SrcVT)) {
-      unsigned CondReg = getRegForValue(TI->getOperand(0));
-      if (!CondReg)
-        return false;
-      bool CondIsKill = hasTrivialKill(TI->getOperand(0));
-
-      // Issue an extract_subreg to get the lower 32-bits.
-      if (SrcVT == MVT::i64) {
-        CondReg = fastEmitInst_extractsubreg(MVT::i32, CondReg, CondIsKill,
-                                             AArch64::sub_32);
-        CondIsKill = true;
-      }
-
-      unsigned ANDReg = emitAnd_ri(MVT::i32, CondReg, CondIsKill, 1);
-      assert(ANDReg && "Unexpected AND instruction emission failure.");
-      emitICmp_ri(MVT::i32, ANDReg, /*IsKill=*/true, 0);
-
-      if (FuncInfo.MBB->isLayoutSuccessor(TBB)) {
-        std::swap(TBB, FBB);
-        CC = AArch64CC::EQ;
-      }
-      BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(AArch64::Bcc))
-          .addImm(CC)
-          .addMBB(TBB);
-
-      finishCondBranch(BI->getParent(), TBB, FBB);
-      return true;
-    }
   } else if (const auto *CI = dyn_cast<ConstantInt>(BI->getCondition())) {
     uint64_t Imm = CI->getZExtValue();
     MachineBasicBlock *Target = (Imm == 0) ? FBB : TBB;
     BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(AArch64::B))
         .addMBB(Target);
 
     // Obtain the branch probability and add the target to the successor list.
     if (FuncInfo.BPI) {
       auto BranchProbability = FuncInfo.BPI->getEdgeProbability(
           BI->getParent(), Target->getBasicBlock());
       FuncInfo.MBB->addSuccessor(Target, BranchProbability);
     } else
       FuncInfo.MBB->addSuccessorWithoutProb(Target);
     return true;
-  } else if (foldXALUIntrinsic(CC, I, BI->getCondition())) {
+  } else {
+    AArch64CC::CondCode CC = AArch64CC::NE;
+    if (foldXALUIntrinsic(CC, I, BI->getCondition())) {
       // Fake request the condition, otherwise the intrinsic might be completely
       // optimized away.
       unsigned CondReg = getRegForValue(BI->getCondition());
       if (!CondReg)
         return false;
 
       // Emit the branch.
       BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(AArch64::Bcc))
         .addImm(CC)
         .addMBB(TBB);
 
       finishCondBranch(BI->getParent(), TBB, FBB);
       return true;
     }
+  }
 
   unsigned CondReg = getRegForValue(BI->getCondition());
   if (CondReg == 0)
     return false;
   bool CondRegIsKill = hasTrivialKill(BI->getCondition());
 
-  // We've been divorced from our compare!  Our block was split, and
-  // now our compare lives in a predecessor block.  We musn't
-  // re-compare here, as the children of the compare aren't guaranteed
-  // live across the block boundary (we *could* check for this).
-  // Regardless, the compare has been done in the predecessor block,
-  // and it left a value for us in a virtual register.  Ergo, we test
-  // the one-bit value left in the virtual register.
-  //
-  // FIXME: Optimize this with TBZW/TBZNW.
-  unsigned ANDReg = emitAnd_ri(MVT::i32, CondReg, CondRegIsKill, 1);
-  assert(ANDReg && "Unexpected AND instruction emission failure.");
-  emitICmp_ri(MVT::i32, ANDReg, /*IsKill=*/true, 0);
-
+  // i1 conditions come as a i32 values, test the lowest bit with tb(n)z.
+  unsigned Opcode = AArch64::TBNZW;
   if (FuncInfo.MBB->isLayoutSuccessor(TBB)) {
     std::swap(TBB, FBB);
-    CC = AArch64CC::EQ;
+    Opcode = AArch64::TBZW;
   }
 
-  BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(AArch64::Bcc))
-      .addImm(CC)
+  const MCInstrDesc &II = TII.get(Opcode);
+  unsigned ConstrainedCondReg
+    = constrainOperandRegClass(II, CondReg, II.getNumDefs());
+  BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, II)
+      .addReg(ConstrainedCondReg, getKillRegState(CondRegIsKill))
+      .addImm(0)
       .addMBB(TBB);
 
   finishCondBranch(BI->getParent(), TBB, FBB);
   return true;
 }
 
 bool AArch64FastISel::selectIndirectBr(const Instruction *I) {
   const IndirectBrInst *BI = cast<IndirectBrInst>(I);

test/CodeGen/AArch64/arm64-fast-isel-br.ll

entry:
   %b.addr = alloca i16, align 2
   %c.addr = alloca i32, align 4
   %d.addr = alloca i64, align 8
   store i8 %a, i8* %a.addr, align 1
   store i16 %b, i16* %b.addr, align 2
   store i32 %c, i32* %c.addr, align 4
   store i64 %d, i64* %d.addr, align 8
   %0 = load i16, i16* %b.addr, align 2
-; CHECK: and w0, w0, #0x1
-; CHECK: cmp w0, #0
-; CHECK: b.eq LBB4_2
+; CHECK: tbz w0, #0, LBB4_2
   %conv = trunc i16 %0 to i1
   br i1 %conv, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   call void @foo1()
   br label %if.end
 
 if.end:                                           ; preds = %if.then, %entry
   %1 = load i32, i32* %c.addr, align 4
-; CHECK: and w[[REG:[0-9]+]], w{{[0-9]+}}, #0x1
-; CHECK: cmp w[[REG]], #0
-; CHECK: b.eq LBB4_4
+; CHECK: tbz w{{[0-9]+}}, #0, LBB4_4
   %conv1 = trunc i32 %1 to i1
   br i1 %conv1, label %if.then3, label %if.end4
 
 if.then3:                                         ; preds = %if.end
   call void @foo1()
   br label %if.end4
 
 if.end4:                                          ; preds = %if.then3, %if.end
   %2 = load i64, i64* %d.addr, align 8
-; CHECK: cmp w{{[0-9]+}}, #0
-; CHECK: b.eq LBB4_6
+; CHECK: tbz w{{[0-9]+}}, #0, LBB4_6
   %conv5 = trunc i64 %2 to i1
   br i1 %conv5, label %if.then7, label %if.end8
 
 if.then7:                                         ; preds = %if.end4
   call void @foo1()
   br label %if.end8
 
 if.end8:                                          ; preds = %if.then7, %if.end4
   %3 = load i8, i8* %a.addr, align 1
   ret i8 %3
 }
 
 declare void @foo1()
 
 ; rdar://15174028
 define i32 @trunc64(i64 %foo) nounwind {
 ; CHECK: trunc64
 ; CHECK: and  [[REG1:x[0-9]+]], x0, #0x1
 ; CHECK: mov  x[[REG2:[0-9]+]], [[REG1]]
-; CHECK: and  [[REG3:w[0-9]+]], w[[REG2]], #0x1
-; CHECK: cmp  [[REG3]], #0
-; CHECK: b.eq LBB5_2
+; CHECK: tbz w[[REG2]], #0, LBB5_2
   %a = and i64 %foo, 1
   %b = trunc i64 %a to i1
   br i1 %b, label %if.then, label %if.else
 
 if.then:
   ret i32 1
 
 if.else:
   ret i32 0
 }

test/CodeGen/AArch64/fast-isel-branch-cond-mask.ll

 ; RUN: llc -mtriple=aarch64-apple-darwin -O0 -fast-isel -fast-isel-abort=0 -verify-machineinstrs < %s | FileCheck %s
 
 define void @test(i64 %a, i64 %b, i2* %c) {
 ; CHECK-LABEL: test
 ; CHECK:       and [[REG1:w[0-9]+]], w8, #0x3
 ; CHECK-NEXT:  strb [[REG1]], {{\[}}x2{{\]}}
-; CHECK:       and [[REG2:w[0-9]+]], w8, #0x1
-; CHECK-NEXT:  cmp [[REG2]], #0
+; CHECK-NEXT:  tbz w9, #0,
  %1 = trunc i64 %a to i2
  %2 = trunc i64 %b to i1
 ; Force fast-isel to fall back to SDAG.
  store i2 %1, i2* %c, align 8
  br i1 %2, label %bb1, label %bb2
 
 bb1:
   ret void
 
 bb2:
   ret void
 }

test/CodeGen/AArch64/fast-isel-branch-cond-split.ll

 ; If the branch is unpredictable, don't add another branch.
 
 ; CHECK-LABEL: test_or_unpredictable
 ; CHECK:       cmp   w0, #0
 ; CHECK-NEXT:  cset  w8, eq
 ; CHECK-NEXT:  cmp   w1, #0
 ; CHECK-NEXT:  cset  w9, eq
 ; CHECK-NEXT:  orr   w8, w8, w9
-; CHECK-NEXT:  and   w8, w8, #0x1
-; CHECK-NEXT:  cmp   w8, #0
-; CHECK-NEXT:  b.ne 
+; CHECK-NEXT:  tbnz w8, #0,
 define i64 @test_or_unpredictable(i32 %a, i32 %b) {
 bb1:
   %0 = icmp eq i32 %a, 0
   %1 = icmp eq i32 %b, 0
   %or.cond = or i1 %0, %1
   br i1 %or.cond, label %bb3, label %bb4, !unpredictable !2
 
 bb3:
   ret i64 0
 
 bb4:
   %2 = call i64 @bar()
   ret i64 %2
 }
 
 ; CHECK-LABEL: test_and_unpredictable
 ; CHECK:       cmp   w0, #0
 ; CHECK-NEXT:  cset  w8, ne
 ; CHECK-NEXT:  cmp   w1, #0
 ; CHECK-NEXT:  cset  w9, ne
 ; CHECK-NEXT:  and   w8, w8, w9
-; CHECK-NEXT:  and   w8, w8, #0x1
-; CHECK-NEXT:  cmp   w8, #0
-; CHECK-NEXT:  b.eq 
+; CHECK-NEXT:  tbz w8, #0,
 define i64 @test_and_unpredictable(i32 %a, i32 %b) {
 bb1:
   %0 = icmp ne i32 %a, 0
   %1 = icmp ne i32 %b, 0
   %or.cond = and i1 %0, %1
   br i1 %or.cond, label %bb4, label %bb3, !unpredictable !2
 
 bb3: