[llvm][AArch64] Allow TB(N)Z to drop signext for sign bit tests.

[fpetrogalli]

Summary

For example if the sign extension is only used in for TBZ, and the value is used elsewhere with a zero extension, this can eliminate a sign extension.

[resistor]

Do we need to handle the same pattern for non IN_REG variations?

[fpetrogalli]

Thank you for the review @resistor. I have implememted the
optimization also looking through a sign_extend node. Tests to cover
this have been added, but only for extending i32 into i64 because
illegal values like i8 are always sign extended inreg.

[peterwaller-arm]

Nit: I think the commit message could be improved such that the benefit is clear? Something referencing the fact you're modifying the behaviour of TBZ.

Suggestion: "Allow TB(N)Z to drop signext for sign bit tests"

With the body explaining why/circumstances this helps: "For example if the sign extension is only used in for TBZ, and the value is used elsewhere with a zero extension, this can eliminate a sign extension".

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5810

sp. returns

[fpetrogalli]

Fix typo in comment.

[samparker]

Thanks, LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerGlobalTLSAddress(SDValue Op,
   if (Subtarget->isTargetELF())
     return LowerELFGlobalTLSAddress(Op, DAG);
   if (Subtarget->isTargetWindows())
     return LowerWindowsGlobalTLSAddress(Op, DAG);
 
   llvm_unreachable("Unexpected platform trying to use TLS");
 }
 
+// Looks through \param Val to determine the bit that can be used to
+// check the sign of the value. It returns the unextended value and
+// the sign bit position.
+std::pair<SDValue, uint64_t> lookThroughSignExtension(SDValue Val) {
+  if (Val.getOpcode() == ISD::SIGN_EXTEND_INREG)
+    return {Val.getOperand(0),
+            cast<VTSDNode>(Val.getOperand(1))->getVT().getFixedSizeInBits() -
+                1};
+
+  if (Val.getOpcode() == ISD::SIGN_EXTEND)
+    return {Val.getOperand(0),
+            Val.getOperand(0)->getValueType(0).getFixedSizeInBits() - 1};
+
+  return {Val, Val.getValueSizeInBits() - 1};
+}
+
 SDValue AArch64TargetLowering::LowerBR_CC(SDValue Op, SelectionDAG &DAG) const {
   SDValue Chain = Op.getOperand(0);
   ISD::CondCode CC = cast<CondCodeSDNode>(Op.getOperand(1))->get();
   SDValue LHS = Op.getOperand(2);
   SDValue RHS = Op.getOperand(3);
   SDValue Dest = Op.getOperand(4);
   SDLoc dl(Op);
 
if (RHSC && RHSC->getZExtValue() == 0 && ProduceNonFlagSettingCondBr) {
                              Dest);
         }
 
         return DAG.getNode(AArch64ISD::CBNZ, dl, MVT::Other, Chain, LHS, Dest);
       } else if (CC == ISD::SETLT && LHS.getOpcode() != ISD::AND) {
         // Don't combine AND since emitComparison converts the AND to an ANDS
         // (a.k.a. TST) and the test in the test bit and branch instruction
         // becomes redundant.  This would also increase register pressure.
-        uint64_t Mask = LHS.getValueSizeInBits() - 1;
+        uint64_t SignBitPos;
+        std::tie(LHS, SignBitPos) = lookThroughSignExtension(LHS);
         return DAG.getNode(AArch64ISD::TBNZ, dl, MVT::Other, Chain, LHS,
-                           DAG.getConstant(Mask, dl, MVT::i64), Dest);
+                           DAG.getConstant(SignBitPos, dl, MVT::i64), Dest);
       }
     }
     if (RHSC && RHSC->getSExtValue() == -1 && CC == ISD::SETGT &&
         LHS.getOpcode() != ISD::AND && ProduceNonFlagSettingCondBr) {
       // Don't combine AND since emitComparison converts the AND to an ANDS
       // (a.k.a. TST) and the test in the test bit and branch instruction
       // becomes redundant.  This would also increase register pressure.
-      uint64_t Mask = LHS.getValueSizeInBits() - 1;
+      uint64_t SignBitPos;
+      std::tie(LHS, SignBitPos) = lookThroughSignExtension(LHS);
       return DAG.getNode(AArch64ISD::TBZ, dl, MVT::Other, Chain, LHS,
-                         DAG.getConstant(Mask, dl, MVT::i64), Dest);
+                         DAG.getConstant(SignBitPos, dl, MVT::i64), Dest);
     }
 
     SDValue CCVal;
     SDValue Cmp = getAArch64Cmp(LHS, RHS, CC, CCVal, DAG, dl);
     return DAG.getNode(AArch64ISD::BRCOND, dl, MVT::Other, Chain, Dest, CCVal,
                        Cmp);
   }
 

llvm/test/CodeGen/AArch64/check-sign-bit-before-extension.ll

+; RUN: llc -mtriple aarch64-gnu-linux -o -  -asm-verbose=0 %s | FileCheck %s
+
+; These tests make sure that the `cmp` instruction is rendered with an
+; instruction that checks the sign bit of the original unextended data
+; (%in) instead of the sign bit of the sign extended one that is
+; created by the type legalization process.
+;
+; The tests are subdivided in tests that determine the sign bit
+; looking through a `sign_extend_inreg` and tests that determine the
+; sign bit looking through a `sign_extend`.
+
+; CHECK-LABEL: f_i8_sign_extend_inreg:
+; CHECK: tbnz w0, #7, .LBB
+define i32 @f_i8_sign_extend_inreg(i8 %in, i32 %a, i32 %b) nounwind {
+entry:
+  %cmp = icmp sgt i8 %in, -1
+  %ext = zext i8 %in to i32
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i32 %ext, %a
+  ret i32 %retA
+
+B:
+  %retB = add i32 %ext, %b
+  ret i32 %retB
+}
+
+; CHECK-LABEL: f_i16_sign_extend_inreg:
+; CHECK: tbnz w0, #15, .LBB
+define i32 @f_i16_sign_extend_inreg(i16 %in, i32 %a, i32 %b) nounwind {
+entry:
+  %cmp = icmp sgt i16 %in, -1
+  %ext = zext i16 %in to i32
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i32 %ext, %a
+  ret i32 %retA
+
+B:
+  %retB = add i32 %ext, %b
+  ret i32 %retB
+}
+
+; CHECK-LABEL: f_i32_sign_extend_inreg:
+; CHECK: tbnz w0, #31, .LBB
+define i64 @f_i32_sign_extend_inreg(i32 %in, i64 %a, i64 %b) nounwind {
+entry:
+  %cmp = icmp sgt i32 %in, -1
+  %ext = zext i32 %in to i64
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i64 %ext, %a
+  ret i64 %retA
+
+B:
+  %retB = add i64 %ext, %b
+  ret i64 %retB
+}
+
+; CHECK-LABEL: g_i8_sign_extend_inreg:
+; CHECK: tbnz w0, #7, .LBB
+define i32 @g_i8_sign_extend_inreg(i8 %in, i32 %a, i32 %b) nounwind {
+entry:
+  %cmp = icmp slt i8 %in, 0
+  %ext = zext i8 %in to i32
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i32 %ext, %a
+  ret i32 %retA
+
+B:
+  %retB = add i32 %ext, %b
+  ret i32 %retB
+}
+
+; CHECK-LABEL: g_i16_sign_extend_inreg:
+; CHECK: tbnz w0, #15, .LBB
+define i32 @g_i16_sign_extend_inreg(i16 %in, i32 %a, i32 %b) nounwind {
+entry:
+  %cmp = icmp slt i16 %in, 0
+  %ext = zext i16 %in to i32
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i32 %ext, %a
+  ret i32 %retA
+
+B:
+  %retB = add i32 %ext, %b
+  ret i32 %retB
+}
+
+; CHECK-LABEL: g_i32_sign_extend_inreg:
+; CHECK: tbnz w0, #31, .LBB
+define i64 @g_i32_sign_extend_inreg(i32 %in, i64 %a, i64 %b) nounwind {
+entry:
+  %cmp = icmp slt i32 %in, 0
+  %ext = zext i32 %in to i64
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i64 %ext, %a
+  ret i64 %retA
+
+B:
+  %retB = add i64 %ext, %b
+  ret i64 %retB
+}
+
+; CHECK-LABEL: f_i32_sign_extend_i64:
+; CHECK: tbnz w0, #31, .LBB
+define i64 @f_i32_sign_extend_i64(i32 %in, i64 %a, i64 %b) nounwind {
+entry:
+  %inext = sext i32 %in to i64
+  %cmp = icmp sgt i64 %inext, -1
+  %ext = zext i32 %in to i64
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i64 %ext, %a
+  ret i64 %retA
+
+B:
+  %retB = add i64 %ext, %b
+  ret i64 %retB
+}
+
+; CHECK-LABEL: g_i32_sign_extend_i64:
+; CHECK: tbnz w0, #31, .LBB
+define i64 @g_i32_sign_extend_i64(i32 %in, i64 %a, i64 %b) nounwind {
+entry:
+  %inext = sext i32 %in to i64
+  %cmp = icmp slt i64 %inext, 0
+  %ext = zext i32 %in to i64
+  br i1 %cmp, label %A, label %B
+
+A:
+  %retA = add i64 %ext, %a
+  ret i64 %retA
+
+B:
+  %retB = add i64 %ext, %b
+  ret i64 %retB
+}