[INSTCOMBINE] Transform reduction(shuffle V, poison, unique_mask) to reduction(V).

[ABataev]

Summary

After SLP + LTO we may have have reduction(shuffle V, poison,
mask). This can be simplified to just reduction(V) if the mask is only
for single vector and just all elements from this vector are permuted,

without reusing, replacing with undefs and/or other values, etc.

[lebedev.ri]

Is this really true for FP reductions? Don't they need reassoc flag for this?

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2000–2001

Not very useful comment

2014

This doesn't deal with non-canonical shuffle variant properly (and the test is missing)

[ABataev]

In D105053#2845269, @lebedev.ri wrote:

Is this really true for FP reductions? Don't they need reassoc flag for this?

I assume you're right, will add a check.

[ABataev]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2014

What do you mean by non-canonical shuffle?

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2014

match(Arg, m_Shuffle(m_Undef(), m_Value(V), m_Mask(Mask)))

[ABataev]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2014

Ah, yes, will fix it, thanks!

[ABataev]

Address comments.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2013

We canonicalize undef/poison operand as the 2nd input vector and we're within InstCombine, so checking for that is unnecessary?

https://github.com/llvm/llvm-project/blob/b458bb8c04cd5ed025884d424f386a00c9c6857e/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp#L2427

The tests have coverage now to verify that unless I missed something in the earlier draft.

llvm/test/Transforms/InstCombine/reduction-shufflevector.ll

9

Identity shuffle should be reduced even without this patch.
It would be better to pre-commit the test file, so we can make sure that we are testing as expected.

159

It's independent of this patch, but we might want to mark this and similar tests with a TODO - if a reduction has a known undef/poison element in the input, the result can be simplified.

[ABataev]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2013

Ok, will check it and remove this extra check.

llvm/test/Transforms/InstCombine/reduction-shufflevector.ll

9

Yes, I believe so, just added a check to try all possible combinations.
Did not precommit the test initially to be sure that the patch/the test are correct. Will precommit it, if it is fine.

159

Yes, thought about it too, just forgot to add it, thanks!

[spatel]

llvm/test/Transforms/InstCombine/reduction-shufflevector.ll

9

Sure, please precommit - the transform can be confirmed by Alive2:
https://alive2.llvm.org/ce/z/CZH3eH

[ABataev]

Rebase + remove the check for non-canonical shuffle.

[spatel]

LGTM - see inline for a few minor points.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2010

Nit: we're not using FVTy, so it would be easier to read (less indented) if we just bail out on scalable vectors:

if (!isa<FixedVectorType>(Arg->getType())) break;

You could also combine the next set of conditions with this to early exit and remove another level of indent.

2015

Unusual spelling: Indices or Indexes?

llvm/test/Transforms/InstCombine/reduction-shufflevector.ll

90

This is correct (no FMF required) based on LangRef, but it would be good to add some extra flags on at least one FP test, so we verify that those are all propagated through the transform.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/FloatingPointMode.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/Loads.h"
if (RetTy == Builder.getInt1Ty())
         } else {
           assert(IID == Intrinsic::vector_reduce_or &&
                  "Expected or reduction.");
           Res = Builder.CreateIsNotNull(Res);
         }
         replaceInstUsesWith(CI, Res);
         return eraseInstFromFunction(CI);
       }
+    LLVM_FALLTHROUGH;
+  }
+  case Intrinsic::vector_reduce_add:
+  case Intrinsic::vector_reduce_mul:
+  case Intrinsic::vector_reduce_xor:
+  case Intrinsic::vector_reduce_umax:
+  case Intrinsic::vector_reduce_umin:
+  case Intrinsic::vector_reduce_smax:
+  case Intrinsic::vector_reduce_smin:
+  case Intrinsic::vector_reduce_fmax:
+  case Intrinsic::vector_reduce_fmin:
+  case Intrinsic::vector_reduce_fadd:
+  case Intrinsic::vector_reduce_fmul: {
+    bool CanBeReassociated = (IID != Intrinsic::vector_reduce_fadd &&
+                              IID != Intrinsic::vector_reduce_fmul) ||
+                             II->hasAllowReassoc();
+    const unsigned ArgIdx = (IID == Intrinsic::vector_reduce_fadd ||
+                             IID == Intrinsic::vector_reduce_fmul)
+                                ? 1
+                                : 0;
+    Value *Arg = II->getArgOperand(ArgIdx);
+    Value *V;
+    ArrayRef<int> Mask;
+    if (!isa<FixedVectorType>(Arg->getType()) || !CanBeReassociated ||
+        !match(Arg, m_Shuffle(m_Value(V), m_Undef(), m_Mask(Mask))) ||
+        !cast<ShuffleVectorInst>(Arg)->isSingleSource())
+      break;
+    int Sz = Mask.size();
+    SmallBitVector UsedIndices(Sz);
+    for (int Idx : Mask) {
+      if (Idx == UndefMaskElem || UsedIndices.test(Idx))
+        break;
+      UsedIndices.set(Idx);
+    }
+    // Can remove shuffle iff just shuffled elements, no repeats, undefs, or
+    // other changes.
+    if (UsedIndices.all()) {
+      replaceUse(II->getOperandUse(ArgIdx), V);
+      return nullptr;
+    }
     break;
   }
   default: {
     // Handle target specific intrinsics
     Optional<Instruction *> V = targetInstCombineIntrinsic(*II);
     if (V.hasValue())
       return V.getValue();
     break;

llvm/test/Transforms/InstCombine/reduction-shufflevector.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i32 @reduce_add(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_add(
 ; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %res = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_or(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_or(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.or.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.or.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> poison, <4 x i32> %x, <4 x i32> <i32 7, i32 6, i32 5, i32 4>
   %res = call i32 @llvm.vector.reduce.or.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_and(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_and(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 0, i32 2, i32 1, i32 3>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 0, i32 2, i32 1, i32 3>
   %res = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_xor(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_xor(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 1, i32 2, i32 3, i32 0>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.xor.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.xor.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> poison, <4 x i32> %x, <4 x i32> <i32 5, i32 6, i32 7, i32 4>
   %res = call i32 @llvm.vector.reduce.xor.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_umax(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_umax(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 2, i32 1, i32 3, i32 0>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.umax.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.umax.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 2, i32 1, i32 3, i32 0>
   %res = call i32 @llvm.vector.reduce.umax.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_umin(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_umin(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 2, i32 3, i32 0, i32 1>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.umin.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.umin.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 2, i32 3, i32 0, i32 1>
   %res = call i32 @llvm.vector.reduce.umin.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_smax(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_smax(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.smax.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.smax.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
   %res = call i32 @llvm.vector.reduce.smax.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_smin(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_smin(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
-; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.smin.v4i32(<4 x i32> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.smin.v4i32(<4 x i32> [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
   %res = call i32 @llvm.vector.reduce.smin.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define float @reduce_fmax(<4 x float> %x) {
 ; CHECK-LABEL: @reduce_fmax(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> poison, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
-; CHECK-NEXT:    [[RES:%.*]] = call float @llvm.vector.reduce.fmax.v4f32(<4 x float> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call nnan nsz float @llvm.vector.reduce.fmax.v4f32(<4 x float> [[X:%.*]])
 ; CHECK-NEXT:    ret float [[RES]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> poison, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
-  %res = call float @llvm.vector.reduce.fmax.v4f32(<4 x float> %shuf)
+  %res = call nsz nnan float @llvm.vector.reduce.fmax.v4f32(<4 x float> %shuf)
   ret float %res
 }
 
 define float @reduce_fmin(<4 x float> %x) {
 ; CHECK-LABEL: @reduce_fmin(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> poison, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
-; CHECK-NEXT:    [[RES:%.*]] = call float @llvm.vector.reduce.fmin.v4f32(<4 x float> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call float @llvm.vector.reduce.fmin.v4f32(<4 x float> [[X:%.*]])
 ; CHECK-NEXT:    ret float [[RES]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> poison, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
   %res = call float @llvm.vector.reduce.fmin.v4f32(<4 x float> %shuf)
   ret float %res
 }
 
 define float @reduce_fadd(float %a, <4 x float> %x) {
 ; CHECK-LABEL: @reduce_fadd(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> undef, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
-; CHECK-NEXT:    [[RES:%.*]] = call reassoc float @llvm.vector.reduce.fadd.v4f32(float [[A:%.*]], <4 x float> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call reassoc float @llvm.vector.reduce.fadd.v4f32(float [[A:%.*]], <4 x float> [[X:%.*]])
 ; CHECK-NEXT:    ret float [[RES]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> %x, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
   %res = call reassoc float @llvm.vector.reduce.fadd.v4f32(float %a, <4 x float> %shuf)
   ret float %res
 }
 
 define float @reduce_fmul(float %a, <4 x float> %x) {
 ; CHECK-LABEL: @reduce_fmul(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> poison, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
-; CHECK-NEXT:    [[RES:%.*]] = call reassoc float @llvm.vector.reduce.fmul.v4f32(float [[A:%.*]], <4 x float> [[SHUF]])
+; CHECK-NEXT:    [[RES:%.*]] = call reassoc float @llvm.vector.reduce.fmul.v4f32(float [[A:%.*]], <4 x float> [[X:%.*]])
 ; CHECK-NEXT:    ret float [[RES]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> zeroinitializer, <4 x i32> <i32 0, i32 3, i32 1, i32 2>
   %res = call reassoc float @llvm.vector.reduce.fmul.v4f32(float %a, <4 x float> %shuf)
   ret float %res
 }
 
 ; Failed cases

 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
 ; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> [[SHUF]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
   %res = call i32 @llvm.vector.reduce.and.v4i32(<4 x i32> %shuf)
   ret i32 %res
 }
 
 define i32 @reduce_xor_failed(<4 x i32> %x) {
 ; CHECK-LABEL: @reduce_xor_failed(
 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> poison, <4 x i32> <i32 1, i32 2, i32 3, i32 undef>
 ; CHECK-NEXT:    [[RES:%.*]] = call i32 @llvm.vector.reduce.xor.v4i32(<4 x i32> [[SHUF]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> poison, <4 x i32> <i32 1, i32 2, i32 3, i32 undef>
   %res = call i32 @llvm.vector.reduce.xor.v4i32(<4 x i32> %shuf)