[llvm] [InstCombine] Narrow the ctpop when upper bits are known zero (PR #196537)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Rutansh Suthar (RutanshS)

<details>
<summary>Changes</summary>

Fixes #<!-- -->196434

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+12) 
- (modified) llvm/test/Transforms/InstCombine/ctpop.ll (+39) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index da258b3b34e27..9968d2d9371d7 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -696,6 +696,18 @@ static Instruction *foldCtpop(IntrinsicInst &II, InstCombinerImpl &IC) {
   KnownBits Known(BitWidth);
   IC.computeKnownBits(Op0, Known, &II);
 
+  // If the upper half of the bits are known zero, narrow:
+  // ctpop(x) --> zext(ctpop(trunc(x)))
+  if (BitWidth >= 64) {
+    unsigned HalfWidth = BitWidth / 2;
+    if (Known.countMinLeadingZeros() >= HalfWidth) {
+      Type *NarrowTy = IntegerType::get(II.getContext(), HalfWidth);
+      Value *Trunc = IC.Builder.CreateTrunc(Op0, NarrowTy);
+      Value *NarrowPop =
+          IC.Builder.CreateUnaryIntrinsic(Intrinsic::ctpop, Trunc);
+      return CastInst::Create(Instruction::ZExt, NarrowPop, Ty);
+    }
+  }
   // If all bits are zero except for exactly one fixed bit, then the result
   // must be 0 or 1, and we can get that answer by shifting to LSB:
   // ctpop (X & 32) --> (X & 32) >> 5
diff --git a/llvm/test/Transforms/InstCombine/ctpop.ll b/llvm/test/Transforms/InstCombine/ctpop.ll
index 217a3c308a2dd..f96cd1f96eea7 100644
--- a/llvm/test/Transforms/InstCombine/ctpop.ll
+++ b/llvm/test/Transforms/InstCombine/ctpop.ll
@@ -503,3 +503,42 @@ define i32 @ctpop_non_zero_with_existing_range_attr(i32 range(i32 1, 255) %x) {
   %ctpop = call range(i32 0, 9) i32 @llvm.ctpop.i32(i32 %x)
   ret i32 %ctpop
 }
+
+define i32 @ctpop_narrow_i64_to_i32_via_assume(i64 %x) {
+; CHECK-LABEL: @ctpop_narrow_i64_to_i32_via_assume(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[X:%.*]], 4294967296
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc nuw i64 [[X]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = call range(i32 0, 33) i32 @llvm.ctpop.i32(i32 [[TMP1]])
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %cmp = icmp ult i64 %x, 4294967296
+  call void @llvm.assume(i1 %cmp)
+  %pop = call i64 @llvm.ctpop.i64(i64 %x)
+  %trunc = trunc i64 %pop to i32
+  ret i32 %trunc
+}
+
+define i64 @ctpop_no_narrow_no_assume(i64 %x) {
+; CHECK-LABEL: @ctpop_no_narrow_no_assume(
+; CHECK-NEXT:    [[POP:%.*]] = call range(i64 0, 65) i64 @llvm.ctpop.i64(i64 [[X:%.*]])
+; CHECK-NEXT:    ret i64 [[POP]]
+;
+  %pop = call i64 @llvm.ctpop.i64(i64 %x)
+  ret i64 %pop
+}
+
+define i32 @ctpop_no_narrow_not_enough_zeros(i64 %x) {
+; CHECK-LABEL: @ctpop_no_narrow_not_enough_zeros(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[X:%.*]], 8589934592
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[POP:%.*]] = call range(i64 0, 34) i64 @llvm.ctpop.i64(i64 [[X]])
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw nsw i64 [[POP]] to i32
+; CHECK-NEXT:    ret i32 [[TRUNC]]
+;
+  %cmp = icmp ult i64 %x, 8589934592
+  call void @llvm.assume(i1 %cmp)
+  %pop = call i64 @llvm.ctpop.i64(i64 %x)
+  %trunc = trunc i64 %pop to i32
+  ret i32 %trunc
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/196537