[llvm] [LLVM][AArch64] Optimize sign bit tests with TST instruction for SIGN_EXTEND patterns (PR #158061)

[David Green via llvm-commits]

================
@@ -11630,6 +11630,48 @@ SDValue AArch64TargetLowering::LowerSELECT_CC(
       return DAG.getNode(ISD::AND, DL, VT, LHS, Shift);
     }
 
+    // Check for sign bit test patterns that can use TST optimization.
+    // (SELECT_CC setlt, singn_extend_inreg, 0, tval, fval)
+    //                          -> TST %operand, sign_bit; CSEL
+    // (SELECT_CC setlt, singn_extend, 0, tval, fval)
+    //                          -> TST %operand, sign_bit; CSEL
+    if (CC == ISD::SETLT && RHSC && RHSC->isZero() && LHS.hasOneUse() &&
+        (LHS.getOpcode() == ISD::SIGN_EXTEND_INREG ||
+         LHS.getOpcode() == ISD::SIGN_EXTEND)) {
+
+      SDValue OriginalVal = LHS.getOperand(0);
+      EVT OriginalVT = LHS.getOpcode() == ISD::SIGN_EXTEND_INREG
+                           ? cast<VTSDNode>(LHS.getOperand(1))->getVT()
+                           : OriginalVal.getValueType();
+
+      // Apply TST optimization for integer types
+      if (OriginalVT.isInteger()) {
+        // Calculate the sign bit for the original type
+        unsigned BitWidth = OriginalVT.getSizeInBits();
+        APInt SignBit = APInt::getSignedMinValue(BitWidth);
+        EVT TestVT = (BitWidth <= 32) ? MVT::i32 : MVT::i64;
+        unsigned TestBitWidth = TestVT.getSizeInBits();
+        if (BitWidth < TestBitWidth) {
+          SignBit = SignBit.zext(TestBitWidth);
+        }
+
+        SDValue SignBitConst = DAG.getConstant(SignBit, DL, TestVT);
+        SDValue TestOperand = OriginalVal;
+        if (OriginalVal.getValueType() != TestVT) {
+          TestOperand = DAG.getNode(ISD::ZERO_EXTEND, DL, TestVT, OriginalVal);
+        }
+
+        SDValue TST =
+            DAG.getNode(AArch64ISD::ANDS, DL, DAG.getVTList(TestVT, MVT::i32),
+                        TestOperand, SignBitConst);
+
+        SDValue Flags = TST.getValue(1);
+        return DAG.getNode(AArch64ISD::CSEL, DL, TVal.getValueType(), TVal,
+                           FVal, DAG.getConstant(AArch64CC::MI, DL, MVT::i32),
----------------
davemgreen wrote:

I think you need to update the condition flag to test if the result of the ANDS is not zero (that original bit is not set). MI checks N which is only set if the top bit of the whole register is zero.

https://github.com/llvm/llvm-project/pull/158061