[LoongArch] Perform SELECT_CC combine (#155994)

Fold `((srl (and X, 1<<C), C), 0, eq/ne)` -> `((shl X, GRLen-1-C), 0, ge/lt)`

Author: heiher

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

@@ -5282,7 +5282,7 @@ static SDValue performBITREV_WCombine(SDNode *N, SelectionDAG &DAG,
                      Src.getOperand(0));
 }
 
-// Perform combines for BR_CC conditions.
+// Perform common combines for BR_CC and SELECT_CC conditions.
 static bool combine_CC(SDValue &LHS, SDValue &RHS, SDValue &CC, const SDLoc &DL,
                        SelectionDAG &DAG, const LoongArchSubtarget &Subtarget) {
   ISD::CondCode CCVal = cast<CondCodeSDNode>(CC)->get();
@@ -5319,6 +5319,29 @@ static bool combine_CC(SDValue &LHS, SDValue &RHS, SDValue &CC, const SDLoc &DL,
     return true;
   }
 
+  // Fold ((srl (and X, 1<<C), C), 0, eq/ne) -> ((shl X, GRLen-1-C), 0, ge/lt)
+  if (isNullConstant(RHS) && LHS.getOpcode() == ISD::SRL && LHS.hasOneUse() &&
+      LHS.getOperand(1).getOpcode() == ISD::Constant) {
+    SDValue LHS0 = LHS.getOperand(0);
+    if (LHS0.getOpcode() == ISD::AND &&
+        LHS0.getOperand(1).getOpcode() == ISD::Constant) {
+      uint64_t Mask = LHS0.getConstantOperandVal(1);
+      uint64_t ShAmt = LHS.getConstantOperandVal(1);
+      if (isPowerOf2_64(Mask) && Log2_64(Mask) == ShAmt) {
+        CCVal = CCVal == ISD::SETEQ ? ISD::SETGE : ISD::SETLT;
+        CC = DAG.getCondCode(CCVal);
+
+        ShAmt = LHS.getValueSizeInBits() - 1 - ShAmt;
+        LHS = LHS0.getOperand(0);
+        if (ShAmt != 0)
+          LHS =
+              DAG.getNode(ISD::SHL, DL, LHS.getValueType(), LHS0.getOperand(0),
+                          DAG.getConstant(ShAmt, DL, LHS.getValueType()));
+        return true;
+      }
+    }
+  }
+
   // (X, 1, setne) -> (X, 0, seteq) if we can prove X is 0/1.
   // This can occur when legalizing some floating point comparisons.
   APInt Mask = APInt::getBitsSetFrom(LHS.getValueSizeInBits(), 1);
@@ -5347,6 +5370,57 @@ static SDValue performBR_CCCombine(SDNode *N, SelectionDAG &DAG,
   return SDValue();
 }
 
+static SDValue performSELECT_CCCombine(SDNode *N, SelectionDAG &DAG,
+                                       TargetLowering::DAGCombinerInfo &DCI,
+                                       const LoongArchSubtarget &Subtarget) {
+  // Transform
+  SDValue LHS = N->getOperand(0);
+  SDValue RHS = N->getOperand(1);
+  SDValue CC = N->getOperand(2);
+  ISD::CondCode CCVal = cast<CondCodeSDNode>(CC)->get();
+  SDValue TrueV = N->getOperand(3);
+  SDValue FalseV = N->getOperand(4);
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+
+  // If the True and False values are the same, we don't need a select_cc.
+  if (TrueV == FalseV)
+    return TrueV;
+
+  // (select (x < 0), y, z)  -> x >> (GRLEN - 1) & (y - z) + z
+  // (select (x >= 0), y, z) -> x >> (GRLEN - 1) & (z - y) + y
+  if (isa<ConstantSDNode>(TrueV) && isa<ConstantSDNode>(FalseV) &&
+      isNullConstant(RHS) &&
+      (CCVal == ISD::CondCode::SETLT || CCVal == ISD::CondCode::SETGE)) {
+    if (CCVal == ISD::CondCode::SETGE)
+      std::swap(TrueV, FalseV);
+
+    int64_t TrueSImm = cast<ConstantSDNode>(TrueV)->getSExtValue();
+    int64_t FalseSImm = cast<ConstantSDNode>(FalseV)->getSExtValue();
+    // Only handle simm12, if it is not in this range, it can be considered as
+    // register.
+    if (isInt<12>(TrueSImm) && isInt<12>(FalseSImm) &&
+        isInt<12>(TrueSImm - FalseSImm)) {
+      SDValue SRA =
+          DAG.getNode(ISD::SRA, DL, VT, LHS,
+                      DAG.getConstant(Subtarget.getGRLen() - 1, DL, VT));
+      SDValue AND =
+          DAG.getNode(ISD::AND, DL, VT, SRA,
+                      DAG.getSignedConstant(TrueSImm - FalseSImm, DL, VT));
+      return DAG.getNode(ISD::ADD, DL, VT, AND, FalseV);
+    }
+
+    if (CCVal == ISD::CondCode::SETGE)
+      std::swap(TrueV, FalseV);
+  }
+
+  if (combine_CC(LHS, RHS, CC, DL, DAG, Subtarget))
+    return DAG.getNode(LoongArchISD::SELECT_CC, DL, N->getValueType(0),
+                       {LHS, RHS, CC, TrueV, FalseV});
+
+  return SDValue();
+}
+
 template <unsigned N>
 static SDValue legalizeIntrinsicImmArg(SDNode *Node, unsigned ImmOp,
                                        SelectionDAG &DAG,
@@ -6041,6 +6115,8 @@ SDValue LoongArchTargetLowering::PerformDAGCombine(SDNode *N,
     return performBITREV_WCombine(N, DAG, DCI, Subtarget);
   case LoongArchISD::BR_CC:
     return performBR_CCCombine(N, DAG, DCI, Subtarget);
+  case LoongArchISD::SELECT_CC:
+    return performSELECT_CCCombine(N, DAG, DCI, Subtarget);
   case ISD::INTRINSIC_WO_CHAIN:
     return performINTRINSIC_WO_CHAINCombine(N, DAG, DCI, Subtarget);
   case LoongArchISD::MOVGR2FR_W_LA64:

llvm/test/CodeGen/LoongArch/bittest.ll

@@ -358,10 +358,9 @@ define signext i32 @bit_10_z_select_i32(i32 signext %a, i32 signext %b, i32 sign
 define signext i32 @bit_10_nz_select_i32(i32 signext %a, i32 signext %b, i32 signext %c) {
 ; LA32-LABEL: bit_10_nz_select_i32:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    andi $a0, $a0, 1024
-; LA32-NEXT:    srli.w $a3, $a0, 10
+; LA32-NEXT:    slli.w $a3, $a0, 21
 ; LA32-NEXT:    move $a0, $a1
-; LA32-NEXT:    bne $a3, $zero, .LBB16_2
+; LA32-NEXT:    bltz $a3, .LBB16_2
 ; LA32-NEXT:  # %bb.1:
 ; LA32-NEXT:    move $a0, $a2
 ; LA32-NEXT:  .LBB16_2:
@@ -408,10 +407,9 @@ define signext i32 @bit_11_z_select_i32(i32 signext %a, i32 signext %b, i32 sign
 define signext i32 @bit_11_nz_select_i32(i32 signext %a, i32 signext %b, i32 signext %c) {
 ; LA32-LABEL: bit_11_nz_select_i32:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    andi $a0, $a0, 2048
-; LA32-NEXT:    srli.w $a3, $a0, 11
+; LA32-NEXT:    slli.w $a3, $a0, 20
 ; LA32-NEXT:    move $a0, $a1
-; LA32-NEXT:    bne $a3, $zero, .LBB18_2
+; LA32-NEXT:    bltz $a3, .LBB18_2
 ; LA32-NEXT:  # %bb.1:
 ; LA32-NEXT:    move $a0, $a2
 ; LA32-NEXT:  .LBB18_2:
@@ -459,11 +457,9 @@ define signext i32 @bit_20_z_select_i32(i32 signext %a, i32 signext %b, i32 sign
 define signext i32 @bit_20_nz_select_i32(i32 signext %a, i32 signext %b, i32 signext %c) {
 ; LA32-LABEL: bit_20_nz_select_i32:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    lu12i.w $a3, 256
-; LA32-NEXT:    and $a0, $a0, $a3
-; LA32-NEXT:    srli.w $a3, $a0, 20
+; LA32-NEXT:    slli.w $a3, $a0, 11
 ; LA32-NEXT:    move $a0, $a1
-; LA32-NEXT:    bne $a3, $zero, .LBB20_2
+; LA32-NEXT:    bltz $a3, .LBB20_2
 ; LA32-NEXT:  # %bb.1:
 ; LA32-NEXT:    move $a0, $a2
 ; LA32-NEXT:  .LBB20_2:

llvm/test/CodeGen/LoongArch/select-const.ll

@@ -305,12 +305,9 @@ define i32 @select_ne_10001_10002(i32 signext %a, i32 signext %b) {
 define i32 @select_slt_zero_constant1_constant2(i32 signext %x) {
 ; LA32-LABEL: select_slt_zero_constant1_constant2:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    move $a1, $a0
-; LA32-NEXT:    ori $a0, $zero, 7
-; LA32-NEXT:    bltz $a1, .LBB16_2
-; LA32-NEXT:  # %bb.1:
-; LA32-NEXT:    addi.w $a0, $zero, -3
-; LA32-NEXT:  .LBB16_2:
+; LA32-NEXT:    srai.w $a0, $a0, 31
+; LA32-NEXT:    andi $a0, $a0, 10
+; LA32-NEXT:    addi.w $a0, $a0, -3
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: select_slt_zero_constant1_constant2: