From 39c406a58f96c0f0c3785220988720c8a8affb40 Mon Sep 17 00:00:00 2001
From: guopeilin <guopeilin1@huawei.com>
Date: Mon, 6 Sep 2021 11:11:23 +0800
Subject: [PATCH] [AArch64][GlobalISel] Use ZExtValue for zext(xor) when invert
 tb(n)z

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.

Reviewed By: paquette

Differential Revision: https://reviews.llvm.org/D108755

(cherry picked from commit 5f48c144c58f6d23e850a1978a6fe05887103b17)
---
 .../GISel/AArch64InstructionSelector.cpp      | 12 +++++--
 .../GlobalISel/opt-fold-xor-tbz-tbnz.mir      | 32 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index a98248438e4..a17bd1b14f4 100644
--- a/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -1301,6 +1301,7 @@ static AArch64CC::CondCode changeICMPPredToAArch64CC(CmpInst::Predicate P) {
 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();
 
@@ -1314,6 +1315,9 @@ static Register getTestBitReg(Register Reg, uint64_t &Bit, bool &Invert,
     // 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))
@@ -1342,8 +1346,12 @@ static Register getTestBitReg(Register Reg, uint64_t &Bit, bool &Invert,
         std::swap(ConstantReg, TestReg);
         VRegAndVal = getConstantVRegValWithLookThrough(ConstantReg, MRI);
       }
-      if (VRegAndVal)
-        C = VRegAndVal->Value.getSExtValue();
+      if (VRegAndVal) {
+        if (HasZext)
+          C = VRegAndVal->Value.getZExtValue();
+        else
+          C = VRegAndVal->Value.getSExtValue();
+      }
       break;
     }
     case TargetOpcode::G_ASHR:
diff --git a/test/CodeGen/AArch64/GlobalISel/opt-fold-xor-tbz-tbnz.mir b/test/CodeGen/AArch64/GlobalISel/opt-fold-xor-tbz-tbnz.mir
index 8e19ba41b2c..53ea6830fdc 100644
--- a/test/CodeGen/AArch64/GlobalISel/opt-fold-xor-tbz-tbnz.mir
+++ b/test/CodeGen/AArch64/GlobalISel/opt-fold-xor-tbz-tbnz.mir
@@ -117,6 +117,38 @@ body:             |
     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