[AArch64][GlobalISel] Use ZExtValue for zext(xor) when invert tb(n)z

[guopeilin]

Summary

Currently, we use SExtValue to decide whether to invert tbz or tbnz.
However, for the case zext (xor x, c), we should use ZExt rather
than SExt otherwise we will generate totally opposite branches.

[paquette]

LGTM

[guopeilin]

In D108755#2975020, @paquette wrote:

LGTM

Thanks for reviewing.

[DavidSpickett]

This fixed https://bugs.llvm.org/show_bug.cgi?id=51505, not sure if you were aware of it. Thanks either way!

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

case CmpInst::ICMP_ULE:
     return AArch64CC::LS;
   }
 }
 
 /// Return a register which can be used as a bit to test in a TB(N)Z.
 static Register getTestBitReg(Register Reg, uint64_t &Bit, bool &Invert,
                               MachineRegisterInfo &MRI) {
   assert(Reg.isValid() && "Expected valid register!");
+  bool HasZext = false;
   while (MachineInstr *MI = getDefIgnoringCopies(Reg, MRI)) {
     unsigned Opc = MI->getOpcode();
 
     if (!MI->getOperand(0).isReg() ||
         !MRI.hasOneNonDBGUse(MI->getOperand(0).getReg()))
       break;
 
     // (tbz (any_ext x), b) -> (tbz x, b) if we don't use the extended bits.
     //
     // (tbz (trunc x), b) -> (tbz x, b) is always safe, because the bit number
     // on the truncated x is the same as the bit number on x.
     if (Opc == TargetOpcode::G_ANYEXT || Opc == TargetOpcode::G_ZEXT ||
         Opc == TargetOpcode::G_TRUNC) {
+      if (Opc == TargetOpcode::G_ZEXT)
+        HasZext = true;
+
       Register NextReg = MI->getOperand(1).getReg();
       // Did we find something worth folding?
       if (!NextReg.isValid() || !MRI.hasOneNonDBGUse(NextReg))
         break;
 
       // NextReg is worth folding. Keep looking.
       Reg = NextReg;
       continue;
case TargetOpcode::G_XOR: {
       auto VRegAndVal = getConstantVRegValWithLookThrough(ConstantReg, MRI);
       if (!VRegAndVal) {
         // AND commutes, check the other side for a constant.
         // FIXME: Can we canonicalize the constant so that it's always on the
         // same side at some point earlier?
         std::swap(ConstantReg, TestReg);
         VRegAndVal = getConstantVRegValWithLookThrough(ConstantReg, MRI);
       }
-      if (VRegAndVal)
+      if (VRegAndVal) {
+        if (HasZext)
+          C = VRegAndVal->Value.getZExtValue();
+        else
           C = VRegAndVal->Value.getSExtValue();
+      }
       break;
     }
     case TargetOpcode::G_ASHR:
     case TargetOpcode::G_LSHR:
     case TargetOpcode::G_SHL: {
       TestReg = MI->getOperand(1).getReg();
       auto VRegAndVal =
           getConstantVRegValWithLookThrough(MI->getOperand(2).getReg(), MRI);

llvm/test/CodeGen/AArch64/GlobalISel/opt-fold-xor-tbz-tbnz.mir

bb.0:
     %cmp:gpr(s32) = G_ICMP intpred(eq), %and(s64), %zero
     %cmp_trunc:gpr(s1) = G_TRUNC %cmp(s32)
     G_BRCOND %cmp_trunc(s1), %bb.1
     G_BR %bb.0
   bb.1:
     RET_ReallyLR
 ...
 ---
+name:            dont_flip_eq_zext
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: dont_flip_eq_zext
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32 = COPY $wzr
+  ; CHECK:   [[SUBREG_TO_REG:%[0-9]+]]:gpr64all = SUBREG_TO_REG 0, [[COPY]], %subreg.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[SUBREG_TO_REG]]
+  ; CHECK:   TBNZX [[COPY1]], 63, %bb.1
+  ; CHECK:   B %bb.0
+  ; CHECK: bb.1:
+  ; CHECK:   RET_ReallyLR
+  bb.0:
+    successors: %bb.0(0x40000000), %bb.1(0x40000000)
+
+    %1:gpr(s32) = G_CONSTANT i32 0
+    %3:gpr(s32) = G_CONSTANT i32 -1
+    %4:gpr(s32) = G_XOR %1, %3
+    %5:gpr(s64) = G_ZEXT %4(s32)
+    %15:gpr(s64) = G_CONSTANT i64 0
+    %13:gpr(s32) = G_ICMP intpred(slt), %5(s64), %15
+    %7:gpr(s1) = G_TRUNC %13(s32)
+    G_BRCOND %7(s1), %bb.1
+    G_BR %bb.0
+  bb.1:
+    RET_ReallyLR
+...
+---
 name:            dont_flip_ne
 alignment:       4
 legalized:       true
 regBankSelected: true
 body:             |
   ; CHECK-LABEL: name: dont_flip_ne
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)