[CodeGenPrepare] Convert ctpop(X) ==/!= 1 -> ctpop(X) u</u> 2/1

dtcxzyw · dtcxzyw · commit bf11b6d37833 · 2024-10-10T22:44:23.000+08:00
diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -2111,6 +2111,32 @@ bool CodeGenPrepare::optimizeURem(Instruction *Rem) {
   return false;
 }
 
+/// Some targets have better codegen for `ctpop(X) u< 2` than `ctpop(X) == 1`.
+/// This function converts `ctpop(X) ==/!= 1` into `ctpop(X) u</u> 2/1` if the
+/// result cannot be zero.
+static bool adjustIsPower2Test(CmpInst *Cmp) {
+  ICmpInst::Predicate Pred;
+  if (!match(Cmp, m_ICmp(Pred, m_Intrinsic<Intrinsic::ctpop>(), m_One())))
+    return false;
+  if (!ICmpInst::isEquality(Pred))
+    return false;
+  auto *II = cast<IntrinsicInst>(Cmp->getOperand(0));
+  if (auto Range = II->getRange()) {
+    Type *Ty = II->getType();
+    unsigned BitWidth = Ty->getScalarSizeInBits();
+    if (Range->contains(APInt::getZero(BitWidth)))
+      return false;
+
+    if (Pred == ICmpInst::ICMP_EQ) {
+      Cmp->setPredicate(ICmpInst::ICMP_ULT);
+      Cmp->setOperand(1, ConstantInt::get(Ty, 2));
+    } else
+      Cmp->setPredicate(ICmpInst::ICMP_UGT);
+    return true;
+  }
+  return false;
+}
+
 bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (sinkCmpExpression(Cmp, *TLI))
     return true;
@@ -2130,6 +2156,9 @@ bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (foldFCmpToFPClassTest(Cmp, *TLI, *DL))
     return true;
 
+  if (adjustIsPower2Test(Cmp))
+    return true;
+
   return false;
 }
 
diff --git a/llvm/test/CodeGen/AArch64/arm64-popcnt.ll b/llvm/test/CodeGen/AArch64/arm64-popcnt.ll
@@ -282,23 +282,21 @@ define i1 @ctpop32_eq_one_nonzero(i32 %x) {
 ; CHECK-LABEL: ctpop32_eq_one_nonzero:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w0, #1
-; CHECK-NEXT:    eor w9, w0, w8
-; CHECK-NEXT:    cmp w9, w8
-; CHECK-NEXT:    cset w0, hi
+; CHECK-NEXT:    tst w0, w8
+; CHECK-NEXT:    cset w0, eq
 ; CHECK-NEXT:    ret
 ;
 ; CHECK-NONEON-LABEL: ctpop32_eq_one_nonzero:
 ; CHECK-NONEON:       // %bb.0: // %entry
 ; CHECK-NONEON-NEXT:    sub w8, w0, #1
-; CHECK-NONEON-NEXT:    eor w9, w0, w8
-; CHECK-NONEON-NEXT:    cmp w9, w8
-; CHECK-NONEON-NEXT:    cset w0, hi
+; CHECK-NONEON-NEXT:    tst w0, w8
+; CHECK-NONEON-NEXT:    cset w0, eq
 ; CHECK-NONEON-NEXT:    ret
 ;
 ; CHECK-CSSC-LABEL: ctpop32_eq_one_nonzero:
 ; CHECK-CSSC:       // %bb.0: // %entry
-; CHECK-CSSC-NEXT:    cnt w8, w0
-; CHECK-CSSC-NEXT:    cmp w8, #1
+; CHECK-CSSC-NEXT:    sub w8, w0, #1
+; CHECK-CSSC-NEXT:    tst w0, w8
 ; CHECK-CSSC-NEXT:    cset w0, eq
 ; CHECK-CSSC-NEXT:    ret
 entry:
@@ -311,23 +309,21 @@ define i1 @ctpop32_ne_one_nonzero(i32 %x) {
 ; CHECK-LABEL: ctpop32_ne_one_nonzero:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w0, #1
-; CHECK-NEXT:    eor w9, w0, w8
-; CHECK-NEXT:    cmp w9, w8
-; CHECK-NEXT:    cset w0, ls
+; CHECK-NEXT:    tst w0, w8
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
 ;
 ; CHECK-NONEON-LABEL: ctpop32_ne_one_nonzero:
 ; CHECK-NONEON:       // %bb.0: // %entry
 ; CHECK-NONEON-NEXT:    sub w8, w0, #1
-; CHECK-NONEON-NEXT:    eor w9, w0, w8
-; CHECK-NONEON-NEXT:    cmp w9, w8
-; CHECK-NONEON-NEXT:    cset w0, ls
+; CHECK-NONEON-NEXT:    tst w0, w8
+; CHECK-NONEON-NEXT:    cset w0, ne
 ; CHECK-NONEON-NEXT:    ret
 ;
 ; CHECK-CSSC-LABEL: ctpop32_ne_one_nonzero:
 ; CHECK-CSSC:       // %bb.0: // %entry
-; CHECK-CSSC-NEXT:    cnt w8, w0
-; CHECK-CSSC-NEXT:    cmp w8, #1
+; CHECK-CSSC-NEXT:    sub w8, w0, #1
+; CHECK-CSSC-NEXT:    tst w0, w8
 ; CHECK-CSSC-NEXT:    cset w0, ne
 ; CHECK-CSSC-NEXT:    ret
 entry:
diff --git a/llvm/test/CodeGen/RISCV/rv32zbb.ll b/llvm/test/CodeGen/RISCV/rv32zbb.ll
@@ -1446,15 +1446,14 @@ define i1 @ctpop32_eq_one_nonzero(i32 %x) {
 ; RV32I-LABEL: ctpop32_eq_one_nonzero:
 ; RV32I:       # %bb.0: # %entry
 ; RV32I-NEXT:    addi a1, a0, -1
-; RV32I-NEXT:    xor a0, a0, a1
-; RV32I-NEXT:    sltu a0, a1, a0
+; RV32I-NEXT:    and a0, a0, a1
+; RV32I-NEXT:    seqz a0, a0
 ; RV32I-NEXT:    ret
 ;
 ; RV32ZBB-LABEL: ctpop32_eq_one_nonzero:
 ; RV32ZBB:       # %bb.0: # %entry
 ; RV32ZBB-NEXT:    cpop a0, a0
-; RV32ZBB-NEXT:    addi a0, a0, -1
-; RV32ZBB-NEXT:    seqz a0, a0
+; RV32ZBB-NEXT:    sltiu a0, a0, 2
 ; RV32ZBB-NEXT:    ret
 entry:
   %popcnt = call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
@@ -1466,16 +1465,15 @@ define i1 @ctpop32_ne_one_nonzero(i32 %x) {
 ; RV32I-LABEL: ctpop32_ne_one_nonzero:
 ; RV32I:       # %bb.0: # %entry
 ; RV32I-NEXT:    addi a1, a0, -1
-; RV32I-NEXT:    xor a0, a0, a1
-; RV32I-NEXT:    sltu a0, a1, a0
-; RV32I-NEXT:    xori a0, a0, 1
+; RV32I-NEXT:    and a0, a0, a1
+; RV32I-NEXT:    snez a0, a0
 ; RV32I-NEXT:    ret
 ;
 ; RV32ZBB-LABEL: ctpop32_ne_one_nonzero:
 ; RV32ZBB:       # %bb.0: # %entry
 ; RV32ZBB-NEXT:    cpop a0, a0
-; RV32ZBB-NEXT:    addi a0, a0, -1
-; RV32ZBB-NEXT:    snez a0, a0
+; RV32ZBB-NEXT:    sltiu a0, a0, 2
+; RV32ZBB-NEXT:    xori a0, a0, 1
 ; RV32ZBB-NEXT:    ret
 entry:
   %popcnt = tail call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
diff --git a/llvm/test/CodeGen/RISCV/rv64zbb.ll b/llvm/test/CodeGen/RISCV/rv64zbb.ll
@@ -1622,17 +1622,16 @@ entry:
 define i1 @ctpop32_eq_one_nonzero(i32 %x) {
 ; RV64I-LABEL: ctpop32_eq_one_nonzero:
 ; RV64I:       # %bb.0: # %entry
-; RV64I-NEXT:    addiw a1, a0, -1
-; RV64I-NEXT:    xor a0, a0, a1
+; RV64I-NEXT:    addi a1, a0, -1
+; RV64I-NEXT:    and a0, a0, a1
 ; RV64I-NEXT:    sext.w a0, a0
-; RV64I-NEXT:    sltu a0, a1, a0
+; RV64I-NEXT:    seqz a0, a0
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: ctpop32_eq_one_nonzero:
 ; RV64ZBB:       # %bb.0: # %entry
 ; RV64ZBB-NEXT:    cpopw a0, a0
-; RV64ZBB-NEXT:    addi a0, a0, -1
-; RV64ZBB-NEXT:    seqz a0, a0
+; RV64ZBB-NEXT:    sltiu a0, a0, 2
 ; RV64ZBB-NEXT:    ret
 entry:
   %popcnt = call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
@@ -1643,18 +1642,17 @@ entry:
 define i1 @ctpop32_ne_one_nonzero(i32 %x) {
 ; RV64I-LABEL: ctpop32_ne_one_nonzero:
 ; RV64I:       # %bb.0: # %entry
-; RV64I-NEXT:    addiw a1, a0, -1
-; RV64I-NEXT:    xor a0, a0, a1
+; RV64I-NEXT:    addi a1, a0, -1
+; RV64I-NEXT:    and a0, a0, a1
 ; RV64I-NEXT:    sext.w a0, a0
-; RV64I-NEXT:    sltu a0, a1, a0
-; RV64I-NEXT:    xori a0, a0, 1
+; RV64I-NEXT:    snez a0, a0
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: ctpop32_ne_one_nonzero:
 ; RV64ZBB:       # %bb.0: # %entry
 ; RV64ZBB-NEXT:    cpopw a0, a0
-; RV64ZBB-NEXT:    addi a0, a0, -1
-; RV64ZBB-NEXT:    snez a0, a0
+; RV64ZBB-NEXT:    sltiu a0, a0, 2
+; RV64ZBB-NEXT:    xori a0, a0, 1
 ; RV64ZBB-NEXT:    ret
 entry:
   %popcnt = tail call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
@@ -1666,15 +1664,14 @@ define i1 @ctpop64_eq_one_nonzero(i64 %x) {
 ; RV64I-LABEL: ctpop64_eq_one_nonzero:
 ; RV64I:       # %bb.0: # %entry
 ; RV64I-NEXT:    addi a1, a0, -1
-; RV64I-NEXT:    xor a0, a0, a1
-; RV64I-NEXT:    sltu a0, a1, a0
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    seqz a0, a0
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: ctpop64_eq_one_nonzero:
 ; RV64ZBB:       # %bb.0: # %entry
 ; RV64ZBB-NEXT:    cpop a0, a0
-; RV64ZBB-NEXT:    addi a0, a0, -1
-; RV64ZBB-NEXT:    seqz a0, a0
+; RV64ZBB-NEXT:    sltiu a0, a0, 2
 ; RV64ZBB-NEXT:    ret
 entry:
   %popcnt = call range(i64 1, 65) i64 @llvm.ctpop.i64(i64 %x)
diff --git a/llvm/test/CodeGen/X86/ispow2.ll b/llvm/test/CodeGen/X86/ispow2.ll
@@ -223,29 +223,39 @@ define <4 x i1> @neither_pow2_non_zero_4xv64_x_maybe_z(<4 x i64> %x) {
 
 
 define i1 @ctpop32_eq_one_nonzero(i32 %x) {
-; CHECK-LABEL: ctpop32_eq_one_nonzero:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
-; CHECK-NEXT:    leal -1(%rdi), %eax
-; CHECK-NEXT:    xorl %eax, %edi
-; CHECK-NEXT:    cmpl %eax, %edi
-; CHECK-NEXT:    seta %al
-; CHECK-NEXT:    retq
+; CHECK-NOBMI-LABEL: ctpop32_eq_one_nonzero:
+; CHECK-NOBMI:       # %bb.0: # %entry
+; CHECK-NOBMI-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NOBMI-NEXT:    leal -1(%rdi), %eax
+; CHECK-NOBMI-NEXT:    testl %eax, %edi
+; CHECK-NOBMI-NEXT:    sete %al
+; CHECK-NOBMI-NEXT:    retq
+;
+; CHECK-BMI2-LABEL: ctpop32_eq_one_nonzero:
+; CHECK-BMI2:       # %bb.0: # %entry
+; CHECK-BMI2-NEXT:    blsrl %edi, %eax
+; CHECK-BMI2-NEXT:    sete %al
+; CHECK-BMI2-NEXT:    retq
 entry:
   %popcnt = call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
   %cmp = icmp eq i32 %popcnt, 1
   ret i1 %cmp
 }
 
 define i1 @ctpop32_ne_one_nonzero(i32 %x) {
-; CHECK-LABEL: ctpop32_ne_one_nonzero:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
-; CHECK-NEXT:    leal -1(%rdi), %eax
-; CHECK-NEXT:    xorl %eax, %edi
-; CHECK-NEXT:    cmpl %eax, %edi
-; CHECK-NEXT:    setbe %al
-; CHECK-NEXT:    retq
+; CHECK-NOBMI-LABEL: ctpop32_ne_one_nonzero:
+; CHECK-NOBMI:       # %bb.0: # %entry
+; CHECK-NOBMI-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NOBMI-NEXT:    leal -1(%rdi), %eax
+; CHECK-NOBMI-NEXT:    testl %eax, %edi
+; CHECK-NOBMI-NEXT:    setne %al
+; CHECK-NOBMI-NEXT:    retq
+;
+; CHECK-BMI2-LABEL: ctpop32_ne_one_nonzero:
+; CHECK-BMI2:       # %bb.0: # %entry
+; CHECK-BMI2-NEXT:    blsrl %edi, %eax
+; CHECK-BMI2-NEXT:    setne %al
+; CHECK-BMI2-NEXT:    retq
 entry:
   %popcnt = tail call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 %x)
   %cmp = icmp ne i32 %popcnt, 1
