[CodeGenPrepare] Unfold slow ctpop when used in power-of-two test

s-barannikov · s-barannikov · commit 8dc0faba4527 · 2025-04-20T05:13:07.000+03:00
DAG combiner already does this transformation, but in some cases it does
not have a chance because either CodeGenPrepare or SelectionDAGBuilder
move icmp to a different basic block.
diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -474,6 +474,7 @@ class CodeGenPrepare {
   bool optimizeURem(Instruction *Rem);
   bool combineToUSubWithOverflow(CmpInst *Cmp, ModifyDT &ModifiedDT);
   bool combineToUAddWithOverflow(CmpInst *Cmp, ModifyDT &ModifiedDT);
+  bool unfoldPow2Test(CmpInst *Cmp);
   void verifyBFIUpdates(Function &F);
   bool _run(Function &F);
 };
@@ -1762,6 +1763,40 @@ bool CodeGenPrepare::combineToUSubWithOverflow(CmpInst *Cmp,
   return true;
 }
 
+// Decanonicalizes icmp+ctpop power-of-two test if ctpop is slow.
+bool CodeGenPrepare::unfoldPow2Test(CmpInst *Cmp) {
+  CmpPredicate Pred;
+  Value *X;
+  uint64_t C;
+
+  if (!match(Cmp, m_ICmp(Pred, m_Intrinsic<Intrinsic::ctpop>(m_Value(X)),
+                         m_ConstantInt(C))))
+    return false;
+
+  Type *Ty = X->getType();
+  if (Ty->isVectorTy() || TTI->getPopcntSupport(Ty->getIntegerBitWidth()) ==
+                              TargetTransformInfo::PSK_FastHardware)
+    return false;
+
+  // (ctpop x) u< 2 -> (x & (x - 1)) == 0
+  // (ctpop x) u> 1 -> (x & (x - 1)) != 0
+  if ((Pred == CmpInst::ICMP_ULT && C == 2) ||
+      (Pred == CmpInst::ICMP_UGT && C == 1)) {
+    IRBuilder<> Builder(Cmp);
+    Value *Sub = Builder.CreateAdd(X, Constant::getAllOnesValue(Ty));
+    Value *And = Builder.CreateAnd(X, Sub);
+    CmpInst::Predicate NewPred =
+        Pred == CmpInst::ICMP_ULT ? CmpInst::ICMP_EQ : CmpInst::ICMP_NE;
+    Value *NewCmp =
+        Builder.CreateICmp(NewPred, And, ConstantInt::getNullValue(Ty));
+    Cmp->replaceAllUsesWith(NewCmp);
+    RecursivelyDeleteTriviallyDeadInstructions(Cmp);
+    return true;
+  }
+
+  return false;
+}
+
 /// Sink the given CmpInst into user blocks to reduce the number of virtual
 /// registers that must be created and coalesced. This is a clear win except on
 /// targets with multiple condition code registers (PowerPC), where it might
@@ -2183,6 +2218,9 @@ bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (combineToUSubWithOverflow(Cmp, ModifiedDT))
     return true;
 
+  if (unfoldPow2Test(Cmp))
+    return true;
+
   if (foldICmpWithDominatingICmp(Cmp, *TLI))
     return true;
 
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll b/llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll
@@ -357,49 +357,14 @@ define i64 @ctpop_i64(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_two:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    j .LBB6_2
-; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    sltiu a0, zero, 0
-; RV32I-NEXT:    ret
-; RV32I-NEXT:  .LBB6_2:
-; RV32I-NEXT:    srli a2, a0, 1
-; RV32I-NEXT:    lui a3, 349525
-; RV32I-NEXT:    lui a4, 209715
-; RV32I-NEXT:    srli a5, a1, 1
-; RV32I-NEXT:    addi a3, a3, 1365
-; RV32I-NEXT:    and a2, a2, a3
-; RV32I-NEXT:    and a3, a5, a3
-; RV32I-NEXT:    lui a5, 61681
-; RV32I-NEXT:    addi a4, a4, 819
-; RV32I-NEXT:    addi a5, a5, -241
-; RV32I-NEXT:    sub a0, a0, a2
-; RV32I-NEXT:    sub a1, a1, a3
-; RV32I-NEXT:    srli a2, a0, 2
-; RV32I-NEXT:    and a0, a0, a4
-; RV32I-NEXT:    srli a3, a1, 2
-; RV32I-NEXT:    and a1, a1, a4
-; RV32I-NEXT:    and a2, a2, a4
-; RV32I-NEXT:    and a3, a3, a4
-; RV32I-NEXT:    add a0, a2, a0
-; RV32I-NEXT:    add a1, a3, a1
-; RV32I-NEXT:    srli a2, a0, 4
-; RV32I-NEXT:    srli a3, a1, 4
-; RV32I-NEXT:    add a0, a2, a0
-; RV32I-NEXT:    add a1, a3, a1
-; RV32I-NEXT:    and a0, a0, a5
-; RV32I-NEXT:    and a1, a1, a5
-; RV32I-NEXT:    slli a2, a0, 8
-; RV32I-NEXT:    slli a3, a1, 8
-; RV32I-NEXT:    add a0, a0, a2
-; RV32I-NEXT:    add a1, a1, a3
-; RV32I-NEXT:    slli a2, a0, 16
-; RV32I-NEXT:    slli a3, a1, 16
-; RV32I-NEXT:    add a0, a0, a2
-; RV32I-NEXT:    add a1, a1, a3
-; RV32I-NEXT:    srli a0, a0, 24
-; RV32I-NEXT:    srli a1, a1, 24
-; RV32I-NEXT:    add a0, a1, a0
-; RV32I-NEXT:    sltiu a0, a0, 2
+; RV32I-NEXT:    addi a2, a0, -1
+; RV32I-NEXT:    addi a3, a1, -1
+; RV32I-NEXT:    sltiu a4, a2, -1
+; RV32I-NEXT:    add a3, a3, a4
+; RV32I-NEXT:    and a0, a0, a2
+; RV32I-NEXT:    and a1, a1, a3
+; RV32I-NEXT:    or a0, a0, a1
+; RV32I-NEXT:    seqz a0, a0
 ; RV32I-NEXT:    ret
 ;
 ; RV32ZBB-LABEL: ctpop_i64_ugt_two:
@@ -422,50 +387,14 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_one:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    j .LBB7_2
-; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    snez a0, zero
-; RV32I-NEXT:    ret
-; RV32I-NEXT:  .LBB7_2:
-; RV32I-NEXT:    srli a2, a0, 1
-; RV32I-NEXT:    lui a3, 349525
-; RV32I-NEXT:    lui a4, 209715
-; RV32I-NEXT:    srli a5, a1, 1
-; RV32I-NEXT:    addi a3, a3, 1365
-; RV32I-NEXT:    and a2, a2, a3
-; RV32I-NEXT:    and a3, a5, a3
-; RV32I-NEXT:    lui a5, 61681
-; RV32I-NEXT:    addi a4, a4, 819
-; RV32I-NEXT:    addi a5, a5, -241
-; RV32I-NEXT:    sub a0, a0, a2
-; RV32I-NEXT:    sub a1, a1, a3
-; RV32I-NEXT:    srli a2, a0, 2
-; RV32I-NEXT:    and a0, a0, a4
-; RV32I-NEXT:    srli a3, a1, 2
-; RV32I-NEXT:    and a1, a1, a4
-; RV32I-NEXT:    and a2, a2, a4
-; RV32I-NEXT:    and a3, a3, a4
-; RV32I-NEXT:    add a0, a2, a0
-; RV32I-NEXT:    add a1, a3, a1
-; RV32I-NEXT:    srli a2, a0, 4
-; RV32I-NEXT:    srli a3, a1, 4
-; RV32I-NEXT:    add a0, a2, a0
-; RV32I-NEXT:    add a1, a3, a1
-; RV32I-NEXT:    and a0, a0, a5
-; RV32I-NEXT:    and a1, a1, a5
-; RV32I-NEXT:    slli a2, a0, 8
-; RV32I-NEXT:    slli a3, a1, 8
-; RV32I-NEXT:    add a0, a0, a2
-; RV32I-NEXT:    add a1, a1, a3
-; RV32I-NEXT:    slli a2, a0, 16
-; RV32I-NEXT:    slli a3, a1, 16
-; RV32I-NEXT:    add a0, a0, a2
-; RV32I-NEXT:    add a1, a1, a3
-; RV32I-NEXT:    srli a0, a0, 24
-; RV32I-NEXT:    srli a1, a1, 24
-; RV32I-NEXT:    add a0, a1, a0
-; RV32I-NEXT:    sltiu a0, a0, 2
-; RV32I-NEXT:    xori a0, a0, 1
+; RV32I-NEXT:    addi a2, a0, -1
+; RV32I-NEXT:    addi a3, a1, -1
+; RV32I-NEXT:    sltiu a4, a2, -1
+; RV32I-NEXT:    add a3, a3, a4
+; RV32I-NEXT:    and a0, a0, a2
+; RV32I-NEXT:    and a1, a1, a3
+; RV32I-NEXT:    or a0, a0, a1
+; RV32I-NEXT:    snez a0, a0
 ; RV32I-NEXT:    ret
 ;
 ; RV32ZBB-LABEL: ctpop_i64_ugt_one:
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/rv64zbb.ll b/llvm/test/CodeGen/RISCV/GlobalISel/rv64zbb.ll
@@ -701,31 +701,9 @@ define signext i32 @ctpop_i32(i32 signext %a) nounwind {
 define i1 @ctpop_i32_ult_two(i32 signext %a) nounwind {
 ; RV64I-LABEL: ctpop_i32_ult_two:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi sp, sp, -16
-; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
-; RV64I-NEXT:    srliw a1, a0, 1
-; RV64I-NEXT:    lui a2, 349525
-; RV64I-NEXT:    addi a2, a2, 1365
-; RV64I-NEXT:    and a1, a1, a2
-; RV64I-NEXT:    lui a2, 209715
-; RV64I-NEXT:    addi a2, a2, 819
-; RV64I-NEXT:    subw a0, a0, a1
-; RV64I-NEXT:    srliw a1, a0, 2
-; RV64I-NEXT:    and a0, a0, a2
-; RV64I-NEXT:    and a1, a1, a2
-; RV64I-NEXT:    lui a2, 61681
-; RV64I-NEXT:    add a0, a1, a0
-; RV64I-NEXT:    sraiw a1, a0, 4
-; RV64I-NEXT:    addw a0, a1, a0
-; RV64I-NEXT:    lui a1, 4112
-; RV64I-NEXT:    addiw a2, a2, -241
-; RV64I-NEXT:    and a0, a0, a2
-; RV64I-NEXT:    addiw a1, a1, 257
-; RV64I-NEXT:    call __muldi3
-; RV64I-NEXT:    srliw a0, a0, 24
-; RV64I-NEXT:    sltiu a0, a0, 2
-; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
-; RV64I-NEXT:    addi sp, sp, 16
+; RV64I-NEXT:    addiw a1, a0, -1
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    seqz a0, a0
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: ctpop_i32_ult_two:
diff --git a/llvm/test/CodeGen/RISCV/pr101786.ll b/llvm/test/CodeGen/RISCV/pr101786.ll
@@ -8,37 +8,9 @@ define i64 @test(i64 %x, ptr %p) {
 ; CHECK-NEXT:    li a0, 0
 ; CHECK-NEXT:    bgtz a2, .LBB0_3
 ; CHECK-NEXT:  # %bb.1: # %entry
-; CHECK-NEXT:    srli a3, a2, 1
-; CHECK-NEXT:    lui a4, 349525
-; CHECK-NEXT:    addiw a4, a4, 1365
-; CHECK-NEXT:    slli a5, a4, 32
-; CHECK-NEXT:    add a4, a4, a5
-; CHECK-NEXT:    and a3, a3, a4
-; CHECK-NEXT:    sub a2, a2, a3
-; CHECK-NEXT:    lui a3, 209715
-; CHECK-NEXT:    addiw a3, a3, 819
-; CHECK-NEXT:    slli a4, a3, 32
-; CHECK-NEXT:    add a3, a3, a4
-; CHECK-NEXT:    and a4, a2, a3
-; CHECK-NEXT:    srli a2, a2, 2
+; CHECK-NEXT:    addi a3, a2, -1
 ; CHECK-NEXT:    and a2, a2, a3
-; CHECK-NEXT:    add a2, a4, a2
-; CHECK-NEXT:    srli a3, a2, 4
-; CHECK-NEXT:    add a2, a2, a3
-; CHECK-NEXT:    lui a3, 61681
-; CHECK-NEXT:    addiw a3, a3, -241
-; CHECK-NEXT:    slli a4, a3, 32
-; CHECK-NEXT:    add a3, a3, a4
-; CHECK-NEXT:    and a2, a2, a3
-; CHECK-NEXT:    slli a3, a2, 8
-; CHECK-NEXT:    add a2, a2, a3
-; CHECK-NEXT:    slli a3, a2, 16
-; CHECK-NEXT:    add a2, a2, a3
-; CHECK-NEXT:    slli a3, a2, 32
-; CHECK-NEXT:    add a2, a2, a3
-; CHECK-NEXT:    srli a2, a2, 56
-; CHECK-NEXT:    li a3, 1
-; CHECK-NEXT:    bltu a3, a2, .LBB0_3
+; CHECK-NEXT:    bnez a2, .LBB0_3
 ; CHECK-NEXT:  # %bb.2: # %if.else
 ; CHECK-NEXT:    ld a0, 0(a1)
 ; CHECK-NEXT:  .LBB0_3: # %if.end
diff --git a/llvm/test/CodeGen/RISCV/rv32zbb.ll b/llvm/test/CodeGen/RISCV/rv32zbb.ll
@@ -571,12 +571,12 @@ define i64 @ctpop_i64(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_two:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    addi a2, a0, -1
-; RV32I-NEXT:    and a2, a0, a2
-; RV32I-NEXT:    seqz a0, a0
-; RV32I-NEXT:    sub a0, a1, a0
-; RV32I-NEXT:    and a0, a1, a0
-; RV32I-NEXT:    or a0, a2, a0
+; RV32I-NEXT:    seqz a2, a0
+; RV32I-NEXT:    addi a3, a0, -1
+; RV32I-NEXT:    sub a2, a1, a2
+; RV32I-NEXT:    and a0, a0, a3
+; RV32I-NEXT:    and a1, a1, a2
+; RV32I-NEXT:    or a0, a0, a1
 ; RV32I-NEXT:    seqz a0, a0
 ; RV32I-NEXT:    ret
 ;
@@ -595,12 +595,12 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_one:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    addi a2, a0, -1
-; RV32I-NEXT:    and a2, a0, a2
-; RV32I-NEXT:    seqz a0, a0
-; RV32I-NEXT:    sub a0, a1, a0
-; RV32I-NEXT:    and a0, a1, a0
-; RV32I-NEXT:    or a0, a2, a0
+; RV32I-NEXT:    seqz a2, a0
+; RV32I-NEXT:    addi a3, a0, -1
+; RV32I-NEXT:    sub a2, a1, a2
+; RV32I-NEXT:    and a0, a0, a3
+; RV32I-NEXT:    and a1, a1, a2
+; RV32I-NEXT:    or a0, a0, a1
 ; RV32I-NEXT:    snez a0, a0
 ; RV32I-NEXT:    ret
 ;
diff --git a/llvm/test/CodeGen/X86/pr94829.ll b/llvm/test/CodeGen/X86/pr94829.ll
@@ -4,30 +4,8 @@
 define ptr @test(i64 %x) {
 ; CHECK-LABEL: test:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    movq %rdi, %rax
-; CHECK-NEXT:    shrq %rax
-; CHECK-NEXT:    movabsq $6148914691236517205, %rcx # imm = 0x5555555555555555
-; CHECK-NEXT:    andq %rax, %rcx
-; CHECK-NEXT:    subq %rcx, %rdi
-; CHECK-NEXT:    movabsq $3689348814741910323, %rax # imm = 0x3333333333333333
-; CHECK-NEXT:    movq %rdi, %rcx
-; CHECK-NEXT:    andq %rax, %rcx
-; CHECK-NEXT:    shrq $2, %rdi
-; CHECK-NEXT:    andq %rax, %rdi
-; CHECK-NEXT:    addq %rcx, %rdi
-; CHECK-NEXT:    movq %rdi, %rax
-; CHECK-NEXT:    shrq $4, %rax
-; CHECK-NEXT:    addq %rdi, %rax
-; CHECK-NEXT:    movabsq $1085102592571150095, %rcx # imm = 0xF0F0F0F0F0F0F0F
-; CHECK-NEXT:    andq %rax, %rcx
-; CHECK-NEXT:    movabsq $72340172838076673, %rax # imm = 0x101010101010101
-; CHECK-NEXT:    imulq %rcx, %rax
-; CHECK-NEXT:    shrq $56, %rax
-; CHECK-NEXT:    cmpq $2, %rax
-; CHECK-NEXT:    jb .LBB0_2
-; CHECK-NEXT:  # %bb.1: # %if.else
-; CHECK-NEXT:    cmpl $2, %eax
-; CHECK-NEXT:  .LBB0_2: # %exit1
+; CHECK-NEXT:    leaq -1(%rdi), %rax
+; CHECK-NEXT:    testq %rax, %rdi
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    retq
 entry:
diff --git a/llvm/test/Transforms/CodeGenPrepare/unfold-pow2-test.ll b/llvm/test/Transforms/CodeGenPrepare/unfold-pow2-test.ll
@@ -12,8 +12,9 @@ define i64 @test_ult_2(i64 %x, i64 %y, i64 %a, i64 %b) {
 ; SLOW-LABEL: define i64 @test_ult_2(
 ; SLOW-SAME: i64 [[X:%.*]], i64 [[Y:%.*]], i64 [[A:%.*]], i64 [[B:%.*]]) {
 ; SLOW-NEXT:  [[ENTRY:.*]]:
-; SLOW-NEXT:    [[CTPOP:%.*]] = call i64 @llvm.ctpop.i64(i64 [[X]])
-; SLOW-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[CTPOP]], 1
+; SLOW-NEXT:    [[TMP0:%.*]] = add i64 [[X]], -1
+; SLOW-NEXT:    [[TMP1:%.*]] = and i64 [[X]], [[TMP0]]
+; SLOW-NEXT:    [[CMP1:%.*]] = icmp ne i64 [[TMP1]], 0
 ; SLOW-NEXT:    [[CMP2:%.*]] = icmp sgt i64 [[Y]], 0
 ; SLOW-NEXT:    [[CMP:%.*]] = or i1 [[CMP2]], [[CMP1]]
 ; SLOW-NEXT:    br i1 [[CMP]], label %[[IF_THEN:.*]], label %[[IF_END:.*]]
@@ -56,8 +57,9 @@ define i64 @test_ugt_1(i64 %x, i64 %y, i64 %a, i64 %b) {
 ; SLOW-LABEL: define i64 @test_ugt_1(
 ; SLOW-SAME: i64 [[X:%.*]], i64 [[Y:%.*]], i64 [[A:%.*]], i64 [[B:%.*]]) {
 ; SLOW-NEXT:  [[ENTRY:.*]]:
-; SLOW-NEXT:    [[CTPOP:%.*]] = call i64 @llvm.ctpop.i64(i64 [[X]])
-; SLOW-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[CTPOP]], 1
+; SLOW-NEXT:    [[TMP0:%.*]] = add i64 [[X]], -1
+; SLOW-NEXT:    [[TMP1:%.*]] = and i64 [[X]], [[TMP0]]
+; SLOW-NEXT:    [[CMP1:%.*]] = icmp ne i64 [[TMP1]], 0
 ; SLOW-NEXT:    [[CMP2:%.*]] = icmp sgt i64 [[Y]], 0
 ; SLOW-NEXT:    [[CMP:%.*]] = or i1 [[CMP2]], [[CMP1]]
 ; SLOW-NEXT:    br i1 [[CMP]], label %[[IF_THEN:.*]], label %[[IF_END:.*]]
