[AArch64] Fix fast-isel of cbz of i1, i8, i16

[olista01]

Summary

This fixes a miscompilation in the AArch64 fast-isel which was triggered when a branch is based on an icmp with condition eq or ne, and type i1, i8 or i16. The cbz instruction compares the whole 32-bit register, so values with the bottom 1, 8 or 16 bits clear would cause the wrong branch to be taken.

[t.p.northover]

Hi Oliver,

This looks fine, just one very minor quibble.

lib/Target/AArch64/AArch64FastISel.cpp

2184

BW < 32?

[t.p.northover]

Go ahead and commit if you want (with or without the suggestion, realistically).

[ributzka]

Hi Oliver,

thanks for catching this. I only have a minor comment inline.
It would be great if you also could optimize the i1 case by emitting a 'tbs' instead. Maybe in a separate commit?

Thanks and LGTM.

-Juergen

lib/Target/AArch64/AArch64FastISel.cpp

2185

getValueType doesn't use kill flags - a simple 'true' will do here.

test/CodeGen/AArch64/fast-isel-cbz.ll

3

This case could be optimized to use tbz instead.

[olista01]

Thanks, committed revision 220553 with both of your suggested changes.

lib/Target/AArch64/AArch64FastISel.cpp

bool AArch64FastISel::emitCompareAndBranch(const BranchInst *BI) {
   if (!SrcReg)
     return false;
   bool SrcIsKill = hasTrivialKill(LHS);
 
   if (BW == 64 && !Is64Bit)
     SrcReg = fastEmitInst_extractsubreg(MVT::i32, SrcReg, SrcIsKill,
                                         AArch64::sub_32);
 
+  if ((BW == 1 || BW == 8 || BW == 16) && !IsBitTest) {
+    EVT CmpEVT = TLI.getValueType(Ty, SrcIsKill);
+    SrcReg =
+        emitIntExt(CmpEVT.getSimpleVT(), SrcReg, MVT::i32, /*isZExt*/ true);
+  }
+
   // Emit the combined compare and branch instruction.
   SrcReg = constrainOperandRegClass(II, SrcReg,  II.getNumDefs());
   MachineInstrBuilder MIB =
       BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(Opc))
           .addReg(SrcReg, getKillRegState(SrcIsKill));
   if (IsBitTest)
     MIB.addImm(TestBit);
   MIB.addMBB(TBB);

test/CodeGen/AArch64/fast-isel-cbz.ll

 ; RUN: llc -fast-isel -fast-isel-abort -aarch64-atomic-cfg-tidy=0 -verify-machineinstrs -mtriple=aarch64-apple-darwin < %s | FileCheck %s
 
-define i32 @icmp_eq_i1(i1 signext %a) {
+define i32 @icmp_eq_i1(i1 %a) {
 ; CHECK-LABEL: icmp_eq_i1
-; CHECK:       cbz w0, {{LBB.+_2}}
+; CHECK:       and [[REG:w[0-9]+]], w0, #0x1
+; CHECK:       cbz [[REG]], {{LBB.+_2}}
   %1 = icmp eq i1 %a, 0
   br i1 %1, label %bb1, label %bb2
 bb2:
   ret i32 1
 bb1:
   ret i32 0
 }
 
-define i32 @icmp_eq_i8(i8 signext %a) {
+define i32 @icmp_eq_i8(i8 %a) {
 ; CHECK-LABEL: icmp_eq_i8
-; CHECK:       cbz w0, {{LBB.+_2}}
+; CHECK:       uxtb [[REG:w[0-9]+]], w0
+; CHECK:       cbz [[REG]], {{LBB.+_2}}
   %1 = icmp eq i8 %a, 0
   br i1 %1, label %bb1, label %bb2
 bb2:
   ret i32 1
 bb1:
   ret i32 0
 }
 
-define i32 @icmp_eq_i16(i16 signext %a) {
+define i32 @icmp_eq_i16(i16 %a) {
 ; CHECK-LABEL: icmp_eq_i16
-; CHECK:       cbz w0, {{LBB.+_2}}
+; CHECK:       uxth [[REG:w[0-9]+]], w0
+; CHECK:       cbz [[REG]], {{LBB.+_2}}
   %1 = icmp eq i16 %a, 0
   br i1 %1, label %bb1, label %bb2
 bb2:
   ret i32 1
 bb1:
   ret i32 0
 }