[VectorCombine] convert scalar fneg with insert/extract to vector fneg

[spatel]

Summary

insertelt DestVec, (fneg (extractelt SrcVec, Index)), Index --> shuffle DestVec, (fneg SrcVec), Mask

This is a specialized form of what could be a more general fold for a binop. It's also possible that fneg is overlooked by SLP in this kind of insert/extract pattern since it's a unary op.

This shows up in the motivating example from #issue 58139, but it won't solve it (that probably requires some x86-specific backend changes). There are also some small enhancements (see TODO comments) that can be done as follow-up patches.

[tschuett]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

563

Tiniest nit: You could Mask.reserve(VecTy->getNumElements());

[RKSimon]

LGTM

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

563

If you wanted to you could go even further (assuming we know that Index < VecTy->getNumElements())

SmallVector<int> Mask(VecTy->getNumElements());
std::iota(Mask.begin(), Mask.end(), 0);
Mask[Index] = Index + e;

llvm/test/Transforms/VectorCombine/X86/extract-fneg-insert.ll

7

Are you going to look at this? If not please can you raise a bug.

[spatel]

llvm/test/Transforms/VectorCombine/X86/extract-fneg-insert.ll

7

I'm not sure what the proper fix will be, so let's track it either way:
https://github.com/llvm/llvm-project/issues/58261

[spatel]

Patch updated:
Replace shuffle mask loop with pre-allocated vector and std::iota().

[RKSimon]

LGTM

[uabelho]

Hello,

The following starts crashing with this patch:

opt -passes='vector-combine' bbi-74776.ll -S -o /dev/null

bbi-74776.ll233 BDownload

[spatel]

In D135278#3862060, @uabelho wrote:

Hello,

The following starts crashing with this patch:

opt -passes='vector-combine' bbi-74776.ll -S -o /dev/null

bbi-74776.ll233 BDownload

Thanks! Taking a look now - probably need to restrict the pattern match to true fneg.

[spatel]

In D135278#3862108, @spatel wrote:

In D135278#3862060, @uabelho wrote:

Hello,

The following starts crashing with this patch:

opt -passes='vector-combine' bbi-74776.ll -S -o /dev/null

bbi-74776.ll233 BDownload

Thanks! Taking a look now - probably need to restrict the pattern match to true fneg.

On 2nd thought, we can adjust the capture to handle this kind of pattern:
8d76fbb5f065

[uabelho]

Thanks!

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

bool isExtractExtractCheap(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                              const Instruction &I,
                              ExtractElementInst *&ConvertToShuffle,
                              unsigned PreferredExtractIndex);
   void foldExtExtCmp(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                      Instruction &I);
   void foldExtExtBinop(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                        Instruction &I);
   bool foldExtractExtract(Instruction &I);
+  bool foldInsExtFNeg(Instruction &I);
   bool foldBitcastShuf(Instruction &I);
   bool scalarizeBinopOrCmp(Instruction &I);
   bool foldExtractedCmps(Instruction &I);
   bool foldSingleElementStore(Instruction &I);
   bool scalarizeLoadExtract(Instruction &I);
   bool foldShuffleOfBinops(Instruction &I);
   bool foldShuffleFromReductions(Instruction &I);
   bool foldSelectShuffle(Instruction &I, bool FromReduction = false);
bool VectorCombine::foldExtractExtract(Instruction &I) {
   else
     foldExtExtBinop(Ext0, Ext1, I);
 
   Worklist.push(Ext0);
   Worklist.push(Ext1);
   return true;
 }
 
+/// Try to replace an extract + scalar fneg + insert with a vector fneg +
+/// shuffle.
+bool VectorCombine::foldInsExtFNeg(Instruction &I) {
+  auto *VecTy = dyn_cast<FixedVectorType>(I.getType());
+  if (!VecTy)
+    return false;
+
+  // Match an insert (op (extract)) pattern.
+  Value *DestVec;
+  uint64_t Index;
+  Instruction *FNeg;
+  if (!match(&I, m_InsertElt(m_Value(DestVec), m_OneUse(m_Instruction(FNeg)),
+                             m_ConstantInt(Index))))
+    return false;
+
+  Value *SrcVec;
+  if (!match(FNeg, m_FNeg(m_ExtractElt(m_Value(SrcVec), m_SpecificInt(Index)))))
+    return false;
+
+  if (SrcVec->getType() != VecTy)
+    return false;
+
+  // We are inserting the negated element into the same lane that we extracted
+  // from. This is equivalent to a select-shuffle that chooses all but the
+  // negated element from the destination vector.
+  SmallVector<int> Mask;
+  for (unsigned i = 0, e = VecTy->getNumElements(); i != e; ++i) {
+    if (i == Index)
+      Mask.push_back(Index + e);
+    else
+      Mask.push_back(i);
+  }
+
+  Type *ScalarTy = VecTy->getScalarType();
+  InstructionCost OldCost =
+      TTI.getArithmeticInstrCost(Instruction::FNeg, ScalarTy) +
+      TTI.getVectorInstrCost(I, VecTy, Index);
+
+  // If the extract has one use, it will be eliminated, so count it in the
+  // original cost. If it has more than one use, ignore the cost because it will
+  // be the same before/after.
+  Instruction *Extract = cast<Instruction>(FNeg->getOperand(0));
+  if (Extract->hasOneUse())
+    OldCost += TTI.getVectorInstrCost(*Extract, VecTy, Index);
+
+  InstructionCost NewCost =
+      TTI.getArithmeticInstrCost(Instruction::FNeg, VecTy) +
+      TTI.getShuffleCost(TargetTransformInfo::SK_Select, VecTy, Mask);
+
+  if (NewCost > OldCost)
+    return false;
+
+  // insertelt DestVec, (fneg (extractelt SrcVec, Index)), Index -->
+  // shuffle DestVec, (fneg SrcVec), Mask
+  Value *VecFNeg = Builder.CreateFNegFMF(SrcVec, FNeg);
+  Value *Shuf = Builder.CreateShuffleVector(DestVec, VecFNeg, Mask);
+  replaceValue(I, *Shuf);
+  return true;
+}
+
 /// If this is a bitcast of a shuffle, try to bitcast the source vector to the
 /// destination type followed by shuffle. This can enable further transforms by
 /// moving bitcasts or shuffles together.
 bool VectorCombine::foldBitcastShuf(Instruction &I) {
   Value *V;
   ArrayRef<int> Mask;
   if (!match(&I, m_BitCast(
                      m_OneUse(m_Shuffle(m_Value(V), m_Undef(), m_Mask(Mask))))))
if (!TTI.getNumberOfRegisters(TTI.getRegisterClassForType(/*Vector*/ true)))
     return false;
 
   bool MadeChange = false;
   auto FoldInst = [this, &MadeChange](Instruction &I) {
     Builder.SetInsertPoint(&I);
     if (!ScalarizationOnly) {
       MadeChange |= vectorizeLoadInsert(I);
       MadeChange |= foldExtractExtract(I);
+      MadeChange |= foldInsExtFNeg(I);
       MadeChange |= foldBitcastShuf(I);
       MadeChange |= foldExtractedCmps(I);
       MadeChange |= foldShuffleOfBinops(I);
       MadeChange |= foldShuffleFromReductions(I);
       MadeChange |= foldSelectShuffle(I);
     }
     MadeChange |= scalarizeBinopOrCmp(I);
     MadeChange |= scalarizeLoadExtract(I);

llvm/test/Transforms/VectorCombine/X86/extract-fneg-insert.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=vector-combine -S -mtriple=x86_64-- -mattr=SSE2 | FileCheck %s --check-prefixes=CHECK,SSE
 ; RUN: opt < %s -passes=vector-combine -S -mtriple=x86_64-- -mattr=AVX2 | FileCheck %s --check-prefixes=CHECK,AVX
 
 declare void @use(float)
 
+; TODO: The insert is costed as free, so creating a shuffle appears to be a loss.
+
 define <4 x float> @ext0_v4f32(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @ext0_v4f32(
 ; CHECK-NEXT:    [[E:%.*]] = extractelement <4 x float> [[X:%.*]], i32 0
 ; CHECK-NEXT:    [[N:%.*]] = fneg float [[E]]
 ; CHECK-NEXT:    [[R:%.*]] = insertelement <4 x float> [[Y:%.*]], float [[N]], i32 0
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %e = extractelement <4 x float> %x, i32 0
   %n = fneg float %e
   %r = insertelement <4 x float> %y, float %n, i32 0
   ret <4 x float> %r
 }
 
 ; Eliminating extract/insert is profitable.
 
 define <4 x float> @ext2_v4f32(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @ext2_v4f32(
-; CHECK-NEXT:    [[E:%.*]] = extractelement <4 x float> [[X:%.*]], i32 2
-; CHECK-NEXT:    [[N:%.*]] = fneg float [[E]]
-; CHECK-NEXT:    [[R:%.*]] = insertelement <4 x float> [[Y:%.*]], float [[N]], i32 2
+; CHECK-NEXT:    [[TMP1:%.*]] = fneg <4 x float> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x float> [[Y:%.*]], <4 x float> [[TMP1]], <4 x i32> <i32 0, i32 1, i32 6, i32 3>
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %e = extractelement <4 x float> %x, i32 2
   %n = fneg float %e
   %r = insertelement <4 x float> %y, float %n, i32 2
   ret <4 x float> %r
 }
 
 ; Eliminating extract/insert is still profitable. Flags propagate.
 
 define <2 x double> @ext1_v2f64(<2 x double> %x, <2 x double> %y) {
 ; CHECK-LABEL: @ext1_v2f64(
-; CHECK-NEXT:    [[E:%.*]] = extractelement <2 x double> [[X:%.*]], i32 1
-; CHECK-NEXT:    [[N:%.*]] = fneg nsz double [[E]]
-; CHECK-NEXT:    [[R:%.*]] = insertelement <2 x double> [[Y:%.*]], double [[N]], i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = fneg nsz <2 x double> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <2 x double> [[Y:%.*]], <2 x double> [[TMP1]], <2 x i32> <i32 0, i32 3>
 ; CHECK-NEXT:    ret <2 x double> [[R]]
 ;
   %e = extractelement <2 x double> %x, i32 1
   %n = fneg nsz double %e
   %r = insertelement <2 x double> %y, double %n, i32 1
   ret <2 x double> %r
 }
 
+; The vector fneg would cost twice as much as the scalar op with SSE,
+; so we don't transform there (the shuffle would also be more expensive).
+
 define <8 x float> @ext7_v8f32(<8 x float> %x, <8 x float> %y) {
-; CHECK-LABEL: @ext7_v8f32(
-; CHECK-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 7
-; CHECK-NEXT:    [[N:%.*]] = fneg float [[E]]
-; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x float> [[Y:%.*]], float [[N]], i32 7
-; CHECK-NEXT:    ret <8 x float> [[R]]
+; SSE-LABEL: @ext7_v8f32(
+; SSE-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 7
+; SSE-NEXT:    [[N:%.*]] = fneg float [[E]]
+; SSE-NEXT:    [[R:%.*]] = insertelement <8 x float> [[Y:%.*]], float [[N]], i32 7
+; SSE-NEXT:    ret <8 x float> [[R]]
+;
+; AVX-LABEL: @ext7_v8f32(
+; AVX-NEXT:    [[TMP1:%.*]] = fneg <8 x float> [[X:%.*]]
+; AVX-NEXT:    [[R:%.*]] = shufflevector <8 x float> [[Y:%.*]], <8 x float> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 15>
+; AVX-NEXT:    ret <8 x float> [[R]]
 ;
   %e = extractelement <8 x float> %x, i32 7
   %n = fneg float %e
   %r = insertelement <8 x float> %y, float %n, i32 7
   ret <8 x float> %r
 }
 
+; Same as above with an extra use of the extracted element.
+
 define <8 x float> @ext7_v8f32_use1(<8 x float> %x, <8 x float> %y) {
-; CHECK-LABEL: @ext7_v8f32_use1(
-; CHECK-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 5
-; CHECK-NEXT:    call void @use(float [[E]])
-; CHECK-NEXT:    [[N:%.*]] = fneg float [[E]]
-; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x float> [[Y:%.*]], float [[N]], i32 5
-; CHECK-NEXT:    ret <8 x float> [[R]]
+; SSE-LABEL: @ext7_v8f32_use1(
+; SSE-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 5
+; SSE-NEXT:    call void @use(float [[E]])
+; SSE-NEXT:    [[N:%.*]] = fneg float [[E]]
+; SSE-NEXT:    [[R:%.*]] = insertelement <8 x float> [[Y:%.*]], float [[N]], i32 5
+; SSE-NEXT:    ret <8 x float> [[R]]
+;
+; AVX-LABEL: @ext7_v8f32_use1(
+; AVX-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 5
+; AVX-NEXT:    call void @use(float [[E]])
+; AVX-NEXT:    [[TMP1:%.*]] = fneg <8 x float> [[X]]
+; AVX-NEXT:    [[R:%.*]] = shufflevector <8 x float> [[Y:%.*]], <8 x float> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 13, i32 6, i32 7>
+; AVX-NEXT:    ret <8 x float> [[R]]
 ;
   %e = extractelement <8 x float> %x, i32 5
   call void @use(float %e)
   %n = fneg float %e
   %r = insertelement <8 x float> %y, float %n, i32 5
   ret <8 x float> %r
 }
 
+; Negative test - the transform is likely not profitable if the fneg has another use.
+
 define <8 x float> @ext7_v8f32_use2(<8 x float> %x, <8 x float> %y) {
 ; CHECK-LABEL: @ext7_v8f32_use2(
 ; CHECK-NEXT:    [[E:%.*]] = extractelement <8 x float> [[X:%.*]], i32 3
 ; CHECK-NEXT:    [[N:%.*]] = fneg float [[E]]
 ; CHECK-NEXT:    call void @use(float [[N]])
 ; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x float> [[Y:%.*]], float [[N]], i32 3
 ; CHECK-NEXT:    ret <8 x float> [[R]]
 ;
   %e = extractelement <8 x float> %x, i32 3
   %n = fneg float %e
   call void @use(float %n)
   %r = insertelement <8 x float> %y, float %n, i32 3
   ret <8 x float> %r
 }
 
+; Negative test - can't convert variable index to a shuffle.
+
 define <2 x double> @ext_index_var_v2f64(<2 x double> %x, <2 x double> %y, i32 %index) {
 ; CHECK-LABEL: @ext_index_var_v2f64(
 ; CHECK-NEXT:    [[E:%.*]] = extractelement <2 x double> [[X:%.*]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    [[N:%.*]] = fneg nsz double [[E]]
 ; CHECK-NEXT:    [[R:%.*]] = insertelement <2 x double> [[Y:%.*]], double [[N]], i32 [[INDEX]]
 ; CHECK-NEXT:    ret <2 x double> [[R]]
 ;
   %e = extractelement <2 x double> %x, i32 %index
   %n = fneg nsz double %e
   %r = insertelement <2 x double> %y, double %n, i32 %index
   ret <2 x double> %r
 }
 
+; Negative test - require same extract/insert index for simple shuffle.
+; TODO: We could handle this by adjusting the cost calculation.
+
 define <2 x double> @ext1_v2f64_ins0(<2 x double> %x, <2 x double> %y) {
 ; CHECK-LABEL: @ext1_v2f64_ins0(
 ; CHECK-NEXT:    [[E:%.*]] = extractelement <2 x double> [[X:%.*]], i32 1
 ; CHECK-NEXT:    [[N:%.*]] = fneg nsz double [[E]]
 ; CHECK-NEXT:    [[R:%.*]] = insertelement <2 x double> [[Y:%.*]], double [[N]], i32 0
 ; CHECK-NEXT:    ret <2 x double> [[R]]
 ;
   %e = extractelement <2 x double> %x, i32 1
   %n = fneg nsz double %e
   %r = insertelement <2 x double> %y, double %n, i32 0
   ret <2 x double> %r
 }