[SelectionDAG][RISCV] Treat zext nneg as sext in PromoteIntOp_ZERO_EXTEND if the promoted input is sign extended.

If the zext has the nneg flag and we can prove the promoted input
is sign extended, we can avoid generating an AND that we might not
be able to remove. RISC-V emits a lot of sext_inreg operations during
i32->i64 promotion that makes this likely.

I've restricted this to the case where the promoted type is the same
as the result type so we don't need to create an additional extend.

I've also restricted it to cases where the target has stated a
preference for sext like i32->i64 on RV64. This is largely to avoid
wasting time in computeNumSignBits until we have a test case that
benefits.

Author: topperc

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

@@ -2605,9 +2605,22 @@ SDValue DAGTypeLegalizer::PromoteIntOp_STRICT_UINT_TO_FP(SDNode *N) {
 
 SDValue DAGTypeLegalizer::PromoteIntOp_ZERO_EXTEND(SDNode *N) {
   SDLoc dl(N);
-  SDValue Op = GetPromotedInteger(N->getOperand(0));
-  Op = DAG.getNode(ISD::ANY_EXTEND, dl, N->getValueType(0), Op);
-  return DAG.getZeroExtendInReg(Op, dl, N->getOperand(0).getValueType());
+  SDValue Src = N->getOperand(0);
+  SDValue Op = GetPromotedInteger(Src);
+  EVT VT = N->getValueType(0);
+
+  // If this zext has the nneg flag and the target prefers sext, see if the
+  // promoted input is already sign extended.
+  // TODO: Should we have some way to set nneg on ISD::AND instead?
+  if (N->getFlags().hasNonNeg() && Op.getValueType() == VT &&
+      TLI.isSExtCheaperThanZExt(Src.getValueType(), VT)) {
+    unsigned OpEffectiveBits = DAG.ComputeMaxSignificantBits(Op);
+    if (OpEffectiveBits <= Src.getScalarValueSizeInBits())
+      return Op;
+  }
+
+  Op = DAG.getNode(ISD::ANY_EXTEND, dl, VT, Op);
+  return DAG.getZeroExtendInReg(Op, dl, Src.getValueType());
 }
 
 SDValue DAGTypeLegalizer::PromoteIntOp_VP_ZERO_EXTEND(SDNode *N) {

llvm/test/CodeGen/RISCV/shifts.ll

@@ -832,24 +832,22 @@ define i128 @shl128_shamt32_arith(i128 %a, i32 signext %b) nounwind {
 ;
 ; RV64I-LABEL: shl128_shamt32_arith:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a2, a2, 1
-; RV64I-NEXT:    slli a4, a2, 32
-; RV64I-NEXT:    srli a4, a4, 32
+; RV64I-NEXT:    addiw a4, a2, 1
 ; RV64I-NEXT:    addi a3, a4, -64
+; RV64I-NEXT:    sll a2, a0, a4
 ; RV64I-NEXT:    bltz a3, .LBB17_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    sll a1, a0, a4
+; RV64I-NEXT:    mv a1, a2
 ; RV64I-NEXT:    j .LBB17_3
 ; RV64I-NEXT:  .LBB17_2:
-; RV64I-NEXT:    sll a1, a1, a2
-; RV64I-NEXT:    srli a5, a0, 1
+; RV64I-NEXT:    sll a1, a1, a4
+; RV64I-NEXT:    srli a0, a0, 1
 ; RV64I-NEXT:    not a4, a4
-; RV64I-NEXT:    srl a4, a5, a4
-; RV64I-NEXT:    or a1, a1, a4
+; RV64I-NEXT:    srl a0, a0, a4
+; RV64I-NEXT:    or a1, a1, a0
 ; RV64I-NEXT:  .LBB17_3:
-; RV64I-NEXT:    sll a0, a0, a2
-; RV64I-NEXT:    srai a3, a3, 63
-; RV64I-NEXT:    and a0, a3, a0
+; RV64I-NEXT:    srai a0, a3, 63
+; RV64I-NEXT:    and a0, a0, a2
 ; RV64I-NEXT:    ret
   %c = add i32 %b, 1
   %zext = zext nneg i32 %c to i128