Missing opt with ctlz and shifts of power of 2 constants (#41333)

[snikitav]

This patch does the following simplifications:

cttz(shl(C, X), 1) -> add(cttz(C, 1), X)
cttz(lshr exact(C, X), 1) -> sub(cttz(C, 1), X)
ctlz(lshr(C, X), 1) --> add(ctlz(C, 1), X)
ctlz(shl nuw (C, X), 1) --> sub(ctlz(C, 1), X)

Alive2: https://alive2.llvm.org/ce/z/9KHlKc
Closes #41333

[llvmbot]

@llvm/pr-subscribers-clang

@llvm/pr-subscribers-llvm-transforms

Author: Sizov Nikita (snikitav)

Changes

Closes #41333

---

Full diff: https://github.com/llvm/llvm-project/pull/74175.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+32)
• (added) llvm/test/Transforms/InstCombine/ctlz-cttz-shifts.ll (+36)

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
+}

[dtcxzyw]

Could you please add the alive2 proof to the PR description?

[dtcxzyw]

LGTM. Please wait for additional approval from other reviewers.

[nikic]

LGTM apart from one issue.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[dtcxzyw]

Thank you!