[llvm] [SelectionDAG] Fix incorrect fold condition in foldSetCCWithFunnelShift. (PR #137637)

[via llvm-commits]

https://github.com/Ruhung created https://github.com/llvm/llvm-project/pull/137637

Proposed by [2ed1598](https://github.com/llvm/llvm-project/commit/2ed15984b49a1af87be37ec8bd6ee3ab7f724767):

`fshl (or X, Y), X, C ==/!= 0  -->  (or (shl Y, C), X) ==/!= 0`

This transformation is valid when C != 0, as verified by [Alive2](https://alive2.llvm.org/ce/z/JupMGN).

Fixes #136746 

>From a7f23b7a04f36027dc1504c93b29df3ec45665bb Mon Sep 17 00:00:00 2001
From: Ruhung <jhlee at pllab.cs.nthu.edu.tw>
Date: Mon, 28 Apr 2025 21:45:36 +0800
Subject: [PATCH 1/2] Pre-commit tests

---
 llvm/test/CodeGen/AArch64/setcc-fsh.ll | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/llvm/test/CodeGen/AArch64/setcc-fsh.ll b/llvm/test/CodeGen/AArch64/setcc-fsh.ll
index 08bfe282703ff..3e592f491c99f 100644
--- a/llvm/test/CodeGen/AArch64/setcc-fsh.ll
+++ b/llvm/test/CodeGen/AArch64/setcc-fsh.ll
@@ -248,3 +248,14 @@ define i1 @fshl_or_ne_2(i32 %x, i32 %y) {
   %r = icmp ne i32 %f, 2
   ret i1 %r
 }
+
+define i1 @fshr_0_or_eq_0(i16 %x, i16 %y) {
+; CHECK-LABEL: fshr_0_or_eq_0:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    ret
+  %or = or i16 %x, %y
+  %f = call i16 @llvm.fshr.i16(i16 %or, i16 %x, i16 0)
+  %r = icmp eq i16 %f, 0
+  ret i1 %r
+}

>From 4dbe104e90215cc1bbc5eae8b6e9d1c4152ea289 Mon Sep 17 00:00:00 2001
From: Ruhung <jhlee at pllab.cs.nthu.edu.tw>
Date: Mon, 28 Apr 2025 21:45:59 +0800
Subject: [PATCH 2/2] [SelectionDAG] Fix incorrect fold condition in
 foldSetCCWithFunnelShift.

---
 llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp | 3 ++-
 llvm/test/CodeGen/AArch64/setcc-fsh.ll           | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 6930b54ddb14a..1e9fb1aa2ea61 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -4462,7 +4462,8 @@ static SDValue foldSetCCWithFunnelShift(EVT VT, SDValue N0, SDValue N1,
 
   unsigned BitWidth = N0.getScalarValueSizeInBits();
   auto *ShAmtC = isConstOrConstSplat(N0.getOperand(2));
-  if (!ShAmtC || ShAmtC->getAPIntValue().uge(BitWidth))
+  APInt AmtVal = ShAmtC->getAPIntValue();
+  if (!ShAmtC || AmtVal.uge(BitWidth) || AmtVal.isZero())
     return SDValue();
 
   // Canonicalize fshr as fshl to reduce pattern-matching.
diff --git a/llvm/test/CodeGen/AArch64/setcc-fsh.ll b/llvm/test/CodeGen/AArch64/setcc-fsh.ll
index 3e592f491c99f..f0cf775f5c2fa 100644
--- a/llvm/test/CodeGen/AArch64/setcc-fsh.ll
+++ b/llvm/test/CodeGen/AArch64/setcc-fsh.ll
@@ -252,7 +252,8 @@ define i1 @fshl_or_ne_2(i32 %x, i32 %y) {
 define i1 @fshr_0_or_eq_0(i16 %x, i16 %y) {
 ; CHECK-LABEL: fshr_0_or_eq_0:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    tst w0, #0xffff
+; CHECK-NEXT:    cset w0, eq
 ; CHECK-NEXT:    ret
   %or = or i16 %x, %y
   %f = call i16 @llvm.fshr.i16(i16 %or, i16 %x, i16 0)