[InstCombine] Added optimisation for trunc (Pow2 >> x) to i1

[kper]

Closes #156898

I have added two cases. The first one matches when the constant is exactly power of 2. The second case was to address the general case mentioned in the linked issue. I, however, did not really solve the general case.
We can only emit a icmp ult if all the bits are one and that's only the case when the constant + 1 is a power of 2. Otherwise, we need to create icmp eq for every bit that is one.

Here are a few examples which won't be working with the two cases:

• constant is 9: https://alive2.llvm.org/ce/z/S5FLJZ
• subrange in 56: https://alive2.llvm.org/ce/z/yn_ZNG
• and finally an example as worst case (because it alternates the bits): https://alive2.llvm.org/ce/z/nDitNA

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (kper)

Changes

Closes #156898

I have added two cases. The first one matches when the constant is exactly power of 2. The second case was to address the general case mentioned in the linked issue. I, however, did not really solve the general case.
We can only emit a icmp ult if all the bits are one and that's only the case when the constant + 1 is a power of 2. Otherwise, we need to create icmp eq for every bit that is one.

Here are a few examples which won't be working with the two cases:

• constant is 9: https://alive2.llvm.org/ce/z/S5FLJZ
• subrange in 56: https://alive2.llvm.org/ce/z/yn_ZNG
• and finally an example as worst case (because it alternates the bits): https://alive2.llvm.org/ce/z/nDitNA

I wonder whether I should still implement the general case since it increments the number of instructions?

cc @nikic @andjo403

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp (+23)
• (modified) llvm/test/Transforms/InstCombine/trunc-lshr.ll (+21)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index fdef49e310f81..a3e9969503f02 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -11,11 +11,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/Value.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
 #include <optional>
@@ -969,6 +971,27 @@ Instruction *InstCombinerImpl::visitTrunc(TruncInst &Trunc) {
     Changed = true;
   }
 
+  const APInt *C1;
+  Value *V1;
+  // trunc (lshr i8 C1, V1) to i1 -> icmp eq V1, sqrt(C1) iff C1 is power of 2
+  if (DestWidth == 1 &&
+      match(Src, m_OneUse(m_Shr(m_Power2(C1), m_Value(V1))))) {
+    const APInt Sqrt = C1->sqrt();
+    Value *Right = ConstantInt::get(V1->getType(), Sqrt);
+    Value *Icmp = Builder.CreateICmpEQ(V1, Right);
+    return replaceInstUsesWith(Trunc, Icmp);
+  }
+
+  // trunc (lshr i8 C1, V1) to i1 -> icmp ult V1, sqrt(C1 + 1) iff (C1 + 1) is
+  // power of 2
+  if (DestWidth == 1 && match(Src, m_OneUse(m_Shr(m_APInt(C1), m_Value(V1)))) &&
+      (*C1 + 1).isPowerOf2()) {
+    const APInt Sqrt = (*C1 + 1).sqrt();
+    Value *Right = ConstantInt::get(V1->getType(), Sqrt);
+    Value *Icmp = Builder.CreateICmpULT(V1, Right);
+    return replaceInstUsesWith(Trunc, Icmp);
+  }
+
   return Changed ? &Trunc : nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/trunc-lshr.ll b/llvm/test/Transforms/InstCombine/trunc-lshr.ll
index 4364b09cfa709..84daba3d13b9a 100644
--- a/llvm/test/Transforms/InstCombine/trunc-lshr.ll
+++ b/llvm/test/Transforms/InstCombine/trunc-lshr.ll
@@ -93,3 +93,24 @@ define i1 @test5(i32 %i, ptr %p) {
   ret i1 %op
 }
 
+define i1 @test6(i8 %x) {
+; CHECK-LABEL: define i1 @test6(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TRUNC:%.*]] = icmp eq i8 [[X]], 2
+; CHECK-NEXT:    ret i1 [[TRUNC]]
+;
+  %lshr = lshr i8 4, %x
+  %trunc = trunc i8 %lshr to i1
+  ret i1 %trunc
+}
+
+define i1 @test7(i8 %x) {
+; CHECK-LABEL: define i1 @test7(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TRUNC:%.*]] = icmp ult i8 [[X]], 4
+; CHECK-NEXT:    ret i1 [[TRUNC]]
+;
+  %lshr = lshr i8 15, %x
+  %trunc = trunc i8 %lshr to i1
+  ret i1 %trunc
+}

[andjo403]

the folds also seems to hold for ashr see https://alive2.llvm.org/ce/z/Dm5HEp

[github-actions]

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

[andjo403]

Looks good to me but wait for an other review

[dtcxzyw]

@zyw-bot mfuzz

[dtcxzyw]

LGTM. Thanks.

As a follow-up, are you interested in folding trunc (shr C=0b11111...0000, %x) to i1 -> icmp ugt %x, cttz(C) - 1 as well? Not sure if it is profitable in real-world programs.

[dtcxzyw]

Suggested change

	Value *Right = ConstantInt::get(V1->getType(), (*C1 + 1).countr_zero());
	Value *Right = ConstantInt::get(V1->getType(), C1->countr_one());

[kper]

addressed

[dtcxzyw]

@kper Please avoid force-pushing your branch if unnecessary.

[kper]

@kper Please avoid force-pushing your branch if unnecessary.

ah ok, sry
I just wanted to have one commit

[kper]

LGTM. Thanks.

As a follow-up, are you interested in folding trunc (shr C=0b11111...0000, %x) to i1 -> icmp ugt %x, cttz(C) - 1 as well? Not sure if it is profitable in real-world programs.

Yeah sure :)
I will make a new PR tomorrow and let's see whether it improves something
Can't hurt to try it, thanks for the suggestion