[llvm] [SelectionDAG] Lowering usub.sat(a, 1) to a - (a != 0) (PR #170076)

[guan jian via llvm-commits]

https://github.com/rez5427 created https://github.com/llvm/llvm-project/pull/170076

I recently observed that LLVM generates the following code:
```
	addi	a1, a0, -1
	sltu	a0, a0, a1
	addi	a0, a0, -1
	and	a0, a0, a1
	ret
```
This could be optimized using the snez instruction instead.

>From 7f5c66e59e7401efeab7a77d5780d746a14cb56c Mon Sep 17 00:00:00 2001
From: rez5427 <guanjian at stu.cdut.edu.cn>
Date: Mon, 1 Dec 2025 13:10:49 +0800
Subject: [PATCH] [SelectionDAG] Lowering usub.sat(a, 1) to a - (a != 0)

Signed-off-by: rez5427 <guanjian at stu.cdut.edu.cn>
---
 .../CodeGen/SelectionDAG/TargetLowering.cpp   | 10 +++++++
 llvm/test/CodeGen/AArch64/and-mask-removal.ll |  6 ++--
 llvm/test/CodeGen/RISCV/usub_sat.ll           | 28 +++++++++++++++++++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 521d8f07434e6..b3af59b6528b9 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -10867,6 +10867,16 @@ SDValue TargetLowering::expandAddSubSat(SDNode *Node, SelectionDAG &DAG) const {
   assert(VT == RHS.getValueType() && "Expected operands to be the same type");
   assert(VT.isInteger() && "Expected operands to be integers");
 
+  // usub.sat(a, 1) -> a - zext(a != 0)
+  if (Opcode == ISD::USUBSAT && !VT.isVector() && isOneConstant(RHS)) {
+    SDValue Zero = DAG.getConstant(0, dl, VT);
+    EVT BoolVT =
+        getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(), VT);
+    SDValue IsNonZero = DAG.getSetCC(dl, BoolVT, LHS, Zero, ISD::SETNE);
+    SDValue Subtrahend = DAG.getBoolExtOrTrunc(IsNonZero, dl, VT, BoolVT);
+    return DAG.getNode(ISD::SUB, dl, VT, LHS, Subtrahend);
+  }
+
   // usub.sat(a, b) -> umax(a, b) - b
   if (Opcode == ISD::USUBSAT && isOperationLegal(ISD::UMAX, VT)) {
     SDValue Max = DAG.getNode(ISD::UMAX, dl, VT, LHS, RHS);
diff --git a/llvm/test/CodeGen/AArch64/and-mask-removal.ll b/llvm/test/CodeGen/AArch64/and-mask-removal.ll
index 5046c0571ad2b..855fe5caf97b2 100644
--- a/llvm/test/CodeGen/AArch64/and-mask-removal.ll
+++ b/llvm/test/CodeGen/AArch64/and-mask-removal.ll
@@ -483,9 +483,9 @@ define i64 @pr58109(i8 signext %0) {
 ; CHECK-SD-LABEL: pr58109:
 ; CHECK-SD:       ; %bb.0:
 ; CHECK-SD-NEXT:    add w8, w0, #1
-; CHECK-SD-NEXT:    and w8, w8, #0xff
-; CHECK-SD-NEXT:    subs w8, w8, #1
-; CHECK-SD-NEXT:    csel w0, wzr, w8, lo
+; CHECK-SD-NEXT:    ands w8, w8, #0xff
+; CHECK-SD-NEXT:    cset w9, ne
+; CHECK-SD-NEXT:    sub w0, w8, w9
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: pr58109:
diff --git a/llvm/test/CodeGen/RISCV/usub_sat.ll b/llvm/test/CodeGen/RISCV/usub_sat.ll
index 33056682dcc79..7084885a0ee2c 100644
--- a/llvm/test/CodeGen/RISCV/usub_sat.ll
+++ b/llvm/test/CodeGen/RISCV/usub_sat.ll
@@ -185,3 +185,31 @@ define zeroext i4 @func3(i4 zeroext %x, i4 zeroext %y) nounwind {
   %tmp = call i4 @llvm.usub.sat.i4(i4 %x, i4 %y);
   ret i4 %tmp;
 }
+
+define signext i32 @sat_dec_i32(i32 signext %x) nounwind {
+; RV32I-LABEL: sat_dec_i32:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    snez a1, a0
+; RV32I-NEXT:    sub a0, a0, a1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: sat_dec_i32:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    snez a1, a0
+; RV64I-NEXT:    subw a0, a0, a1
+; RV64I-NEXT:    ret
+;
+; RV32IZbb-LABEL: sat_dec_i32:
+; RV32IZbb:       # %bb.0:
+; RV32IZbb-NEXT:    snez a1, a0
+; RV32IZbb-NEXT:    sub a0, a0, a1
+; RV32IZbb-NEXT:    ret
+;
+; RV64IZbb-LABEL: sat_dec_i32:
+; RV64IZbb:       # %bb.0:
+; RV64IZbb-NEXT:    snez a1, a0
+; RV64IZbb-NEXT:    subw a0, a0, a1
+; RV64IZbb-NEXT:    ret
+  %tmp = call i32 @llvm.usub.sat.i32(i32 %x, i32 1)
+  ret i32 %tmp
+}