[llvm] Missing opt with ctlz and shifts of power of 2 constants (#41333) (PR #74175)

[Sizov Nikita via llvm-commits]

https://github.com/snikitav updated https://github.com/llvm/llvm-project/pull/74175

>From a0bf3e94880484671ccd7cc537c2264e9b958931 Mon Sep 17 00:00:00 2001
From: Sizov Nikita <s.nikita.v at gmail.com>
Date: Sat, 2 Dec 2023 04:53:32 +0300
Subject: [PATCH] Missing opt with ctlz and shifts of power of 2 constants
 (#41333)

---
 .../InstCombine/InstCombineCalls.cpp          | 32 +++++++++++++++++
 .../InstCombine/ctlz-cttz-shifts.ll           | 36 +++++++++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/ctlz-cttz-shifts.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index a991f0906052a30..49b55a7a0cc871e 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -514,6 +514,8 @@ static Instruction *foldCttzCtlz(IntrinsicInst &II, InstCombinerImpl &IC) {
     return IC.replaceInstUsesWith(II, ConstantInt::getNullValue(II.getType()));
   }
 
+  ConstantInt *CInt;
+
   if (IsTZ) {
     // cttz(-x) -> cttz(x)
     if (match(Op0, m_Neg(m_Value(X))))
@@ -549,6 +551,36 @@ static Instruction *foldCttzCtlz(IntrinsicInst &II, InstCombinerImpl &IC) {
 
     if (match(Op0, m_Intrinsic<Intrinsic::abs>(m_Value(X))))
       return IC.replaceOperand(II, 0, X);
+
+    // cttz(shl(%const, %val), 1) --> add(cttz(%const, 1), %val)
+    if (match(Op0, m_Shl(m_ConstantInt(CInt), m_Value(X))) &&
+        match(Op1, m_One())) {
+      Value *ConstCttz =
+          IC.Builder.CreateBinaryIntrinsic(Intrinsic::cttz, CInt, Op1);
+      return IC.replaceInstUsesWith(II, IC.Builder.CreateAdd(ConstCttz, X));
+    }
+
+    // cttz(lshr exact (%const, %val), 0) --> sub(cttz(%const, 0), %val)
+    if (match(Op0, m_Exact(m_LShr(m_ConstantInt(CInt), m_Value(X))))) {
+      Value *ConstCttz =
+          IC.Builder.CreateBinaryIntrinsic(Intrinsic::cttz, CInt, Op1);
+      return IC.replaceInstUsesWith(II, IC.Builder.CreateSub(ConstCttz, X));
+    }
+  } else {
+    // ctlz(lshr(%const, %val), 1) --> add(ctlz(%const, 1), %val)
+    if (match(Op0, m_LShr(m_ConstantInt(CInt), m_Value(X))) &&
+        match(Op1, m_One())) {
+      Value *ConstCtlz =
+          IC.Builder.CreateBinaryIntrinsic(Intrinsic::ctlz, CInt, Op1);
+      return IC.replaceInstUsesWith(II, IC.Builder.CreateAdd(ConstCtlz, X));
+    }
+
+    // ctlz(shl nuw (%const, %val), 0) --> sub(ctlz(%const, 0), %val)
+    if (match(Op0, m_NUWShl(m_ConstantInt(CInt), m_Value(X)))) {
+      Value *ConstCtlz =
+          IC.Builder.CreateBinaryIntrinsic(Intrinsic::ctlz, CInt, Op1);
+      return IC.replaceInstUsesWith(II, IC.Builder.CreateSub(ConstCtlz, X));
+    }
   }
 
   KnownBits Known = IC.computeKnownBits(Op0, 0, &II);
diff --git a/llvm/test/Transforms/InstCombine/ctlz-cttz-shifts.ll b/llvm/test/Transforms/InstCombine/ctlz-cttz-shifts.ll
new file mode 100644
index 000000000000000..32e5ee10835d90c
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/ctlz-cttz-shifts.ll
@@ -0,0 +1,36 @@
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+declare i32 @llvm.ctlz.i32(i32, i1)
+declare i32 @llvm.cttz.i32(i32, i1)
+
+define i32 @lshr_ctlz_true(i32) {
+; CHECK-LABEL: @lshr_ctlz_true
+; CHECK-NEXT:  %ctlz = add
+  %lshr = lshr i32 8387584, %0
+  %ctlz = call i32 @llvm.ctlz.i32(i32 %lshr, i1 true)
+  ret i32 %ctlz
+}
+
+define i32 @shl_ctlz_false(i32) {
+; CHECK-LABEL: @shl_ctlz_false
+; CHECK-NEXT:  %ctlz = sub
+  %shl = shl nuw i32 8387584, %0
+  %ctlz = call i32 @llvm.ctlz.i32(i32 %shl, i1 false)
+  ret i32 %ctlz
+}
+
+define i32 @lshr_exact_cttz_false(i32) {
+; CHECK-LABEL: @lshr_exact_cttz_false
+; CHECK-NEXT:  %cttz = sub
+  %lshr = lshr exact i32 8387584, %0
+  %cttz = call i32 @llvm.cttz.i32(i32 %lshr, i1 false)
+  ret i32 %cttz
+}
+
+define i32 @shl_cttz_true(i32) {
+; CHECK-LABEL: @shl_cttz_true
+; CHECK-NEXT:  %cttz = add
+  %shl = shl i32 8387584, %0
+  %cttz = call i32 @llvm.cttz.i32(i32 %shl, i1 true)
+  ret i32 %cttz
+}