[llvm] [DAGCombine] Fix type mismatch in `(shl X, cttz(Y)) -> (mul (Y & -Y), X)` (PR #94008)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/94008

>From 703bea7b4b7626c7e902dc483b069ca35b867266 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 1 Jun 2024 02:35:37 +0800
Subject: [PATCH 1/2] [DAGCombine] Fix type mismatch in `(shl X, cttz(Y)) ->
 (mul (Y & -Y), X)`

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 11 +++--
 llvm/test/CodeGen/PowerPC/pr85066.ll          | 45 +++++++++++++++++++
 2 files changed, 52 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/CodeGen/PowerPC/pr85066.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 42e861e61201c..2084f9727f9bb 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -10109,13 +10109,16 @@ SDValue DAGCombiner::visitSHL(SDNode *N) {
 
   // fold (shl X, cttz(Y)) -> (mul (Y & -Y), X) if cttz is unsupported on the
   // target.
-  if ((N1.getOpcode() == ISD::CTTZ || N1.getOpcode() == ISD::CTTZ_ZERO_UNDEF) &&
-      N1.hasOneUse() && !TLI.isOperationLegalOrCustom(ISD::CTTZ, VT) &&
+  if (((N1.getOpcode() == ISD::CTTZ &&
+        VT.getScalarSizeInBits() >= ShiftVT.getScalarSizeInBits()) ||
+       N1.getOpcode() == ISD::CTTZ_ZERO_UNDEF) &&
+      N1.hasOneUse() && !TLI.isOperationLegalOrCustom(ISD::CTTZ, ShiftVT) &&
       TLI.isOperationLegalOrCustom(ISD::MUL, VT)) {
     SDValue Y = N1.getOperand(0);
     SDLoc DL(N);
-    SDValue NegY = DAG.getNegative(Y, DL, VT);
-    SDValue And = DAG.getNode(ISD::AND, DL, VT, Y, NegY);
+    SDValue NegY = DAG.getNegative(Y, DL, ShiftVT);
+    SDValue And =
+        DAG.getZExtOrTrunc(DAG.getNode(ISD::AND, DL, ShiftVT, Y, NegY), DL, VT);
     return DAG.getNode(ISD::MUL, DL, VT, And, N0);
   }
 
diff --git a/llvm/test/CodeGen/PowerPC/pr85066.ll b/llvm/test/CodeGen/PowerPC/pr85066.ll
new file mode 100644
index 0000000000000..9b47d6e6508fa
--- /dev/null
+++ b/llvm/test/CodeGen/PowerPC/pr85066.ll
@@ -0,0 +1,45 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=powerpc64le -verify-machineinstrs < %s | FileCheck %s
+
+; Tests from pr85066
+define i64 @test_shl_zext_cttz(i16 %x) {
+; CHECK-LABEL: test_shl_zext_cttz:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    oris 3, 3, 1
+; CHECK-NEXT:    neg 4, 3
+; CHECK-NEXT:    and 3, 3, 4
+; CHECK-NEXT:    clrldi 3, 3, 32
+; CHECK-NEXT:    blr
+entry:
+  %cttz = tail call i16 @llvm.cttz.i16(i16 %x, i1 false)
+  %zext = zext i16 %cttz to i64
+  %res = shl i64 1, %zext
+  ret i64 %res
+}
+
+define i64 @test_shl_zext_cttz_zero_is_poison(i16 %x) {
+; CHECK-LABEL: test_shl_zext_cttz_zero_is_poison:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    neg 4, 3
+; CHECK-NEXT:    and 3, 3, 4
+; CHECK-NEXT:    clrldi 3, 3, 32
+; CHECK-NEXT:    blr
+entry:
+  %cttz = tail call i16 @llvm.cttz.i16(i16 %x, i1 true)
+  %zext = zext i16 %cttz to i64
+  %res = shl i64 1, %zext
+  ret i64 %res
+}
+
+define i16 @test_shl_trunc_cttz(i32 %x) {
+; CHECK-LABEL: test_shl_trunc_cttz:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    neg 4, 3
+; CHECK-NEXT:    and 3, 3, 4
+; CHECK-NEXT:    blr
+entry:
+  %cttz = tail call i32 @llvm.cttz.i32(i32 %x, i1 false)
+  %trunc = trunc i32 %cttz to i16
+  %res = shl i16 1, %trunc
+  ret i16 %res
+}

>From b6faa3368e9ace52b85a65e3fcac65dea0df1a80 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 1 Jun 2024 16:42:34 +0800
Subject: [PATCH 2/2] [SDAG] Address review comments. NFC.

---
 llvm/test/CodeGen/PowerPC/pr85066.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/PowerPC/pr85066.ll b/llvm/test/CodeGen/PowerPC/pr85066.ll
index 9b47d6e6508fa..335038ec4734c 100644
--- a/llvm/test/CodeGen/PowerPC/pr85066.ll
+++ b/llvm/test/CodeGen/PowerPC/pr85066.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
-; RUN: llc -mtriple=powerpc64le -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=powerpc64le < %s | FileCheck %s
 
 ; Tests from pr85066
 define i64 @test_shl_zext_cttz(i16 %x) {