[RISCV] Add computeKnownBitsForTargetNode for RISCVISD::SRAW. (#156191)

This node reads the lower 32 bits, shifts it right arithmetically
then sign extends to i64. If we know some of the lower 32 bits we
can propagate that information.
    
For the test case I had to find something that didn't get optimized
before type legalizaton and didn't get type legalized to a sign
extended value. The bswap gets type legalized to (lshr (bswap), 32).

2 files changed, 34 insertions, 0 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 9c26801..9115c13 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -21379,6 +21379,15 @@ void RISCVTargetLowering::computeKnownBitsForTargetNode(const SDValue Op,
     Known = Known.sext(BitWidth);
     break;
   }
+  case RISCVISD::SRAW: {
+    KnownBits Known2;
+    Known = DAG.computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
+    Known2 = DAG.computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
+    Known = KnownBits::ashr(Known.trunc(32), Known2.trunc(5).zext(32));
+    // Restore the original width by sign extending.
+    Known = Known.sext(BitWidth);
+    break;
+  }
   case RISCVISD::CTZW: {
     KnownBits Known2 = DAG.computeKnownBits(Op.getOperand(0), Depth + 1);
     unsigned PossibleTZ = Known2.trunc(32).countMaxTrailingZeros();
diff --git a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
index 90735d8..b23f4f5 100644
--- a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
+++ b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
@@ -235,3 +235,28 @@ define signext i32 @test15(i32 signext %x, i32 signext %y) {
   %c = shl i32 %b, 1
   ret i32 %c
 }
+
+define signext i32 @test16(i32 signext %x, i32 signext %y) {
+; RV64I-LABEL: test16:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    ori a2, a0, 192
+; RV64I-NEXT:    srli a3, a0, 8
+; RV64I-NEXT:    lui a4, 16
+; RV64I-NEXT:    srliw a0, a0, 24
+; RV64I-NEXT:    addi a4, a4, -256
+; RV64I-NEXT:    and a3, a3, a4
+; RV64I-NEXT:    and a4, a2, a4
+; RV64I-NEXT:    or a0, a3, a0
+; RV64I-NEXT:    slli a4, a4, 8
+; RV64I-NEXT:    slli a2, a2, 24
+; RV64I-NEXT:    or a2, a2, a4
+; RV64I-NEXT:    or a0, a2, a0
+; RV64I-NEXT:    sraw a0, a0, a1
+; RV64I-NEXT:    slli a0, a0, 1
+; RV64I-NEXT:    ret
+  %a = or i32 %x, 192
+  %d = call i32 @llvm.bswap.i32(i32 %a)
+  %b = ashr i32 %d, %y
+  %c = shl i32 %b, 1
+  ret i32 %c
+}