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

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/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).

>From 3cf55b8e72ded7ee82f2cae73f8dd4459290d821 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Sat, 30 Aug 2025 09:18:36 -0700
Subject: [PATCH 1/2] Pre-commit test

---
 llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll | 25 +++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
index 90735d88494b5..a75007de83310 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:    slliw 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
+}

>From 0e6fd7bffed29f42f9b0913a0b45fd6fd8d2f47c Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Sat, 30 Aug 2025 09:38:10 -0700
Subject: [PATCH 2/2] [RISCV] Add computeKnownBitsForTargetNode for
 RISCVISD::SRAW.

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.

Test case is a little contrived because I had to find something
that didn't get optimized before type legalizaton and didn't get
type legalized to a sign extending operation.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 9 +++++++++
 llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 9c268012b24d2..9115c1385d6df 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 a75007de83310..b23f4f5812000 100644
--- a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
+++ b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
@@ -252,7 +252,7 @@ define signext i32 @test16(i32 signext %x, i32 signext %y) {
 ; RV64I-NEXT:    or a2, a2, a4
 ; RV64I-NEXT:    or a0, a2, a0
 ; RV64I-NEXT:    sraw a0, a0, a1
-; RV64I-NEXT:    slliw a0, a0, 1
+; RV64I-NEXT:    slli a0, a0, 1
 ; RV64I-NEXT:    ret
   %a = or i32 %x, 192
   %d = call i32 @llvm.bswap.i32(i32 %a)