[InstCombine] Generate better code for std::bit_floor from libstdc++

[kazu]

Summary

Without this patch, std::bit_floor<uint32_t> in libstdc++ is compiled
as:

%eq0 = icmp eq i32 %x, 0
%lshr = lshr i32 %x, 1
%ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
%sub = sub i32 32, %ctlz
%shl = shl i32 1, %sub
%sel = select i1 %eq0, i32 0, i32 %shl

With this patch:

%eq0 = icmp eq i32 %x, 0
%ctlz = call i32 @llvm.ctlz.i32(i32 %x, i1 false)
%lshr = lshr i32 -2147483648, %1
%sel = select i1 %eq0, i32 0, i32 %lshr

This patch recognizes the specific pattern emitted for std::bit_floor
in libstdc++.

https://alive2.llvm.org/ce/z/piMdFX

This patch fixes:

https://github.com/llvm/llvm-project/issues/61183

[kazu]

Combine some "if" statements.

[kazu]

Please take a look. Thanks!

[RKSimon]

please can you add vector test coverage?

[kazu]

Added a vector test and some negative tests.

[kazu]

In D145890#4209037, @RKSimon wrote:

please can you add vector test coverage?

I just checked in:

https://github.com/llvm/llvm-project/commit/f9f4767af9f3d89792d67ae8c5f65913ff263b89

[kazu]

Please take look. Thanks!

[kazu]

Rebased.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3296

'32' -> 'BitWidth'.

3304

I would prefer this as a more generic function for replacing 1 << (BitWidth - ctlz(X >> 1)) -> (SignBit >> ctlz(X) which we can do anytime we know X is non-zero: https://alive2.llvm.org/ce/z/PDxghx

Then use the helper for the select case, but then it can also be used with an IsKnownNonZero query.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3323

Instead of matching exact BitWidth can you implement this more generically: https://alive2.llvm.org/ce/z/q2hCjC

1 << (C - ctlz(X >> 1)) -> ((1 << (C-1)) >> ctlz(X)

[kazu]

Addressed feedback.

[kazu]

Please take a look. Thanks!

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3312

Doesn't need to be added now, but please add TODO for C==0 case if we have a mask i.e:
https://alive2.llvm.org/ce/z/epLwX_

In general, you could either match C - CTZL with C > 0 && C <= BitWidth or match (C - CTLZ) % BitWidth for any C (then set bit (C->getZExtValue() - 1) % BitWidth.

3318

APInt ShiftedBit = APInt::getOneBitSet(BitWidth, C->getZExtValue() - 1);

3325

This can be exact:
https://alive2.llvm.org/ce/z/Enurag

I think:

BinaryOperator *BO = BinaryOperator::CreateLShr(ConstantInt::get(NType, ShiftedBit), NewCTLZ);
  BO->setIsExact();
  return BO;

Should work.

[kazu]

Start using getOneBitSet.

Add a TODO comment.

[kazu]

Please take a look. Thanks!

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3325

If I do this, I get:

llvm/include/llvm/Transforms/Utils/InstructionWorklist.h:60: void llvm::InstructionWorklist::push(llvm::Instruction *): Assertion `I->getParent() && "Instruction not inserted yet?"' failed.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3325

If I do this, I get:

llvm/include/llvm/Transforms/Utils/InstructionWorklist.h:60: void llvm::InstructionWorklist::push(llvm::Instruction *): Assertion `I->getParent() && "Instruction not inserted yet?"' failed.

Ahh, because you use replaceoperand below I believe.

Then just cast what you are turning to a BinaryOperator and setIsExact.

[kazu]

Call setIsExact on Instruction.

[kazu]

Please take a look. Thanks!

[goldstein.w.n]

LGTM.

I'm not a maintainer so please way a few days to push to give others a chance
to look it over.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3302

One more note:

X doesn't actually need to be nonzero. Either X nonzero OR ctlz with zero->poison (second arg is true):
https://alive2.llvm.org/ce/z/gDJ_yT

Can you just add that is a todo?

[kazu]

Add a comment to relax the requirement that X be nonzero.

[kazu]

Thank you for the review!

I'm not a maintainer so please way a few days to push to give others a chance
to look it over.

Will do.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

static Instruction *foldBitCeil(SelectInst &SI, IRBuilderBase &Builder) {
   // hardware as part of the shift instruction.
   Value *Neg = Builder.CreateNeg(Ctlz);
   Value *Masked =
       Builder.CreateAnd(Neg, ConstantInt::get(SelType, BitWidth - 1));
   return BinaryOperator::Create(Instruction::Shl, ConstantInt::get(SelType, 1),
                                 Masked);
 }
 
+// Transform:
+//
+//   1 << (C - ctlz(X >> 1))
+//
+// into
+//
+//   (1 << (C - 1)) >> ctlz(X)
+//
+// The caller must guarantee that X is nonzero.
+//
+// TODO: Relax the requirement that X be nonzero.  We just need to require X to
+// be nonzero or the second argument of CTLZ to be true (that is, returning
+// poison on zero).
+static Instruction *foldBitFloorNonzero(Value *N, Value *X,
+                                        InstCombiner::BuilderTy &Builder) {
+  Type *NType = N->getType();
+  unsigned BitWidth = NType->getScalarSizeInBits();
+
+  // Match C - ctlz(X >> 1), where C is in (0, BitWidth].
+  // TODO: Handle C in [0, BitWidth] (with 0 included in the range), in which
+  // case 1 << C - ctlz(X >> 1) is equivalent to
+  // (1 << ((C - 1) & (BitWidth - 1))) >> ctlz(X).
+  const APInt *C = nullptr;
+  Value *CTLZ;
+  if (!match(N, m_OneUse(m_Shl(m_One(),
+                               m_OneUse(m_Sub(m_APInt(C), m_Value(CTLZ)))))) ||
+      !(C->ugt(0) && C->ule(BitWidth)) ||
+      !match(CTLZ, m_OneUse(m_Intrinsic<Intrinsic::ctlz>(
+                       m_OneUse(m_LShr(m_Specific(X), m_One())), m_Zero()))))
+    return nullptr;
+
+  APInt ShiftedBit = APInt::getOneBitSet(BitWidth, C->getZExtValue() - 1);
+
+  // Build ShiftedBit >> CTLZ.
+  Value *NewCTLZ =
+      Builder.CreateIntrinsic(Intrinsic::ctlz, {CTLZ->getType()},
+                              {X, cast<Instruction>(CTLZ)->getOperand(1)});
+  auto *Shift = cast<Instruction>(
+      Builder.CreateLShr(ConstantInt::get(NType, ShiftedBit), NewCTLZ));
+  Shift->setIsExact();
+  return Shift;
+}
+
+// Transform:
+//
+//   X == 0 ? 0 : (1 << (C1 - ctlz(X >> 1)))
+//
+// into
+//
+//   X == 0 ? 0 : (C2 >> ctlz(X))
+//
+// where C2 is computed by foldBitFloorNonzero based on C1.  The caller is
+// responsible for replacing one of the select operands.
+static Instruction *foldBitFloor(SelectInst &SI,
+                                 InstCombiner::BuilderTy &Builder) {
+  Value *TrueVal = SI.getTrueValue();
+  Value *FalseVal = SI.getFalseValue();
+
+  ICmpInst::Predicate Pred;
+  Value *Cond0;
+  if (!match(SI.getCondition(), m_ICmp(Pred, m_Value(Cond0), m_Zero())) ||
+      !ICmpInst::isEquality(Pred))
+    return nullptr;
+
+  if (Pred == ICmpInst::ICMP_NE)
+    std::swap(TrueVal, FalseVal);
+
+  if (!match(TrueVal, m_Zero()))
+    return nullptr;
+
+  return foldBitFloorNonzero(FalseVal, Cond0, Builder);
+}
+
 Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   if (Value *V = simplifySelectInst(CondVal, TrueVal, FalseVal,
                                     SQ.getWithInstruction(&SI)))
Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   //   (~x) | y  -->  ~(x & (~y))
   //   (~x) & y  -->  ~(x | (~y))
   if (sinkNotIntoOtherHandOfLogicalOp(SI))
     return &SI;
 
   if (Instruction *I = foldBitCeil(SI, Builder))
     return I;
 
+  if (Instruction *I = foldBitFloor(SI, Builder))
+    return replaceOperand(SI, match(SI.getTrueValue(), m_Zero()) ? 2 : 1, I);
+
   return nullptr;
 }

llvm/test/Transforms/InstCombine/bit_floor.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 define i32 @bit_floor_32(i32 %x) {
 ; CHECK-LABEL: @bit_floor_32(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0:![0-9]+]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[SHL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.ctlz.i32(i32 [[X]], i1 false), !range [[RNG0:![0-9]+]]
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i32 -2147483648, [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %eq0 = icmp eq i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %eq0, i32 0, i32 %shl
   ret i32 %sel
 }
 
 define i64 @bit_floor_64(i64 %x) {
 ; CHECK-LABEL: @bit_floor_64(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i64 [[X:%.*]], 0
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i64 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i64 @llvm.ctlz.i64(i64 [[LSHR]], i1 false), !range [[RNG1:![0-9]+]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i64 64, [[CTLZ]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i64 1, [[SUB]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i64 0, i64 [[SHL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 false), !range [[RNG1:![0-9]+]]
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i64 -9223372036854775808, [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    ret i64 [[SEL]]
 ;
   %eq0 = icmp eq i64 %x, 0
   %lshr = lshr i64 %x, 1
   %ctlz = tail call i64 @llvm.ctlz.i64(i64 %lshr, i1 false)
   %sub = sub i64 64, %ctlz
   %shl = shl i64 1, %sub
   %sel = select i1 %eq0, i64 0, i64 %shl
   ret i64 %sel
 }
 
 ; Commutted select operands should still be recognized.
 define i32 @bit_floor_commuted_operands(i32 %x) {
 ; CHECK-LABEL: @bit_floor_commuted_operands(
 ; CHECK-NEXT:    [[NE0_NOT:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[NE0_NOT]], i32 0, i32 [[SHL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.ctlz.i32(i32 [[X]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact i32 -2147483648, [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[NE0_NOT]], i32 0, i32 [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ne0 = icmp ne i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %ne0, i32 %shl, i32 0
   ret i32 %sel
 }
 
 ; Negative test: lshr used twice
 define i32 @bit_floor_lshr_used_twice(i32 %x, ptr %p) {
 ; CHECK-LABEL: @bit_floor_lshr_used_twice(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i32 [[X:%.*]], 0
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG2:![0-9]+]]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[SHL]]
 ; CHECK-NEXT:    store i32 [[LSHR]], ptr [[P:%.*]], align 4
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %eq0 = icmp eq i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %eq0, i32 0, i32 %shl
   store i32 %lshr, ptr %p, align 4
   ret i32 %sel
 }
 
 ; Negative test: ctlz used twice
 define i32 @bit_floor_ctlz_used_twice(i32 %x, ptr %p) {
 ; CHECK-LABEL: @bit_floor_ctlz_used_twice(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i32 [[X:%.*]], 0
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG2]]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[SHL]]
 ; CHECK-NEXT:    store i32 [[CTLZ]], ptr [[P:%.*]], align 4
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %eq0 = icmp eq i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %eq0, i32 0, i32 %shl
   store i32 %ctlz, ptr %p, align 4
   ret i32 %sel
 }
 
 ; Negative test: sub used twice
 define i32 @bit_floor_sub_used_twice(i32 %x, ptr %p) {
 ; CHECK-LABEL: @bit_floor_sub_used_twice(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i32 [[X:%.*]], 0
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG2]]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[SHL]]
 ; CHECK-NEXT:    store i32 [[SUB]], ptr [[P:%.*]], align 4
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %eq0 = icmp eq i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %eq0, i32 0, i32 %shl
   store i32 %sub, ptr %p, align 4
   ret i32 %sel
 }
 
 ; Negative test: shl used twice
 define i32 @bit_floor_shl_used_twice(i32 %x, ptr %p) {
 ; CHECK-LABEL: @bit_floor_shl_used_twice(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i32 [[X:%.*]], 0
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[CTLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[LSHR]], i1 false), !range [[RNG2]]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i32 32, [[CTLZ]]
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i32 0, i32 [[SHL]]
 ; CHECK-NEXT:    store i32 [[SHL]], ptr [[P:%.*]], align 4
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %eq0 = icmp eq i32 %x, 0
   %lshr = lshr i32 %x, 1
   %ctlz = tail call i32 @llvm.ctlz.i32(i32 %lshr, i1 false)
   %sub = sub i32 32, %ctlz
   %shl = shl i32 1, %sub
   %sel = select i1 %eq0, i32 0, i32 %shl
   store i32 %shl, ptr %p, align 4
   ret i32 %sel
 }
 
 ; a vector version of @bit_floor_32 above
 define <4 x i32> @bit_floor_v4i32(<4 x i32> %x) {
 ; CHECK-LABEL: @bit_floor_v4i32(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq <4 x i32> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <4 x i32> [[X]], <i32 1, i32 1, i32 1, i32 1>
-; CHECK-NEXT:    [[CTLZ:%.*]] = tail call <4 x i32> @llvm.ctlz.v4i32(<4 x i32> [[LSHR]], i1 false), !range [[RNG0]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw <4 x i32> <i32 32, i32 32, i32 32, i32 32>, [[CTLZ]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl nuw <4 x i32> <i32 1, i32 1, i32 1, i32 1>, [[SUB]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[EQ0]], <4 x i32> zeroinitializer, <4 x i32> [[SHL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.ctlz.v4i32(<4 x i32> [[X]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr exact <4 x i32> <i32 -2147483648, i32 -2147483648, i32 -2147483648, i32 -2147483648>, [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[EQ0]], <4 x i32> zeroinitializer, <4 x i32> [[TMP2]]
 ; CHECK-NEXT:    ret <4 x i32> [[SEL]]
 ;
   %eq0 = icmp eq <4 x i32> %x, <i32 0, i32 0, i32 0, i32 0>
   %lshr = lshr <4 x i32> %x, <i32 1, i32 1, i32 1, i32 1>
   %ctlz = tail call <4 x i32> @llvm.ctlz.v4i32(<4 x i32> %lshr, i1 false)
   %sub = sub <4 x i32> <i32 32, i32 32, i32 32, i32 32>, %ctlz
   %shl = shl <4 x i32> <i32 1, i32 1, i32 1, i32 1>, %sub
   %sel = select <4 x i1> %eq0, <4 x i32> <i32 0, i32 0, i32 0, i32 0>, <4 x i32> %shl
   ret <4 x i32> %sel
 }
 
 declare i32 @llvm.ctlz.i32(i32, i1 immarg)
 declare i64 @llvm.ctlz.i64(i64, i1 immarg)
 declare <4 x i32> @llvm.ctlz.v4i32(<4 x i32>, i1)