Custom IR printer for SLP that prints the cost after each vectorized instruction

[vporpo]

Summary

NOTE: This is WIP and not ready for reaview.

If we use a custom printer we could annotate the instructions with their cost.
This could prove useful for keeping track of cost model changes.
We could use that in existing lit tests and get something like:

%6 = add <4 x i32> %shuffle, %4            Cost: -3
%7 = bitcast i32* %idxS0 to <4 x i32>*
store <4 x i32> %6, <4 x i32>* %7, align 8 Cost: -3
ret void

There are some issues though:

• The cost of inserts/extracts is not currently being mapped well to the generated instructions, so this needs more work
• The dump contains dead instructions that are not being removed by SLP.

[RKSimon]

I like this, although why not just put it as a trailing comment so we can include it usable IR?

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
+#include "llvm/Support/FormattedStream.h"
 #ifdef EXPENSIVE_CHECKS
 #include "llvm/IR/Verifier.h"
 #endif
+#include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/DOTGraphTraits.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/GraphWriter.h"
cl::opt<bool> RunSLPVectorization("vectorize-slp", cl::init(true), cl::Hidden,
                                   cl::desc("Run the SLP vectorization passes"));
 
 static cl::opt<int>
     SLPCostThreshold("slp-threshold", cl::init(0), cl::Hidden,
                      cl::desc("Only vectorize if you gain more than this "
                               "number "));
 
 static cl::opt<bool>
+    SLPDumpIRWithCosts("slp-dump-ir-with-costs", cl::init(false), cl::Hidden,
+                       cl::desc("Dump IR annotated with SLP tree costs"));
+
+static cl::opt<bool>
 ShouldVectorizeHor("slp-vectorize-hor", cl::init(true), cl::Hidden,
                    cl::desc("Attempt to vectorize horizontal reductions"));
 
 static cl::opt<bool> ShouldStartVectorizeHorAtStore(
     "slp-vectorize-hor-store", cl::init(false), cl::Hidden,
     cl::desc(
         "Attempt to vectorize horizontal reductions feeding into a store"));
 
#endif
   /// It's like Instruction::eraseFromParent() except that the actual deletion
   /// is delayed until BoUpSLP is destructed.
   void eraseInstruction(Instruction *I) {
     DeletedInstructions.insert(I);
   }
 
   ~BoUpSLP();
 
+#ifndef NDEBUG
+  llvm::Optional<InstructionCost>
+  getVectorizedInstrCost(const Instruction *I) const {
+    auto It = VectorizedInstrCostMap.find(I);
+    return It != VectorizedInstrCostMap.end()
+               ? llvm::Optional<InstructionCost>(It->second)
+               : None;
+  }
+#endif
+
 private:
   /// Check if the operands on the edges \p Edges of the \p UserTE allows
   /// reordering (i.e. the operands can be reordered because they have only one
   /// user and reordarable).
   /// \param ReorderableGathers List of all gather nodes that require reordering
   /// (e.g., gather of extractlements or partially vectorizable loads).
   /// \param GatherOps List of gather operand nodes for \p UserTE that require
   /// reordering, subset of \p NonVectorized.
struct TreeEntry {
     /// Thus we need a reference back to the container to translate the indices
     /// to entries.
     VecTreeTy &Container;
 
     /// The TreeEntry index containing the user of this entry.  We can actually
     /// have multiple users so the data structure is not truly a tree.
     SmallVector<EdgeInfo, 1> UserTreeIndices;
 
+    /// The cost gain (vector - scalar) if we decide to vectorize this entry.
+    InstructionCost Cost;
+
     /// The index of this treeEntry in VectorizableTree.
     int Idx = -1;
 
   private:
     /// The operands of each instruction in each lane Operands[op_index][lane].
     /// Note: This helps avoid the replication of the code that performs the
     /// reordering of operands during buildTree_rec() and vectorizeTree().
     SmallVector<ValueList, 2> Operands;
int findLaneForValue(Value *V) const {
       assert(FoundLane < Scalars.size() && "Couldn't find extract lane");
       if (!ReuseShuffleIndices.empty()) {
         FoundLane = std::distance(ReuseShuffleIndices.begin(),
                                   find(ReuseShuffleIndices, FoundLane));
       }
       return FoundLane;
     }
 
+    void setCost(InstructionCost Cost) { this->Cost = Cost; }
+
+    InstructionCost getCost() const { return Cost; }
+
 #ifndef NDEBUG
     /// Debug printer.
     LLVM_DUMP_METHOD void dump() const {
       dbgs() << Idx << ".\n";
       for (unsigned OpI = 0, OpE = Operands.size(); OpI != OpE; ++OpI) {
         dbgs() << "Operand " << OpI << ":\n";
         for (const Value *V : Operands[OpI])
           dbgs().indent(2) << *V << "\n";
       }
       dbgs() << "Scalars: \n";
       for (Value *V : Scalars)
         dbgs().indent(2) << *V << "\n";
       dbgs() << "State: ";
       switch (State) {
       case Vectorize:
         dbgs() << "Vectorize\n";
         break;
       case ScatterVectorize:
         dbgs() << "ScatterVectorize\n";
         break;
       case NeedToGather:
         dbgs() << "NeedToGather\n";
         break;
       }
+      dbgs() << "Cost: "
+             << (Cost.getValue() ? std::to_string(*Cost.getValue()) : "-")
+             << "\n";
       dbgs() << "MainOp: ";
       if (MainOp)
         dbgs() << *MainOp << "\n";
       else
         dbgs() << "NULL\n";
       dbgs() << "AltOp: ";
       if (AltOp)
         dbgs() << *AltOp << "\n";
#endif
   SmallDenseMap<Value*, TreeEntry *> ScalarToTreeEntry;
 
   /// Maps a value to the proposed vectorizable size.
   SmallDenseMap<Value *, unsigned> InstrElementSize;
 
   /// A list of scalars that we found that we need to keep as scalars.
   ValueSet MustGather;
 
+#ifndef NDEBUG
+  /// Helper map for holding the cost of vectorized instructions. Used only in
+  /// lit tests with -slp-dump-ir-with-costs.
+  SmallDenseMap<const Instruction *, InstructionCost> VectorizedInstrCostMap;
+#endif
+
   /// This POD struct describes one external user in the vectorized tree.
   struct ExternalUser {
     ExternalUser(Value *S, llvm::User *U, int L)
         : Scalar(S), User(U), Lane(L) {}
 
     // Which scalar in our function.
     Value *Scalar;
 
LLVM_DEBUG(dbgs() << "SLP: Calculating cost for tree of size "
                     << VectorizableTree.size() << ".\n");
 
   unsigned BundleWidth = VectorizableTree[0]->Scalars.size();
 
   for (unsigned I = 0, E = VectorizableTree.size(); I < E; ++I) {
     TreeEntry &TE = *VectorizableTree[I];
 
     InstructionCost C = getEntryCost(&TE, VectorizedVals);
+    TE.setCost(C);
     Cost += C;
     LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
                       << " for bundle that starts with " << *TE.Scalars[0]
                       << ".\n"
                       << "SLP: Current total cost = " << Cost << "\n");
   }
 
   SmallPtrSet<Value *, 16> ExtractCostCalculated;
Value *BoUpSLP::createBuildVector(ArrayRef<Value *> VL) {
 
   ShuffleInstructionBuilder ShuffleBuilder(Builder, VF, GatherShuffleSeq,
                                            CSEBlocks);
   Value *Vec = gather(VL);
   if (!ReuseShuffleIndicies.empty()) {
     ShuffleBuilder.addMask(ReuseShuffleIndicies);
     Vec = ShuffleBuilder.finalize(Vec);
   }
+#ifndef NDEBUG
+  if (SLPDumpIRWithCosts)
+    if (Instruction *I = dyn_cast<Instruction>(Vec))
+      VectorizedInstrCostMap[I] = getGatherCost(VL);
+#endif
   return Vec;
 }
 
 Value *BoUpSLP::vectorizeTree(TreeEntry *E) {
   IRBuilder<>::InsertPointGuard Guard(Builder);
 
   if (E->VectorizedValue) {
     LLVM_DEBUG(dbgs() << "SLP: Diamond merged for " << *E->Scalars[0] << ".\n");
switch (ShuffleOrOp) {
     default:
     llvm_unreachable("unknown inst");
   }
   return nullptr;
 }
 
 Value *BoUpSLP::vectorizeTree() {
   ExtraValueToDebugLocsMap ExternallyUsedValues;
-  return vectorizeTree(ExternallyUsedValues);
+  Value *V = vectorizeTree(ExternallyUsedValues);
+#ifndef NDEBUG
+  if (SLPDumpIRWithCosts) {
+    // Collect the cost of all vectorized instructions.
+    for (const auto &TE : VectorizableTree) {
+      if (Value *VectorizedValue = TE->VectorizedValue)
+        if (Instruction *I = dyn_cast<Instruction>(VectorizedValue))
+          VectorizedInstrCostMap[I] = TE->getCost();
+    }
+  }
+#endif
+  return V;
 }
 
 Value *
 BoUpSLP::vectorizeTree(ExtraValueToDebugLocsMap &ExternallyUsedValues) {
   // All blocks must be scheduled before any instructions are inserted.
   for (auto &BSIter : BlocksSchedules) {
     scheduleBlock(BSIter.second.get());
   }
if (!GEPs.empty()) {
       Changed |= vectorizeGEPIndices(BB, R);
     }
   }
 
   if (Changed) {
     R.optimizeGatherSequence();
     LLVM_DEBUG(dbgs() << "SLP: vectorized \"" << F.getName() << "\"\n");
   }
+
+  if (SLPDumpIRWithCosts) {
+    /// Helper class for annotating vector instructions with their entry cost.
+    class SLPAssemblyAnnotationWriter : public AssemblyAnnotationWriter {
+      BoUpSLP *R;
+
+    public:
+      SLPAssemblyAnnotationWriter(BoUpSLP *R) : R(R) {}
+      void printInfoComment(const Value &V,
+                            formatted_raw_ostream &OS) override {
+        if (const Instruction *I = dyn_cast<Instruction>(&V))
+          if (Optional<InstructionCost> CostOpt = R->getVectorizedInstrCost(I))
+            OS << " Cost:" << CostOpt->getValue();
+      }
+    };
+    SLPAssemblyAnnotationWriter SLPAAW(&R);
+    F.print(outs(), &SLPAAW, true, true);
+  }
   return Changed;
 }
 
 bool SLPVectorizerPass::vectorizeStoreChain(ArrayRef<Value *> Chain, BoUpSLP &R,
                                             unsigned Idx) {
   LLVM_DEBUG(dbgs() << "SLP: Analyzing a store chain of length " << Chain.size()
                     << "\n");
   const unsigned Sz = R.getVectorElementSize(Chain[0]);

llvm/test/Transforms/SLPVectorizer/X86/broadcast.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -slp-vectorizer -S -mtriple=x86_64-unknown-linux -mcpu=corei7-avx -slp-threshold=-999 < %s | FileCheck %s
+; RUN: opt -slp-vectorizer -S -mtriple=x86_64-unknown-linux -mcpu=corei7-avx -slp-threshold=-999 -disable-output -slp-dump-ir-with-costs < %s | FileCheck %s
 
 
 ; S[0] = %v1 + %v2
 ; S[1] = %v2 + %v1
 ; S[2] = %v2 + %v1
 ; S[3] = %v1 + %v2
 ;
 ; We broadcast %v1 and %v2
 ;
 
 define void @bcast_vals(i64 *%A, i64 *%B, i64 *%S) {
 ; CHECK-LABEL: @bcast_vals(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A0:%.*]] = load i64, i64* [[A:%.*]], align 8
 ; CHECK-NEXT:    [[B0:%.*]] = load i64, i64* [[B:%.*]], align 8
 ; CHECK-NEXT:    [[V1:%.*]] = sub i64 [[A0]], 1
 ; CHECK-NEXT:    [[V2:%.*]] = sub i64 [[B0]], 1
 ; CHECK-NEXT:    [[IDXS0:%.*]] = getelementptr inbounds i64, i64* [[S:%.*]], i64 0
+; CHECK-NEXT:    [[IDXS1:%.*]] = getelementptr inbounds i64, i64* [[S]], i64 1
+; CHECK-NEXT:    [[IDXS2:%.*]] = getelementptr inbounds i64, i64* [[S]], i64 2
+; CHECK-NEXT:    [[IDXS3:%.*]] = getelementptr inbounds i64, i64* [[S]], i64 3
+; CHECK-NEXT:    [[ADD3:%.*]] = add i64 [[V1]], [[V2]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i64> poison, i64 [[V1]], i32 0
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i64> [[TMP0]], <4 x i64> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i64> [[TMP0]], <4 x i64> poison, <4 x i32> zeroinitializer Cost:5
 ; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i64> poison, i64 [[V2]], i32 0
-; CHECK-NEXT:    [[SHUFFLE1:%.*]] = shufflevector <4 x i64> [[TMP1]], <4 x i64> poison, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP2:%.*]] = add <4 x i64> [[SHUFFLE]], [[SHUFFLE1]]
+; CHECK-NEXT:    [[SHUFFLE1:%.*]] = shufflevector <4 x i64> [[TMP1]], <4 x i64> poison, <4 x i32> zeroinitializer Cost:5
+; CHECK-NEXT:    [[TMP2:%.*]] = add <4 x i64> [[SHUFFLE]], [[SHUFFLE1]] Cost:0
+; CHECK-NEXT:    [[ADD2:%.*]] = add i64 [[V2]], [[V1]]
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[V2]], [[V1]]
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[V1]], [[V2]]
+; CHECK-NEXT:    store i64 [[ADD3]], i64* [[IDXS3]], align 8
 ; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i64* [[IDXS0]] to <4 x i64>*
-; CHECK-NEXT:    store <4 x i64> [[TMP2]], <4 x i64>* [[TMP3]], align 8
+; CHECK-NEXT:    store <4 x i64> [[TMP2]], <4 x i64>* [[TMP3]], align 8 Cost:-2
+; CHECK-NEXT:    store i64 [[ADD2]], i64* [[IDXS2]], align 8
+; CHECK-NEXT:    store i64 [[ADD1]], i64* [[IDXS1]], align 8
+; CHECK-NEXT:    store i64 [[ADD0]], i64* [[IDXS0]], align 8
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %A0 = load i64, i64 *%A, align 8
   %B0 = load i64, i64 *%B, align 8
 
   %v1 = sub i64 %A0, 1
   %v2 = sub i64 %B0, 1

 
 ;
 define void @bcast_vals2(i16 *%A, i16 *%B, i16 *%C, i16 *%D, i16 *%E, i32 *%S) {
 ; CHECK-LABEL: @bcast_vals2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A0:%.*]] = load i16, i16* [[A:%.*]], align 8
 ; CHECK-NEXT:    [[V1:%.*]] = sext i16 [[A0]] to i32
 ; CHECK-NEXT:    [[IDXS0:%.*]] = getelementptr inbounds i32, i32* [[S:%.*]], i64 0
+; CHECK-NEXT:    [[IDXS1:%.*]] = getelementptr inbounds i32, i32* [[S]], i64 1
+; CHECK-NEXT:    [[IDXS2:%.*]] = getelementptr inbounds i32, i32* [[S]], i64 2
+; CHECK-NEXT:    [[IDXS3:%.*]] = getelementptr inbounds i32, i32* [[S]], i64 3
 ; CHECK-NEXT:    [[B0:%.*]] = load i16, i16* [[B:%.*]], align 8
 ; CHECK-NEXT:    [[C0:%.*]] = load i16, i16* [[C:%.*]], align 8
 ; CHECK-NEXT:    [[D0:%.*]] = load i16, i16* [[D:%.*]], align 8
 ; CHECK-NEXT:    [[E0:%.*]] = load i16, i16* [[E:%.*]], align 8
+; CHECK-NEXT:    [[V4:%.*]] = sext i16 [[D0]] to i32
 ; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i16> poison, i16 [[B0]], i32 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i16> [[TMP0]], i16 [[C0]], i32 1
 ; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x i16> [[TMP1]], i16 [[E0]], i32 2
-; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <4 x i16> [[TMP2]], i16 [[D0]], i32 3
-; CHECK-NEXT:    [[TMP4:%.*]] = sext <4 x i16> [[TMP3]] to <4 x i32>
+; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <4 x i16> [[TMP2]], i16 [[D0]], i32 3 Cost:4
+; CHECK-NEXT:    [[TMP4:%.*]] = sext <4 x i16> [[TMP3]] to <4 x i32> Cost:1
+; CHECK-NEXT:    [[V5:%.*]] = sext i16 [[E0]] to i32
+; CHECK-NEXT:    [[V3:%.*]] = sext i16 [[C0]] to i32
+; CHECK-NEXT:    [[V2:%.*]] = sext i16 [[B0]] to i32
+; CHECK-NEXT:    [[ADD3:%.*]] = add i32 [[V1]], [[V4]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = insertelement <4 x i32> poison, i32 [[V1]], i32 0
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i32> [[TMP5]], <4 x i32> poison, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP6:%.*]] = add <4 x i32> [[SHUFFLE]], [[TMP4]]
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i32> [[TMP5]], <4 x i32> poison, <4 x i32> zeroinitializer Cost:2
+; CHECK-NEXT:    [[TMP6:%.*]] = add <4 x i32> [[SHUFFLE]], [[TMP4]] Cost:-3
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[V5]], [[V1]]
+; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[V3]], [[V1]]
+; CHECK-NEXT:    [[ADD0:%.*]] = add i32 [[V1]], [[V2]]
+; CHECK-NEXT:    store i32 [[ADD3]], i32* [[IDXS3]], align 8
 ; CHECK-NEXT:    [[TMP7:%.*]] = bitcast i32* [[IDXS0]] to <4 x i32>*
-; CHECK-NEXT:    store <4 x i32> [[TMP6]], <4 x i32>* [[TMP7]], align 8
+; CHECK-NEXT:    store <4 x i32> [[TMP6]], <4 x i32>* [[TMP7]], align 8 Cost:-3
+; CHECK-NEXT:    store i32 [[ADD2]], i32* [[IDXS2]], align 8
+; CHECK-NEXT:    store i32 [[ADD1]], i32* [[IDXS1]], align 8
+; CHECK-NEXT:    store i32 [[ADD0]], i32* [[IDXS0]], align 8
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %A0 = load i16, i16 *%A, align 8
   %B0 = load i16, i16 *%B, align 8
   %C0 = load i16, i16 *%C, align 8
   %D0 = load i16, i16 *%D, align 8
   %E0 = load i16, i16 *%E, align 8