[InstCombine] Fold and-reduce idiom

[mkazantsev]

Summary

This patch introduces folding of and-reduce idiom and generates code
that is easier to read and which is lest costly in terms of icmp operations.
The folding is

icmp eq (bitcast(icmp ne (lhs, rhs)), 0)

into

icmp eq(bitcast(lhs), bitcast(rhs))

See https://github.com/llvm/llvm-project/issues/53419

[mkazantsev]

Sanity check: https://alive2.llvm.org/ce/z/5JhbEz

[craig.topper]

Description says "add-reduce"

[mkazantsev]

This should partly (on IR level) address https://github.com/llvm/llvm-project/issues/53419

[mkazantsev]

In D118317#3274826, @craig.topper wrote:

Description says "add-reduce"

Thanks, fixed!

[RKSimon]

TBH I'm not sure about converting legal vector types (e.g. <8 x i64>) into massive integer (i512).

[mkazantsev]

You're right, need to check if the new scalar type is a legal type. I'll add DL & negative tests & rebase on top. Thanks!

[lebedev.ri]

This comment was removed by lebedev.ri.

[mkazantsev]

Added DL check that we only produce legal scalar types & rebased on top of new (negative) tests.

[lebedev.ri]

I think it would make most sense to do at backend level, but i suppose it would be fine to do it here too.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5885

Add a comment explaining what this does?

5894

This should probably be InnerICmp

5897–5898

I think the LHSTy should just be hoisted here, and used instead.

5907

I suspect this should be relaxed to something like "is less than a legal integer"

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5907

I'm not sure what happens if smth like i63 lives through the pipeline to codegen. Is it ok?

[mkazantsev]

Addressed Roman's comments:

• LHS type hoisted
• Comments added
• Vars renamed
• As for relaxation of legal type to "not wider than max legal type" - I could not find proper API easily. Maybe it needs some extra infrastructure work, and need to check legality. So I added TODO, lets postpone this. I think in most practical cases vector width is power of 2 and all such types are legal.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5885

5887

5907

Not sure, but at least something like "is a power of two that is smaller than the largest legal integer" should be fine.

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5885

Ahaha, good catch! :)

5907

I didn't find good API to get this largest legal integer. Let's postpone.

[mkazantsev]

Fixed typos.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5910

Why not just use NumBits < getMinVectorRegisterBitWidth() ?

[spatel]

llvm/test/Transforms/InstCombine/icmp-vec.ll

431

If this had a power-of-2 type with this pattern, would we handle it? Would be good to add a few more test comments, so the intent is clear.
https://alive2.llvm.org/ce/z/rDeaif

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5910

And what if some platform has scalar types wider than vector types? Theoretically, it's possible I think. Or, say,i66 is less than i128, should be be allowed?

I suggest to leave this thing alone for now. We can have a separate discussion on what is the less restrictive check here. I really don't want this to be blocked by this particular thing.

llvm/test/Transforms/InstCombine/icmp-vec.ll

431

Yup. It was going to make this as 2nd step of this activity and support or-base ne pattern.

431

I mean, it's a slightly different pattern with -1 in the original cmp, but we can support it in a very similar manner.

[spatel]

llvm/test/Transforms/InstCombine/icmp-vec.ll

431

I think we'll always canonicalize it to a 0 constant (as shown in the current CHECK lines) - only the predicate is different, so it would be a very small adjustment to make this work?

[mkazantsev]

llvm/test/Transforms/InstCombine/icmp-vec.ll

431

Sure. Do you mind if I do it in a separate follow-up patch? Need to precommit tests anyways.

[lebedev.ri]

The legality check probably needs refinement, but other than that this seems fine to me.
@spatel ?

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5922

I'm guessing you are expecting to handle other patterns?
Otherwise, early-return.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5902

Is it worth using the pattern match helpers to match all of this?

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5902

That's what I tried initially, but for single use checks we'd still need to extract all intermediate values and cast them to instructions.

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5902

Let's give it another try, maybe will be better after all refactoring we've done...

[mkazantsev]

NFC refactoring. Used pattern match following @RKSimon's comment. Looks not so ugly, and it makes support of further transform proposed by @spatel easier.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5922

Yes, there will be others. One of them was mentioned by Sanjay already.

[mkazantsev]

NFC: replaced m_Constant with m_Zero as it seems that it'll be enough for all planned patterns.

[spatel]

LGTM - if you don't plan to extend the predicate match to isEquality soon, please put a TODO comment on it.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5915

You could roll these checks into the match line above with m_OneUse()...and that removes the need to capture the bitcast and inner icmp?

[mkazantsev]

I'll extend the predicate in the follow-up later today, thanks!

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5915

I didn't know we had this, thanks!

[hans]

This is causing asserts to fire in Chromium builds. See https://bugs.chromium.org/p/chromium/issues/detail?id=1291929#c2 for a reproducer.

I've reverted back to green in fabaca10b86f77f7d2d34db91fa6b284da924395 until it can be fixed.

[mkazantsev]

Thanks for the repro, I'll take a look.

[spatel]

I didn't look at the crasher, but in hindsight, it seems clear that we need to check the bitcast source types more thoroughly.
This will crash:

define i1 @eq_cast_eq-1(<2 x i4*> %x, <2 x i4*> %y) {
  %ic = icmp eq <2 x i4*> %x, %y
  %b = bitcast <2 x i1> %ic to i2
  %r = icmp eq i2 %b, -1
  ret i1 %r
}

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if (match(I.getOperand(0), m_Instruction(Op)) &&
       match(I.getOperand(1), m_Zero()) &&
       Op->isLaunderOrStripInvariantGroup()) {
     return ICmpInst::Create(Instruction::ICmp, I.getPredicate(),
                             Op->getOperand(0), I.getOperand(1));
   }
   return nullptr;
 }
 
+/// This function folds patterns produced by lowering of reduce idioms, such as
+/// llvm.vector.reduce.and which are lowered into instruction chains. This code
+/// attempts to generate fewer number of scalar comparisons instead of vector
+/// comparisons when possible.
+static Instruction *foldReductionIdiom(ICmpInst &I,
+                                       InstCombiner::BuilderTy &Builder,
+                                       const DataLayout &DL) {
+  if (I.getType()->isVectorTy())
+    return nullptr;
+  ICmpInst::Predicate OuterPred, InnerPred;
+  Value *LHS, *RHS;
+
+  // Match lowering of @llvm.vector.reduce.and. Turn
+  ///   %vec_ne = icmp ne <8 x i8> %lhs, %rhs
+  ///   %scalar_ne = bitcast <8 x i1> %vec_ne to i8
+  ///   %all_eq = icmp eq i8 %scalar_ne, 0
+  ///
+  /// into
+  ///
+  ///   %lhs.scalar = bitcast <8 x i8> %lhs to i64
+  ///   %rhs.scalar = bitcast <8 x i8> %rhs to i64
+  ///   %all_eq = icmp eq i64 %lhs.scalar, %rhs.scalar
+  if (!match(&I, m_ICmp(OuterPred,
+                        m_OneUse(m_BitCast(m_OneUse(
+                            m_ICmp(InnerPred, m_Value(LHS), m_Value(RHS))))),
+                        m_Zero())))
+    return nullptr;
+  auto *LHSTy = dyn_cast<FixedVectorType>(LHS->getType());
+  if (!LHSTy || !LHSTy->getElementType()->isIntegerTy())
+    return nullptr;
+  unsigned NumBits =
+      LHSTy->getNumElements() * LHSTy->getElementType()->getIntegerBitWidth();
+  // TODO: Relax this to "not wider than max legal integer type"?
+  if (!DL.isLegalInteger(NumBits))
+    return nullptr;
+
+  // TODO: Generalize to isEquality and support other patterns.
+  if (OuterPred == ICmpInst::ICMP_EQ && InnerPred == ICmpInst::ICMP_NE) {
+    auto *ScalarTy = Builder.getIntNTy(NumBits);
+    LHS = Builder.CreateBitCast(LHS, ScalarTy, LHS->getName() + ".scalar");
+    RHS = Builder.CreateBitCast(RHS, ScalarTy, RHS->getName() + ".scalar");
+    return ICmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, LHS, RHS,
+                            I.getName());
+  }
+
+  return nullptr;
+}
+
 Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
   bool Changed = false;
   const SimplifyQuery Q = SQ.getWithInstruction(&I);
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   unsigned Op0Cplxity = getComplexity(Op0);
   unsigned Op1Cplxity = getComplexity(Op1);
 
   /// Orders the operands of the compare so that they are listed from most
Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
 
   if (I.getType()->isVectorTy())
     if (Instruction *Res = foldVectorCmp(I, Builder))
       return Res;
 
   if (Instruction *Res = foldICmpInvariantGroup(I))
     return Res;
 
+  if (Instruction *Res = foldReductionIdiom(I, Builder, DL))
+    return Res;
+
   return Changed ? &I : nullptr;
 }
 
 /// Fold fcmp ([us]itofp x, cst) if possible.
 Instruction *InstCombinerImpl::foldFCmpIntToFPConst(FCmpInst &I,
                                                     Instruction *LHSI,
                                                     Constant *RHSC) {
   if (!isa<ConstantFP>(RHSC)) return nullptr;

llvm/test/Transforms/InstCombine/icmp-vec.ll

 ;
   %cmp.ne = icmp ne <vscale x 2 x i64> %x, zeroinitializer
   %sel = select <vscale x 2 x i1> %cmp.ne, <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> undef, i1 true, i32 0), <vscale x 2 x i1> undef, <vscale x 2 x i32> zeroinitializer), <vscale x 2 x i1> %falseval
   ret <vscale x 2 x i1> %sel
 }
 
 define i1 @eq_cast_eq-1(<2 x i4> %x, <2 x i4> %y) {
 ; CHECK-LABEL: @eq_cast_eq-1(
-; CHECK-NEXT:    [[IC:%.*]] = icmp ne <2 x i4> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <2 x i1> [[IC]] to i2
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i2 [[TMP1]], 0
+; CHECK-NEXT:    [[X_SCALAR:%.*]] = bitcast <2 x i4> [[X:%.*]] to i8
+; CHECK-NEXT:    [[Y_SCALAR:%.*]] = bitcast <2 x i4> [[Y:%.*]] to i8
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i8 [[X_SCALAR]], [[Y_SCALAR]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %ic = icmp eq <2 x i4> %x, %y
   %b = bitcast <2 x i1> %ic to i2
   %r = icmp eq i2 %b, -1
   ret i1 %r
 }
 
 define i1 @ne_cast_eq-1(<3 x i7> %x, <3 x i7> %y) {
 ; CHECK-LABEL: @ne_cast_eq-1(
 ; CHECK-NEXT:    [[IC:%.*]] = icmp eq <3 x i7> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <3 x i1> [[IC]] to i3
 ; CHECK-NEXT:    [[R:%.*]] = icmp eq i3 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %ic = icmp ne <3 x i7> %x, %y
   %b = bitcast <3 x i1> %ic to i3
   %r = icmp eq i3 %b, -1
   ret i1 %r
 }
 
 define i1 @eq_cast_ne-1(<2 x i7> %x, <2 x i7> %y) {
 ; CHECK-LABEL: @eq_cast_ne-1(
 ; CHECK-NEXT:    [[IC:%.*]] = icmp ne <2 x i7> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <2 x i1> [[IC]] to i2
 ; CHECK-NEXT:    [[R:%.*]] = icmp ne i2 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %ic = icmp eq <2 x i7> %x, %y
   %b = bitcast <2 x i1> %ic to i2

llvm/test/Transforms/InstCombine/reduction-and-sext-zext-i1.ll

;
   %ext = extractelement <8 x i64> %zext, i32 0
   store i64 %ext, i64* @glob1, align 8
   ret i64 %res
 }
 
 define i1 @reduce_and_pointer_cast(i8* %arg, i8* %arg1) {
 ; CHECK-LABEL: @reduce_and_pointer_cast(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[PTR1:%.*]] = bitcast i8* [[ARG1:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[PTR2:%.*]] = bitcast i8* [[ARG:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[LHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR1]], align 8
-; CHECK-NEXT:    [[RHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR2]], align 8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <8 x i8> [[LHS]], [[RHS]]
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <8 x i1> [[CMP]] to i8
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[TMP0]], 0
-; CHECK-NEXT:    ret i1 [[TMP1]]
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[ARG1:%.*]] to i64*
+; CHECK-NEXT:    [[LHS1:%.*]] = load i64, i64* [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[ARG:%.*]] to i64*
+; CHECK-NEXT:    [[RHS2:%.*]] = load i64, i64* [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i64 [[LHS1]], [[RHS2]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
 bb:
   %ptr1 = bitcast i8* %arg1 to <8 x i8>*
   %ptr2 = bitcast i8* %arg to <8 x i8>*
   %lhs = load <8 x i8>, <8 x i8>* %ptr1
   %rhs = load <8 x i8>, <8 x i8>* %ptr2
   %cmp = icmp eq <8 x i8> %lhs, %rhs
   %all_eq = call i1 @llvm.vector.reduce.and.v8i32(<8 x i1> %cmp)
bb:
   %cmp = icmp eq <8 x i16> %lhs, %rhs
   %all_eq = call i1 @llvm.vector.reduce.and.v8i32(<8 x i1> %cmp)
   ret i1 %all_eq
 }
 
 define i1 @reduce_and_pointer_cast_ne(i8* %arg, i8* %arg1) {
 ; CHECK-LABEL: @reduce_and_pointer_cast_ne(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[PTR1:%.*]] = bitcast i8* [[ARG1:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[PTR2:%.*]] = bitcast i8* [[ARG:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[LHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR1]], align 8
-; CHECK-NEXT:    [[RHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR2]], align 8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <8 x i8> [[LHS]], [[RHS]]
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <8 x i1> [[CMP]] to i8
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i8 [[TMP0]], 0
-; CHECK-NEXT:    ret i1 [[TMP1]]
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[ARG1:%.*]] to i64*
+; CHECK-NEXT:    [[LHS1:%.*]] = load i64, i64* [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[ARG:%.*]] to i64*
+; CHECK-NEXT:    [[RHS2:%.*]] = load i64, i64* [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i64 [[LHS1]], [[RHS2]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
 bb:
   %ptr1 = bitcast i8* %arg1 to <8 x i8>*
   %ptr2 = bitcast i8* %arg to <8 x i8>*
   %lhs = load <8 x i8>, <8 x i8>* %ptr1
   %rhs = load <8 x i8>, <8 x i8>* %ptr2
   %cmp = icmp eq <8 x i8> %lhs, %rhs
   %all_eq = call i1 @llvm.vector.reduce.and.v8i32(<8 x i1> %cmp)

llvm/test/Transforms/InstCombine/reduction-or-sext-zext-i1.ll

;
   %ext = extractelement <8 x i64> %zext, i32 0
   store i64 %ext, i64* @glob1, align 8
   ret i64 %res
 }
 
 define i1 @reduce_or_pointer_cast(i8* %arg, i8* %arg1) {
 ; CHECK-LABEL: @reduce_or_pointer_cast(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[PTR1:%.*]] = bitcast i8* [[ARG1:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[PTR2:%.*]] = bitcast i8* [[ARG:%.*]] to <8 x i8>*
-; CHECK-NEXT:    [[LHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR1]], align 8
-; CHECK-NEXT:    [[RHS:%.*]] = load <8 x i8>, <8 x i8>* [[PTR2]], align 8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <8 x i8> [[LHS]], [[RHS]]
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <8 x i1> [[CMP]] to i8
-; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq i8 [[TMP0]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[ARG1:%.*]] to i64*
+; CHECK-NEXT:    [[LHS1:%.*]] = load i64, i64* [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[ARG:%.*]] to i64*
+; CHECK-NEXT:    [[RHS2:%.*]] = load i64, i64* [[TMP1]], align 8
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq i64 [[LHS1]], [[RHS2]]
 ; CHECK-NEXT:    ret i1 [[DOTNOT]]
 ;
 bb:
   %ptr1 = bitcast i8* %arg1 to <8 x i8>*
   %ptr2 = bitcast i8* %arg to <8 x i8>*
   %lhs = load <8 x i8>, <8 x i8>* %ptr1
   %rhs = load <8 x i8>, <8 x i8>* %ptr2
   %cmp = icmp ne <8 x i8> %lhs, %rhs