[llvm] [RISCV] Add SRAW to ComputeNumSignBitsForTargetNode. (PR #155564)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/155564

>From 0ad34a6b6e8308bb0e3abc397766a10d12d739ea Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 26 Aug 2025 23:02:22 -0700
Subject: [PATCH 1/4] Pre-commit tests

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

diff --git a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
index b3c22a5322cb4..b3b304f1e7e21 100644
--- a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
+++ b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
@@ -220,3 +220,18 @@ define signext i32 @test14(ptr %0, ptr %1, i64 %2) {
   %12 = add i32 %9, %11
   ret i32 %12
 }
+
+; Test that we can propage sign bits through sraw. We should use an slli
+; instead of slliw.
+define signext i32 @test15(i32 signext %x, i32 signext %y) {
+; RV64I-LABEL: test15:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    srli a0, a0, 1
+; RV64I-NEXT:    sraw a0, a0, a1
+; RV64I-NEXT:    slliw a0, a0, 1
+; RV64I-NEXT:    ret
+  %a = ashr i32 %x, 1
+  %b = ashr i32 %a, %y
+  %c = shl i32 %b, 1
+  ret i32 %c
+}

>From 50561aea6185890816522c7800bbaa7df0877487 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 26 Aug 2025 23:41:40 -0700
Subject: [PATCH 2/4] [RISCV] Add SRAW to ComputeNumSignBitsForTargetNode.

---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 8 +++++++-
 llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index b84bd1ce0ac50..c0e3a18c1096a 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -21451,8 +21451,14 @@ unsigned RISCVTargetLowering::ComputeNumSignBitsForTargetNode(
     if (Tmp < 33) return 1;
     return 33;
   }
+  case RISCVISD::SRAW: {
+    unsigned Tmp =
+        DAG.ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth + 1);
+    // sraw produces at least 33 sign bits. If the input already has more than
+    // 33 sign bits sraw, will preserve them.
+    return std::max(Tmp, 33U);
+  }
   case RISCVISD::SLLW:
-  case RISCVISD::SRAW:
   case RISCVISD::SRLW:
   case RISCVISD::DIVW:
   case RISCVISD::DIVUW:
diff --git a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
index b3b304f1e7e21..6617de051a09d 100644
--- a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
+++ b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
@@ -228,7 +228,7 @@ define signext i32 @test15(i32 signext %x, i32 signext %y) {
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    srli a0, a0, 1
 ; RV64I-NEXT:    sraw a0, a0, a1
-; RV64I-NEXT:    slliw a0, a0, 1
+; RV64I-NEXT:    slli a0, a0, 1
 ; RV64I-NEXT:    ret
   %a = ashr i32 %x, 1
   %b = ashr i32 %a, %y

>From 74138a9c862d67812b01c7a586d61ff5c1ef247b Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 26 Aug 2025 23:43:03 -0700
Subject: [PATCH 3/4] Update llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll

---
 llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
index 6617de051a09d..90735d88494b5 100644
--- a/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
+++ b/llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll
@@ -221,7 +221,7 @@ define signext i32 @test14(ptr %0, ptr %1, i64 %2) {
   ret i32 %12
 }
 
-; Test that we can propage sign bits through sraw. We should use an slli
+; Test that we can propagate sign bits through sraw. We should use an slli
 ; instead of slliw.
 define signext i32 @test15(i32 signext %x, i32 signext %y) {
 ; RV64I-LABEL: test15:

>From 6cd7df2838ec64d302929285614c0063cd6aafd8 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Wed, 27 Aug 2025 09:05:34 -0700
Subject: [PATCH 4/4] fixup! Update comments

---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index c0e3a18c1096a..a3d6f77656360 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -21456,6 +21456,8 @@ unsigned RISCVTargetLowering::ComputeNumSignBitsForTargetNode(
         DAG.ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth + 1);
     // sraw produces at least 33 sign bits. If the input already has more than
     // 33 sign bits sraw, will preserve them.
+    // TODO: A more precise answer could be calculated depending on known bits
+    // in the shift amount.
     return std::max(Tmp, 33U);
   }
   case RISCVISD::SLLW:
@@ -21469,9 +21471,7 @@ unsigned RISCVTargetLowering::ComputeNumSignBitsForTargetNode(
   case RISCVISD::FCVT_WU_RV64:
   case RISCVISD::STRICT_FCVT_W_RV64:
   case RISCVISD::STRICT_FCVT_WU_RV64:
-    // TODO: As the result is sign-extended, this is conservatively correct. A
-    // more precise answer could be calculated for SRAW depending on known
-    // bits in the shift amount.
+    // TODO: As the result is sign-extended, this is conservatively correct.
     return 33;
   case RISCVISD::VMV_X_S: {
     // The number of sign bits of the scalar result is computed by obtaining the