[RISCV] Fold Zba-expanded (mul (shr exact X, C1), C2)

Author: pfusik

llvm/include/llvm/CodeGen/SDPatternMatch.h

@@ -903,6 +903,11 @@ template <typename LHS, typename RHS>
 inline BinaryOpc_match<LHS, RHS> m_Srl(const LHS &L, const RHS &R) {
   return BinaryOpc_match<LHS, RHS>(ISD::SRL, L, R);
 }
+template <typename LHS, typename RHS>
+inline auto m_ExactSr(const LHS &L, const RHS &R) {
+  return m_AnyOf(BinaryOpc_match<LHS, RHS>(ISD::SRA, L, R, SDNodeFlags::Exact),
+                 BinaryOpc_match<LHS, RHS>(ISD::SRL, L, R, SDNodeFlags::Exact));
+}
 
 template <typename LHS, typename RHS>
 inline BinaryOpc_match<LHS, RHS> m_Rotl(const LHS &L, const RHS &R) {

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -16829,10 +16829,13 @@ static SDValue getShlAddShlAdd(SDNode *N, SelectionDAG &DAG, unsigned ShX,
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
   SDValue X = N->getOperand(0);
-  // Put the shift first if we can fold a zext into the shift forming a slli.uw.
+  // Put the shift first if we can fold:
+  // a. a zext into the shift forming a slli.uw
+  // b. an exact shift right forming one shorter shift or no shift at all
   using namespace SDPatternMatch;
   if (Shift != 0 &&
-      sd_match(X, m_And(m_Value(), m_SpecificInt(UINT64_C(0xffffffff))))) {
+      sd_match(X, m_AnyOf(m_And(m_Value(), m_SpecificInt(UINT64_C(0xffffffff))),
+                          m_ExactSr(m_Value(), m_ConstInt())))) {
     X = DAG.getNode(ISD::SHL, DL, VT, X, DAG.getConstant(Shift, DL, VT));
     Shift = 0;
   }

llvm/test/CodeGen/RISCV/rv64zba.ll

@@ -5026,16 +5026,12 @@ define i64 @exactashr1mul6(i64 %a) {
 ;
 ; RV64ZBA-LABEL: exactashr1mul6:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    srli a0, a0, 1
 ; RV64ZBA-NEXT:    sh1add a0, a0, a0
-; RV64ZBA-NEXT:    slli a0, a0, 1
 ; RV64ZBA-NEXT:    ret
 ;
 ; RV64XANDESPERF-LABEL: exactashr1mul6:
 ; RV64XANDESPERF:       # %bb.0:
-; RV64XANDESPERF-NEXT:    srli a0, a0, 1
 ; RV64XANDESPERF-NEXT:    nds.lea.h a0, a0, a0
-; RV64XANDESPERF-NEXT:    slli a0, a0, 1
 ; RV64XANDESPERF-NEXT:    ret
   %c = ashr exact i64 %a, 1
   %d = mul i64 %c, 6
@@ -5052,18 +5048,16 @@ define i64 @exactlshr3mul22(i64 %a) {
 ;
 ; RV64ZBA-LABEL: exactlshr3mul22:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    srli a0, a0, 3
+; RV64ZBA-NEXT:    srli a0, a0, 2
 ; RV64ZBA-NEXT:    sh2add a1, a0, a0
 ; RV64ZBA-NEXT:    sh1add a0, a1, a0
-; RV64ZBA-NEXT:    slli a0, a0, 1
 ; RV64ZBA-NEXT:    ret
 ;
 ; RV64XANDESPERF-LABEL: exactlshr3mul22:
 ; RV64XANDESPERF:       # %bb.0:
-; RV64XANDESPERF-NEXT:    srli a0, a0, 3
+; RV64XANDESPERF-NEXT:    srli a0, a0, 2
 ; RV64XANDESPERF-NEXT:    nds.lea.w a1, a0, a0
 ; RV64XANDESPERF-NEXT:    nds.lea.h a0, a0, a1
-; RV64XANDESPERF-NEXT:    slli a0, a0, 1
 ; RV64XANDESPERF-NEXT:    ret
   %c = lshr exact i64 %a, 3
   %d = mul i64 %c, 22
@@ -5080,16 +5074,14 @@ define i64 @exactashr1mul36(i64 %a) {
 ;
 ; RV64ZBA-LABEL: exactashr1mul36:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    srli a0, a0, 1
+; RV64ZBA-NEXT:    slli a0, a0, 1
 ; RV64ZBA-NEXT:    sh3add a0, a0, a0
-; RV64ZBA-NEXT:    slli a0, a0, 2
 ; RV64ZBA-NEXT:    ret
 ;
 ; RV64XANDESPERF-LABEL: exactashr1mul36:
 ; RV64XANDESPERF:       # %bb.0:
-; RV64XANDESPERF-NEXT:    srli a0, a0, 1
+; RV64XANDESPERF-NEXT:    slli a0, a0, 1
 ; RV64XANDESPERF-NEXT:    nds.lea.d a0, a0, a0
-; RV64XANDESPERF-NEXT:    slli a0, a0, 2
 ; RV64XANDESPERF-NEXT:    ret
   %c = ashr exact i64 %a, 1
   %d = mul i64 %c, 36