[VectorCombine] Add option to only run scalarization transforms.

[fhahn]

Summary

This patch adds a pass option to only run transforms that scalarize
vector operations and do not create new vector instructions.

When running VectorCombine early in the pipeline introducing new vector
operations can have negative effects, like blocking loop or SLP
vectorization. To avoid regressions, restrict the early VectorCombine
run (when using -enable-matrix) to only perform scalarization and not
introduce new vector operations.

This is done as option to the pass directly, which is then set when
adding the pass to the pipeline. This is done for the new pass manager
only.

[spatel]

What if we just did better in VectorCombine?

We'd need to chain a bunch of combines together on paper or just implement this:
https://llvm.org/PR52178
...to know if it gets us to the minimal set of shuffles in IR and/or codegen for the 'hadd' example, but it might be enough?

[fhahn]

In D111800#3064573, @spatel wrote:

What if we just did better in VectorCombine?

We'd need to chain a bunch of combines together on paper or just implement this:
https://llvm.org/PR52178
...to know if it gets us to the minimal set of shuffles in IR and/or codegen for the 'hadd' example, but it might be enough?

Yes, I think that would get us a bit further, especially on the ARM64 test case. For X86 the shuffles/add chains are a bit more difficult to tackle: it converts the 4 scalar adds to 4 vector adds which each process a single lane. I'm not sure if we will be able to cover this in VectorCombine.

[spatel]

In D111800#3064694, @fhahn wrote:

In D111800#3064573, @spatel wrote:

What if we just did better in VectorCombine?

We'd need to chain a bunch of combines together on paper or just implement this:
https://llvm.org/PR52178
...to know if it gets us to the minimal set of shuffles in IR and/or codegen for the 'hadd' example, but it might be enough?

Yes, I think that would get us a bit further, especially on the ARM64 test case. For X86 the shuffles/add chains are a bit more difficult to tackle: it converts the 4 scalar adds to 4 vector adds which each process a single lane. I'm not sure if we will be able to cover this in VectorCombine.

I just drafted a patch for PR51278, and it got the hadd example down to:

In D111800#3064694, @fhahn wrote:

In D111800#3064573, @spatel wrote:

What if we just did better in VectorCombine?

We'd need to chain a bunch of combines together on paper or just implement this:
https://llvm.org/PR52178
...to know if it gets us to the minimal set of shuffles in IR and/or codegen for the 'hadd' example, but it might be enough?

Yes, I think that would get us a bit further, especially on the ARM64 test case. For X86 the shuffles/add chains are a bit more difficult to tackle: it converts the 4 scalar adds to 4 vector adds which each process a single lane. I'm not sure if we will be able to cover this in VectorCombine.

I drafted a patch for PR52178 , and I see that we get the first pair folded, but we're stuck on the next pair:

define <4 x float> @reverse_hadd_v4f32(<4 x float> %a, <4 x float> %b) local_unnamed_addr #0 {
  %shift = shufflevector <4 x float> %a, <4 x float> poison, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
  %shift1 = shufflevector <4 x float> %a, <4 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>
  %1 = shufflevector <4 x float> %shift, <4 x float> %shift1, <4 x i32> <i32 undef, i32 undef, i32 6, i32 0>
  %2 = shufflevector <4 x float> %a, <4 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 2, i32 0>
  %3 = fadd <4 x float> %1, %2
  %shift2 = shufflevector <4 x float> %b, <4 x float> poison, <4 x i32> <i32 undef, i32 0, i32 undef, i32 undef>
  %4 = fadd <4 x float> %shift2, %b
  %5 = shufflevector <4 x float> %3, <4 x float> %4, <4 x i32> <i32 undef, i32 5, i32 2, i32 3>
  %shift3 = shufflevector <4 x float> %b, <4 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>
  %6 = fadd <4 x float> %shift3, %b
  %7 = shufflevector <4 x float> %5, <4 x float> %6, <4 x i32> <i32 6, i32 1, i32 2, i32 3>
  ret <4 x float> %7
}

[spatel]

I posted the proposed vector-combine change here:
D111901

But this patch LGTM anyway. We don't want vector-combine to make things harder for LV or SLP.
Those passes should have a better view of total costs, so it seems right to have vector-combine only trying its minor changes after the others have run.

[fhahn]

In D111800#3067330, @spatel wrote:

I drafted a patch for PR52178 , and I see that we get the first pair folded, but we're stuck on the next pair:

Thanks for putting up that patch! I'll take a closer look in the next few days. Is it possible that we need to add additional instructions to the worklist to catch the earlier cases?

In D111800#3067330, @spatel wrote:

I posted the proposed vector-combine change here:
D111901

But this patch LGTM anyway. We don't want vector-combine to make things harder for LV or SLP.
Those passes should have a better view of total costs, so it seems right to have vector-combine only trying its minor changes after the others have run.

Sounds good, happy to iterate on improvements to vector-combine separately :)

llvm/include/llvm/Transforms/Vectorize/VectorCombine.h

 #ifndef LLVM_TRANSFORMS_VECTORIZE_VECTORCOMBINE_H
 #define LLVM_TRANSFORMS_VECTORIZE_VECTORCOMBINE_H
 
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// Optimize scalar/vector interactions in IR using target cost models.
-struct VectorCombinePass : public PassInfoMixin<VectorCombinePass> {
+class VectorCombinePass : public PassInfoMixin<VectorCombinePass> {
+  /// If true only perform scalarization combines and do not introduce new
+  /// vector operations.
+  bool ScalarizationOnly;
+
 public:
+  VectorCombinePass(bool ScalarizationOnly = false)
+      : ScalarizationOnly(ScalarizationOnly) {}
+
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
 };
-
 }
 #endif // LLVM_TRANSFORMS_VECTORIZE_VECTORCOMBINE_H

llvm/lib/Passes/PassBuilderPipelines.cpp

FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM2),
                                               /*UseBlockFrequencyInfo=*/false));
 
   // Delete small array after loop unroll.
   FPM.addPass(SROA());
 
   // The matrix extension can introduce large vector operations early, which can
   // benefit from running vector-combine early on.
   if (EnableMatrix)
-    FPM.addPass(VectorCombinePass());
+    FPM.addPass(VectorCombinePass(/*ScalarizationOnly=*/true));
 
   // Eliminate redundancies.
   FPM.addPass(MergedLoadStoreMotionPass());
   if (RunNewGVN)
     FPM.addPass(NewGVNPass());
   else
     FPM.addPass(GVN());
 

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

static cl::opt<unsigned> MaxInstrsToScan(
     cl::desc("Max number of instructions to scan for vector combining."));
 
 static const unsigned InvalidIndex = std::numeric_limits<unsigned>::max();
 
 namespace {
 class VectorCombine {
 public:
   VectorCombine(Function &F, const TargetTransformInfo &TTI,
-                const DominatorTree &DT, AAResults &AA, AssumptionCache &AC)
-      : F(F), Builder(F.getContext()), TTI(TTI), DT(DT), AA(AA), AC(AC) {}
+                const DominatorTree &DT, AAResults &AA, AssumptionCache &AC,
+                bool ScalarizationOnly)
+      : F(F), Builder(F.getContext()), TTI(TTI), DT(DT), AA(AA), AC(AC),
+        ScalarizationOnly(ScalarizationOnly) {}
 
   bool run();
 
 private:
   Function &F;
   IRBuilder<> Builder;
   const TargetTransformInfo &TTI;
   const DominatorTree &DT;
   AAResults &AA;
   AssumptionCache &AC;
+
+  /// If true only perform scalarization combines and do not introduce new
+  /// vector operations.
+  bool ScalarizationOnly;
+
   InstructionWorklist Worklist;
 
   bool vectorizeLoadInsert(Instruction &I);
   ExtractElementInst *getShuffleExtract(ExtractElementInst *Ext0,
                                         ExtractElementInst *Ext1,
                                         unsigned PreferredExtractIndex) const;
   bool isExtractExtractCheap(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                              const Instruction &I,
bool VectorCombine::run() {
 
   // Don't attempt vectorization if the target does not support vectors.
   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 |= foldBitcastShuf(I);
-    MadeChange |= scalarizeBinopOrCmp(I);
       MadeChange |= foldExtractedCmps(I);
+    }
+    MadeChange |= scalarizeBinopOrCmp(I);
     MadeChange |= scalarizeLoadExtract(I);
     MadeChange |= foldSingleElementStore(I);
   };
   for (BasicBlock &BB : F) {
     // Ignore unreachable basic blocks.
     if (!DT.isReachableFromEntry(&BB))
       continue;
     // Use early increment range so that we can erase instructions in loop.
public:
 
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
     auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
-    VectorCombine Combiner(F, TTI, DT, AA, AC);
+    VectorCombine Combiner(F, TTI, DT, AA, AC, false);
     return Combiner.run();
   }
 };
 } // namespace
 
 char VectorCombineLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(VectorCombineLegacyPass, "vector-combine",
                       "Optimize scalar/vector ops", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(VectorCombineLegacyPass, "vector-combine",
                     "Optimize scalar/vector ops", false, false)
 Pass *llvm::createVectorCombinePass() {
   return new VectorCombineLegacyPass();
 }
 
 PreservedAnalyses VectorCombinePass::run(Function &F,
                                          FunctionAnalysisManager &FAM) {
   auto &AC = FAM.getResult<AssumptionAnalysis>(F);
   TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
   DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
   AAResults &AA = FAM.getResult<AAManager>(F);
-  VectorCombine Combiner(F, TTI, DT, AA, AC);
+  VectorCombine Combiner(F, TTI, DT, AA, AC, ScalarizationOnly);
   if (!Combiner.run())
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }

llvm/test/Transforms/PhaseOrdering/AArch64/matrix-extract-insert.ll

 
 ; Function Attrs: argmemonly nofree nosync nounwind willreturn
 declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #1
 
 ; Function Attrs: nounwind ssp uwtable mustprogress
 
 define <4 x float> @reverse_hadd_v4f32(<4 x float> %a, <4 x float> %b) {
 ; CHECK-LABEL: @reverse_hadd_v4f32(
-; CHECK-NEXT:    [[SHIFT:%.*]] = shufflevector <4 x float> [[A:%.*]], <4 x float> poison, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
-; CHECK-NEXT:    [[TMP1:%.*]] = fadd <4 x float> [[SHIFT]], [[A]]
-; CHECK-NEXT:    [[SHIFT1:%.*]] = shufflevector <4 x float> [[A]], <4 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>
-; CHECK-NEXT:    [[TMP2:%.*]] = fadd <4 x float> [[SHIFT1]], [[A]]
-; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x float> [[TMP1]], <4 x float> [[TMP2]], <4 x i32> <i32 undef, i32 undef, i32 6, i32 0>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shufflevector <4 x float> [[B:%.*]], <4 x float> poison, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
-; CHECK-NEXT:    [[TMP4:%.*]] = fadd <4 x float> [[SHIFT2]], [[B]]
-; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x float> [[TMP3]], <4 x float> [[TMP4]], <4 x i32> <i32 undef, i32 4, i32 2, i32 3>
-; CHECK-NEXT:    [[SHIFT3:%.*]] = shufflevector <4 x float> [[B]], <4 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>
-; CHECK-NEXT:    [[TMP6:%.*]] = fadd <4 x float> [[SHIFT3]], [[B]]
-; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <4 x float> [[TMP5]], <4 x float> [[TMP6]], <4 x i32> <i32 6, i32 1, i32 2, i32 3>
-; CHECK-NEXT:    ret <4 x float> [[TMP7]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x float> [[A:%.*]], <4 x float> undef, <2 x i32> <i32 2, i32 0>
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x float> [[A]], <4 x float> undef, <2 x i32> <i32 3, i32 1>
+; CHECK-NEXT:    [[TMP3:%.*]] = fadd <2 x float> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <2 x float> [[TMP3]], <2 x float> poison, <4 x i32> <i32 undef, i32 undef, i32 0, i32 1>
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x float> [[B:%.*]], <4 x float> undef, <2 x i32> <i32 2, i32 0>
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <4 x float> [[B]], <4 x float> undef, <2 x i32> <i32 3, i32 1>
+; CHECK-NEXT:    [[TMP7:%.*]] = fadd <2 x float> [[TMP5]], [[TMP6]]
+; CHECK-NEXT:    [[TMP8:%.*]] = shufflevector <2 x float> [[TMP7]], <2 x float> poison, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+; CHECK-NEXT:    [[TMP9:%.*]] = shufflevector <4 x float> [[TMP8]], <4 x float> [[TMP4]], <4 x i32> <i32 0, i32 1, i32 6, i32 7>
+; CHECK-NEXT:    ret <4 x float> [[TMP9]]
 ;
   %vecext = extractelement <4 x float> %a, i32 0
   %vecext1 = extractelement <4 x float> %a, i32 1
   %add = fadd float %vecext, %vecext1
   %vecinit = insertelement <4 x float> undef, float %add, i32 0
   %vecext2 = extractelement <4 x float> %a, i32 2
   %vecext3 = extractelement <4 x float> %a, i32 3
   %add4 = fadd float %vecext2, %vecext3