[llvm] [RISCV][DAGCombine] Combine `sext_inreg (shl X, Y), i32` into `sllw X, Y` (PR #111101)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/111101

Alive2: https://alive2.llvm.org/ce/z/ncf36D


>From 1075d4c6f9582434395f2fc75654b664ec96bfaa Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 4 Oct 2024 12:31:06 +0800
Subject: [PATCH 1/2] [RISCV][DAGCombine] Add pre-commit tests. NFC.

---
 llvm/test/CodeGen/RISCV/rv64zbs.ll | 42 ++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/llvm/test/CodeGen/RISCV/rv64zbs.ll b/llvm/test/CodeGen/RISCV/rv64zbs.ll
index 577005a051b204..68c06c3b810b8f 100644
--- a/llvm/test/CodeGen/RISCV/rv64zbs.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zbs.ll
@@ -1210,3 +1210,45 @@ define i1 @icmp_eq_nonpow2(i32 signext %x) nounwind {
   %cmp = icmp eq i32 %x, 32767
   ret i1 %cmp
 }
+
+define signext i32 @fold_sextinreg_shl_to_sllw(i64 %x) nounwind {
+; RV64I-LABEL: fold_sextinreg_shl_to_sllw:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    andi a0, a0, 31
+; RV64I-NEXT:    li a1, 1
+; RV64I-NEXT:    sll a0, a1, a0
+; RV64I-NEXT:    sext.w a0, a0
+; RV64I-NEXT:    ret
+;
+; RV64ZBS-LABEL: fold_sextinreg_shl_to_sllw:
+; RV64ZBS:       # %bb.0: # %entry
+; RV64ZBS-NEXT:    andi a0, a0, 31
+; RV64ZBS-NEXT:    bset a0, zero, a0
+; RV64ZBS-NEXT:    sext.w a0, a0
+; RV64ZBS-NEXT:    ret
+entry:
+  %mask = and i64 %x, 31
+  %shl = shl i64 1, %mask
+  %trunc = trunc i64 %shl to i32
+  ret i32 %trunc
+}
+
+define signext i32 @fold_sextinreg_shl_to_sllw_large_shamt(i64 %x) nounwind {
+; RV64I-LABEL: fold_sextinreg_shl_to_sllw_large_shamt:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    li a1, 1
+; RV64I-NEXT:    sll a0, a1, a0
+; RV64I-NEXT:    sext.w a0, a0
+; RV64I-NEXT:    ret
+;
+; RV64ZBS-LABEL: fold_sextinreg_shl_to_sllw_large_shamt:
+; RV64ZBS:       # %bb.0: # %entry
+; RV64ZBS-NEXT:    bset a0, zero, a0
+; RV64ZBS-NEXT:    sext.w a0, a0
+; RV64ZBS-NEXT:    ret
+entry:
+  %mask = and i64 %x, 63
+  %shl = shl i64 1, %mask
+  %trunc = trunc i64 %shl to i32
+  ret i32 %trunc
+}

>From e369d3762317112a963c665b26ed3243e9b450db Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 4 Oct 2024 13:26:46 +0800
Subject: [PATCH 2/2] [RISCV][DAGCombine] Combine `sext_inreg (shl X, Y), i32
 into sllw X, Y`

---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 16 ++++++++++++----
 llvm/test/CodeGen/RISCV/rv64zbs.ll          | 19 +++++--------------
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 3e06960092cf32..463887b8b55e61 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -1489,6 +1489,7 @@ RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
                        ISD::INTRINSIC_WO_CHAIN, ISD::ADD, ISD::SUB, ISD::MUL,
                        ISD::AND, ISD::OR, ISD::XOR, ISD::SETCC, ISD::SELECT});
   setTargetDAGCombine(ISD::SRA);
+  setTargetDAGCombine(ISD::SIGN_EXTEND_INREG);
 
   if (Subtarget.hasStdExtFOrZfinx())
     setTargetDAGCombine({ISD::FADD, ISD::FMAXNUM, ISD::FMINNUM, ISD::FMUL});
@@ -1502,8 +1503,7 @@ RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
 
   if (Subtarget.hasStdExtZbkb())
     setTargetDAGCombine(ISD::BITREVERSE);
-  if (Subtarget.hasStdExtZfhminOrZhinxmin())
-    setTargetDAGCombine(ISD::SIGN_EXTEND_INREG);
+
   if (Subtarget.hasStdExtFOrZfinx())
     setTargetDAGCombine({ISD::ZERO_EXTEND, ISD::FP_TO_SINT, ISD::FP_TO_UINT,
                          ISD::FP_TO_SINT_SAT, ISD::FP_TO_UINT_SAT});
@@ -14456,15 +14456,23 @@ performSIGN_EXTEND_INREGCombine(SDNode *N, SelectionDAG &DAG,
                                 const RISCVSubtarget &Subtarget) {
   SDValue Src = N->getOperand(0);
   EVT VT = N->getValueType(0);
+  EVT SrcVT = cast<VTSDNode>(N->getOperand(1))->getVT();
+  unsigned Opc = Src.getOpcode();
 
   // Fold (sext_inreg (fmv_x_anyexth X), i16) -> (fmv_x_signexth X)
   // Don't do this with Zhinx. We need to explicitly sign extend the GPR.
-  if (Src.getOpcode() == RISCVISD::FMV_X_ANYEXTH &&
-      cast<VTSDNode>(N->getOperand(1))->getVT().bitsGE(MVT::i16) &&
+  if (Opc == RISCVISD::FMV_X_ANYEXTH && SrcVT.bitsGE(MVT::i16) &&
       Subtarget.hasStdExtZfhmin())
     return DAG.getNode(RISCVISD::FMV_X_SIGNEXTH, SDLoc(N), VT,
                        Src.getOperand(0));
 
+  // Fold (sext_inreg (shl X, Y), i32) -> (sllw X, Y) iff Y u< 32
+  if (Opc == ISD::SHL && Subtarget.is64Bit() && SrcVT == MVT::i32 &&
+      VT == MVT::i64 && !isa<ConstantSDNode>(Src.getOperand(1)) &&
+      DAG.computeKnownBits(Src.getOperand(1)).countMaxActiveBits() <= 5)
+    return DAG.getNode(RISCVISD::SLLW, SDLoc(N), VT, Src.getOperand(0),
+                       Src.getOperand(1));
+
   return SDValue();
 }
 
diff --git a/llvm/test/CodeGen/RISCV/rv64zbs.ll b/llvm/test/CodeGen/RISCV/rv64zbs.ll
index 68c06c3b810b8f..c1e1e16d0d4ae2 100644
--- a/llvm/test/CodeGen/RISCV/rv64zbs.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zbs.ll
@@ -1212,20 +1212,11 @@ define i1 @icmp_eq_nonpow2(i32 signext %x) nounwind {
 }
 
 define signext i32 @fold_sextinreg_shl_to_sllw(i64 %x) nounwind {
-; RV64I-LABEL: fold_sextinreg_shl_to_sllw:
-; RV64I:       # %bb.0: # %entry
-; RV64I-NEXT:    andi a0, a0, 31
-; RV64I-NEXT:    li a1, 1
-; RV64I-NEXT:    sll a0, a1, a0
-; RV64I-NEXT:    sext.w a0, a0
-; RV64I-NEXT:    ret
-;
-; RV64ZBS-LABEL: fold_sextinreg_shl_to_sllw:
-; RV64ZBS:       # %bb.0: # %entry
-; RV64ZBS-NEXT:    andi a0, a0, 31
-; RV64ZBS-NEXT:    bset a0, zero, a0
-; RV64ZBS-NEXT:    sext.w a0, a0
-; RV64ZBS-NEXT:    ret
+; CHECK-LABEL: fold_sextinreg_shl_to_sllw:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a1, 1
+; CHECK-NEXT:    sllw a0, a1, a0
+; CHECK-NEXT:    ret
 entry:
   %mask = and i64 %x, 31
   %shl = shl i64 1, %mask