[llvm] [RISCV] Check the types are the same for folding (sub 0, (setcc x, 0, setlt)) to (sra x, xlen - 1) (PR #158179)

[Jim Lin via llvm-commits]

https://github.com/tclin914 created https://github.com/llvm/llvm-project/pull/158179

We should check the type of x is the same as `sub` operation. Otherwise the shift amount xlen -1 will exceed the bit size of x.

Fixes https://github.com/llvm/llvm-project/issues/158121.

>From 5ddf07f5f475c1bc36775a8dff124893b8608646 Mon Sep 17 00:00:00 2001
From: Jim Lin <jim at andestech.com>
Date: Fri, 12 Sep 2025 11:08:53 +0800
Subject: [PATCH] [RISCV] Check the types are the same for folding (sub 0,
 (setcc x, 0, setlt)) to (sra x, xlen - 1)

We should check the type of x is the same as `sub` operation. Otherwise
the shift amount xlen -1 will exceed the bit size of x.

Fixes https://github.com/llvm/llvm-project/issues/158121.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp |  3 ++-
 llvm/test/CodeGen/RISCV/pr158121.ll         | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/RISCV/pr158121.ll

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index ae9e2fef88673..0203d1a701147 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -15830,7 +15830,8 @@ static SDValue performSUBCombine(SDNode *N, SelectionDAG &DAG,
   SDValue N1 = N->getOperand(1);
   // fold (sub 0, (setcc x, 0, setlt)) -> (sra x, xlen - 1)
   if (isNullConstant(N0) && N1.getOpcode() == ISD::SETCC && N1.hasOneUse() &&
-      isNullConstant(N1.getOperand(1))) {
+      isNullConstant(N1.getOperand(1)) &&
+      N1.getValueType() == N1.getOperand(0).getValueType()) {
     ISD::CondCode CCVal = cast<CondCodeSDNode>(N1.getOperand(2))->get();
     if (CCVal == ISD::SETLT) {
       SDLoc DL(N);
diff --git a/llvm/test/CodeGen/RISCV/pr158121.ll b/llvm/test/CodeGen/RISCV/pr158121.ll
new file mode 100644
index 0000000000000..2c018444e9c67
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/pr158121.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=riscv64 | FileCheck %s
+
+define i64 @f(ptr %p) {
+; CHECK-LABEL: f:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lb a0, 0(a0)
+; CHECK-NEXT:    srai a0, a0, 63
+; CHECK-NEXT:    ret
+  %load = load i8, ptr %p, align 1
+  %conv1 = zext i8 %load to i32
+  %cmp = icmp ult i32 127, %conv1
+  %conv2 = zext i1 %cmp to i32
+  %sub = sub nsw i32 0, %conv2
+  %conv3 = sext i32 %sub to i64
+  ret i64 %conv3
+}