From f3f717bbfa61600703b3b2149ebdb87fbee2dbac Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Sat, 30 Aug 2025 18:27:18 -0700
Subject: [PATCH] [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).
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp |  9 ++++++++
 llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll | 25 +++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 9c268012b24d..9115c1385d6d 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 90735d88494b..b23f4f581200 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
+}