[TTI] Add hook for contextual cast estimates

[mssimpso]

Summary

This change adds a new hook for estimating the cost of cast instructions within
a given context. The hook differs from the existing context-free hook in that
clients can provide information about the cast's known operand or user.

The change is motivated by sign extensions, which on AArch64 can often be
folded into other instructions and are free. For example, a vector extract
followed by a sign extension can be performed by smov. And a vector subtraction
whose operands are both vector sign extensions can be performed by ssubl. And
so on.

While context is generally not a desirable feature to include in the cost
model, the fact that these extensions are often free may make them special
enough to warrant the added complexity.

The default implementation of the hook uses the context-free estimate. For
AArch64, I've added logic to cover the extract-extend pattern and the
lengthening and widening operations I mentioned above. The intent is to quickly
and simply identify the cases in which the cast is free, and to fall back on
the default implementation otherwise.

For an initial client, I've modified the cast estimates in the SLP vectorizer
to provide the additional context when known and thought to be useful.

[jmolloy]

include/llvm/Analysis/TargetTransformInfo.h

475

I feel that this structure should contain EITHER an anonymous opcode (as it currently does) OR a Value*.

I think it is not worthwhile to have the user reencode their use-def graph for the purposes of passing as context. What we need to be able to say is "I'm going to insert $OPCODE and it's going to use these values". Some of the values may be yet to be inserted - so these could be anonymous too. But encoding the entire use-def graph to an arbitrary depth sounds like wasted effort. The code below for AArch64 seems to look quite deep at the graph too.

502

I'd prefer if an optional Ctx argument were added to existing methods.

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

468

If you expect Dst, you must use cast<> above instead of dyn_cast_or_null. Similarly in all the other place you use dyn_cast_or_null, I don't think a nullptr argument is *actually* allowed.

[mssimpso]

Addressed James's comments

Thanks again for the review James! I took a shot at updating the patch and also added some responses inline.

[mssimpso]

include/llvm/Analysis/TargetTransformInfo.h

475

I agree this seems overly complex. However, I'm having difficulty coming up with a better alternative.

For the scalar cost estimates (where we're getting an estimate for an instruction that already exists), I definitely see the advantage of just looking at the existing Values. We wouldn't have to re-encode the use-def information like you mention. But I think this gets trickier for the vector estimates. For example, considering the extract-extend case (where we sign extend the value we extract from the vector), we would need to create an unlinked instruction for the extract if I understand correctly. But what would it's operand be? We haven't decided to vectorize anything yet. And if we eventually chose not to vectorize, we would have to go back and delete the unlinked instructions we create. I'm definitely open to suggestions here, and any more thoughts you might have would be very helpful!

502

I'm fine doing this. I've added an optional Ctx argument to the existing getCastInstrCost hook like you suggested. The drawback of course is that we have to modify the other targets (and we might cause some minor annoyance for the out-of-tree target maintainers).

[mssimpso]

Rebased.

[hfinkel]

include/llvm/Analysis/TargetTransformInfo.h

475

We definitely need this kind of functionality, but I share James's uncomfortableness with forcing the user to reencode some arbitrary portion of the use/def graph. Part of the problem is that the user of the interface has no idea how much of the graph would be useful to encode. Plus, the things you're reencoding into are essentially instructions that aren't instructions. While we generally have a policy against creating temporary instructions, getting around that by reinventing instructions seems unsatisfying.

What if we defined an abstract graph-traversal class that could be used by a TTI implementation to walk a (mutated) view of the existing SSA use/def graph.

struct InstrContextView {
  struct InstrContext {
     Type *Ty;
     unsigned Opcode;
     unsigned NumOperands;
  };

  // Call this to get the initial instruction context (the starting point of the view on the use/def chain)
  virtual InstrContext *getInstrContext() = 0;
  virtual InstrContext *getOperandInstrContext(const InstrContext *, unsigned OpNum) = 0;
  virtual unsigned getNumInstrContextUsers(const InstrContext *) = 0;
  virtual InstrContext *getUserInstrContext(const InstrContext *, unsigned UserNum) = 0;
};

Then the TTI interfaces take some derived class of InstrContextView, which holds the necessary state to provide whatever view is possible given the caller's knowledge of the potential IR. In this way, the creation of the InstrContext objects can be completely lazy.

The downside is that, because the traversal functions need to be virtual, walking large portions of the graph this was is not feasible. Luckily, I think we never want to do that.

[mssimpso]

Hal,

Thanks very much for the suggestions, and I apologize for the delayed response. I agree the initial approach is unsatisfying. There is definitely a balance to be struck here between providing context to the cost model and having to unnecessarily re-encode information. The question of how much context to provide is also problematic. Although not a high priority for me at the moment, I do plan to evaluate your idea soon. So thanks again for the feedback!

[mssimpso]

Rebasing for evaluation only. The patch still needs to be updated to address reviewer comments.

[mssimpso]

I'm abandoning this change in favor of D18523. But I hope we can eventually revisit the idea of contextualizing the cost model estimates if needed.

include/llvm/Analysis/TargetTransformInfo.h

int getArithmeticInstrCost(
       OperandValueProperties Opd2PropInfo = OP_None) const;
 
   /// \return The cost of a shuffle instruction of kind Kind and of type Tp.
   /// The index and subtype parameters are used by the subvector insertion and
   /// extraction shuffle kinds.
   int getShuffleCost(ShuffleKind Kind, Type *Tp, int Index = 0,
                      Type *SubTp = nullptr) const;
 
+  /// A struct to hold information for contextual cost estimates.
+  struct InstrContext {
+    InstrContext() : InstrTy(nullptr), Opcode(0), User(nullptr) {
+      Operands = makeArrayRef((InstrContext *)nullptr, 0);
+    }
+    InstrContext(Type *Ty, unsigned Opc = 0, InstrContext *Usr = nullptr)
+        : InstrTy(Ty), Opcode(Opc), User(Usr) {
+      Operands = makeArrayRef((InstrContext *)nullptr, 0);
+    }
+    InstrContext(Type *Ty, unsigned Opc, ArrayRef<InstrContext> Ops,
+                 InstrContext *Usr = nullptr)
+        : InstrTy(Ty), Opcode(Opc), Operands(Ops), User(Usr) {}
+
+    /// The type of the instruction.
+    Type *InstrTy;
+
+    /// The opcode of the instruction.
+    unsigned Opcode;
+
+    /// A list of InstrContexts describing the operands of the instruction.
+    ArrayRef<InstrContext> Operands;
+
+    /// An InstrContext describing a single user of the instruction.
+    InstrContext *User;
+  };
+
   /// \return The expected cost of cast instructions, such as bitcast, trunc,
   /// zext, etc.
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) const;
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       InstrContext *Ctx = nullptr) const;
 
   /// \return The expected cost of control-flow related instructions such as
   /// Phi, Ret, Br.
   int getCFInstrCost(unsigned Opcode) const;
 
   /// \returns The expected cost of compare and select instructions.
   int getCmpSelInstrCost(unsigned Opcode, Type *ValTy,
                          Type *CondTy = nullptr) const;
 
   /// \return The expected cost of vector Insert and Extract.
public:
 
   /// \return The cost of Gather or Scatter operation
   /// \p Opcode - is a type of memory access Load or Store
   /// \p DataTy - a vector type of the data to be loaded or stored
   /// \p Ptr - pointer [or vector of pointers] - address[es] in memory
   /// \p VariableMask - true when the memory access is predicated with a mask
   ///                   that is not a compile-time constant
   /// \p Alignment - alignment of single element
   int getGatherScatterOpCost(unsigned Opcode, Type *DataTy, Value *Ptr,
                              bool VariableMask, unsigned Alignment) const;
 
   /// \return The cost of the interleaved memory operation.
   /// \p Opcode is the memory operation code
   /// \p VecTy is the vector type of the interleaved access.
   /// \p Factor is the interleave factor
   /// \p Indices is the indices for interleaved load members (as interleaved
   ///    load allows gaps)
public:
   virtual unsigned getMaxInterleaveFactor(unsigned VF) = 0;
   virtual unsigned
   getArithmeticInstrCost(unsigned Opcode, Type *Ty, OperandValueKind Opd1Info,
                          OperandValueKind Opd2Info,
                          OperandValueProperties Opd1PropInfo,
                          OperandValueProperties Opd2PropInfo) = 0;
   virtual int getShuffleCost(ShuffleKind Kind, Type *Tp, int Index,
                              Type *SubTp) = 0;
-  virtual int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) = 0;
+  virtual int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                               InstrContext *Ctx) = 0;
   virtual int getCFInstrCost(unsigned Opcode) = 0;
   virtual int getCmpSelInstrCost(unsigned Opcode, Type *ValTy,
                                  Type *CondTy) = 0;
   virtual int getVectorInstrCost(unsigned Opcode, Type *Val,
                                  unsigned Index) = 0;
   virtual int getMemoryOpCost(unsigned Opcode, Type *Src, unsigned Alignment,
                               unsigned AddressSpace) = 0;
   virtual int getMaskedMemoryOpCost(unsigned Opcode, Type *Src,
getArithmeticInstrCost(unsigned Opcode, Type *Ty, OperandValueKind Opd1Info,
                          OperandValueProperties Opd2PropInfo) override {
     return Impl.getArithmeticInstrCost(Opcode, Ty, Opd1Info, Opd2Info,
                                        Opd1PropInfo, Opd2PropInfo);
   }
   int getShuffleCost(ShuffleKind Kind, Type *Tp, int Index,
                      Type *SubTp) override {
     return Impl.getShuffleCost(Kind, Tp, Index, SubTp);
   }
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) override {
-    return Impl.getCastInstrCost(Opcode, Dst, Src);
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       InstrContext *Ctx) override {
+    return Impl.getCastInstrCost(Opcode, Dst, Src, Ctx);
   }
   int getCFInstrCost(unsigned Opcode) override {
     return Impl.getCFInstrCost(Opcode);
   }
   int getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy) override {
     return Impl.getCmpSelInstrCost(Opcode, ValTy, CondTy);
   }
   int getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index) override {

include/llvm/Analysis/TargetTransformInfoImpl.h

unsigned getArithmeticInstrCost(unsigned Opcode, Type *Ty,
     return 1;
   }
 
   unsigned getShuffleCost(TTI::ShuffleKind Kind, Type *Ty, int Index,
                           Type *SubTp) {
     return 1;
   }
 
-  unsigned getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) { return 1; }
+  unsigned getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                            TTI::InstrContext *Ctx) {
+    return 1;
+  }
 
   unsigned getCFInstrCost(unsigned Opcode) { return 1; }
 
   unsigned getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy) {
     return 1;
   }
 
   unsigned getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index) {

include/llvm/CodeGen/BasicTTIImpl.h

public:
   unsigned getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index,
                           Type *SubTp) {
     if (Kind == TTI::SK_Alternate) {
       return getAltShuffleOverhead(Tp);
     }
     return 1;
   }
 
-  unsigned getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) {
+  unsigned getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                            TTI::InstrContext *Ctx) {
     const TargetLoweringBase *TLI = getTLI();
     int ISD = TLI->InstructionOpcodeToISD(Opcode);
     assert(ISD && "Invalid opcode");
     std::pair<unsigned, MVT> SrcLT = TLI->getTypeLegalizationCost(DL, Src);
     std::pair<unsigned, MVT> DstLT = TLI->getTypeLegalizationCost(DL, Dst);
 
     // Check for NOOP conversions.
     if (SrcLT.first == DstLT.first &&
if (Dst->isVectorTy() && Src->isVectorTy()) {
           return SrcLT.first * 1;
       }
 
       // If we are converting vectors and the operation is illegal, or
       // if the vectors are legalized to different types, estimate the
       // scalarization costs.
       unsigned Num = Dst->getVectorNumElements();
       unsigned Cost = static_cast<T *>(this)->getCastInstrCost(
-          Opcode, Dst->getScalarType(), Src->getScalarType());
+          Opcode, Dst->getScalarType(), Src->getScalarType(), Ctx);
 
       // Return the cost of multiple scalar invocation plus the cost of
       // inserting and extracting the values.
       return getScalarizationOverhead(Dst, true, true) + Num * Cost;
     }
 
     // We already handled vector-to-vector and scalar-to-scalar conversions.
     // This

lib/Analysis/TargetTransformInfo.cpp

 
 int TargetTransformInfo::getShuffleCost(ShuffleKind Kind, Type *Ty, int Index,
                                         Type *SubTp) const {
   int Cost = TTIImpl->getShuffleCost(Kind, Ty, Index, SubTp);
   assert(Cost >= 0 && "TTI should not produce negative costs!");
   return Cost;
 }
 
-int TargetTransformInfo::getCastInstrCost(unsigned Opcode, Type *Dst,
-                                          Type *Src) const {
-  int Cost = TTIImpl->getCastInstrCost(Opcode, Dst, Src);
+int TargetTransformInfo::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                                          InstrContext *Ctx) const {
+  int Cost = TTIImpl->getCastInstrCost(Opcode, Dst, Src, Ctx);
   assert(Cost >= 0 && "TTI should not produce negative costs!");
   return Cost;
 }
 
 int TargetTransformInfo::getCFInstrCost(unsigned Opcode) const {
   int Cost = TTIImpl->getCFInstrCost(Opcode);
   assert(Cost >= 0 && "TTI should not produce negative costs!");
   return Cost;

lib/Target/AArch64/AArch64TargetTransformInfo.h

if (Vector) {
         return 128;
       return 0;
     }
     return 64;
   }
 
   unsigned getMaxInterleaveFactor(unsigned VF);
 
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src);
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       TTI::InstrContext *Ctx);
+
+  int getCastInstrCostInCtx(unsigned Opcode, Type *Dst, Type *Src,
+                            TTI::InstrContext *Ctx);
 
   int getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index);
 
   int getArithmeticInstrCost(
       unsigned Opcode, Type *Ty,
       TTI::OperandValueKind Opd1Info = TTI::OK_AnyValue,
       TTI::OperandValueKind Opd2Info = TTI::OK_AnyValue,
       TTI::OperandValueProperties Opd1PropInfo = TTI::OP_None,

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

 AArch64TTIImpl::getPopcntSupport(unsigned TyWidth) {
   assert(isPowerOf2_32(TyWidth) && "Ty width must be power of 2");
   if (TyWidth == 32 || TyWidth == 64)
     return TTI::PSK_FastHardware;
   // TODO: AArch64TargetLowering::LowerCTPOP() supports 128bit popcount.
   return TTI::PSK_Software;
 }
 
-int AArch64TTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) {
+int AArch64TTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                                     TTI::InstrContext *Ctx) {
+
+  // Try and use the contextual cost estimate. Note that the contextual version
+  // calls the implementation here if nothing useful can be done, so we can
+  // directly return the cost.
+  if (Ctx)
+    return getCastInstrCostInCtx(Opcode, Dst, Src, Ctx);
+
   int ISD = TLI->InstructionOpcodeToISD(Opcode);
   assert(ISD && "Invalid opcode");
 
   EVT SrcTy = TLI->getValueType(DL, Src);
   EVT DstTy = TLI->getValueType(DL, Dst);
 
   if (!SrcTy.isSimple() || !DstTy.isSimple())
-    return BaseT::getCastInstrCost(Opcode, Dst, Src);
+    return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 
   static const TypeConversionCostTblEntry
   ConversionTbl[] = {
     { ISD::TRUNCATE, MVT::v4i16, MVT::v4i32,  1 },
     { ISD::TRUNCATE, MVT::v4i32, MVT::v4i64,  0 },
     { ISD::TRUNCATE, MVT::v8i8,  MVT::v8i32,  3 },
     { ISD::TRUNCATE, MVT::v16i8, MVT::v16i32, 6 },
 
ConversionTbl[] = {
     { ISD::FP_TO_UINT, MVT::v2i8,  MVT::v2f64, 2 },
   };
 
   if (const auto *Entry = ConvertCostTableLookup(ConversionTbl, ISD,
                                                  DstTy.getSimpleVT(),
                                                  SrcTy.getSimpleVT()))
     return Entry->Cost;
 
-  return BaseT::getCastInstrCost(Opcode, Dst, Src);
+  return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
+}
+
+int AArch64TTIImpl::getCastInstrCostInCtx(unsigned Opcode, Type *Dst, Type *Src,
+                                          TTI::InstrContext *Ctx) {
+
+  assert(!Ctx->Operands.empty() && "Invalid context");
+
+  // Contextual cost estimate. We can base the estimate on information about
+  // the value being cast.
+  switch (Ctx->Operands[0].Opcode) {
+  default:
+    break;
+
+  // If the cast value has been extracted from a vector, we can usually perform
+  // sign and zero extensions automatically with smov and umov. We check for
+  // the free cases below, and fall back to the context-free estimate
+  // otherwise.
+  case Instruction::ExtractElement:
+
+    // We only contextualize sext and zext casts.
+    if (Opcode != Instruction::SExt && Opcode != Instruction::ZExt)
+      break;
+
+    // See if we've been given enough context to do anything useful. We need to
+    // know the first operand of the extractelement.
+    auto &Operands = Ctx->Operands[0].Operands;
+    if (Operands.empty())
+      break;
+
+    // Make sure the first operand of the extractelement is a vector type. This
+    // must be the case if we have a valid context.
+    assert(isa<VectorType>(Operands[0].InstrTy) && "Invalid context");
+    auto *Vec = cast<VectorType>(Operands[0].InstrTy);
+
+    // Because the cast is either zext or sext, we have to be working with all
+    // integer types. Again, this must be the case if we have a valid context.
+    assert(isa<IntegerType>(Dst) && Src == Vec->getElementType() &&
+           "Invalid context");
+
+    // Now we legalize the types. If the resulting type is still a vector and
+    // the destination type is legal, we may get the extension for free.
+    auto VecLT = TLI->getTypeLegalizationCost(DL, Vec);
+    auto DstVT = TLI->getValueType(DL, Dst);
+    if (!VecLT.second.isVector() || !TLI->isTypeLegal(DstVT))
+      break;
+
+    // Ensure we have a valid extract-extend combination. The source and
+    // element types should be the same, and the destination type should be
+    // larger than the element type.
+    auto SrcVT = TLI->getValueType(DL, Src);
+    auto ElmVT = VecLT.second.getVectorElementType();
+    if (SrcVT != ElmVT || DstVT.getSizeInBits() < ElmVT.getSizeInBits())
+      break;
+
+    switch (Opcode) {
+    default:
+      llvm_unreachable("Ctx->Opcode should be either SExt or ZExt");
+
+    // For sign extensions, we only need a smov, which performs the extension
+    // automatically.
+    case Instruction::SExt:
+      return 0;
+
+    // For zero extensions, the extend is performed automatically by a umov
+    // unless the destination type is i64 and the element type is i8 or i16.
+    case Instruction::ZExt:
+      if (DstVT.getSizeInBits() != 64u || ElmVT.getSizeInBits() == 32u)
+        return 0;
+    }
+  }
+
+  // At this point, we've looked at the operand of the cast if interesting to
+  // do so. If we don't yet know that the cast will be free, we can try to look
+  // at the user of the cast. If there is no information about what the user of
+  // the cast looks like, there is nothing else to do.
+  if (!Ctx->User)
+    return getCastInstrCost(Opcode, Dst, Src, nullptr);
+
+  // Otherwise, we can can base the estimate on information about the user.
+  switch (Ctx->User->Opcode) {
+  default:
+    break;
+
+  // AArch64 can perform widening and narrowing vector arithmetic. If the cast
+  // value is a vector extend that is used by an arithmetic operation, it may
+  // happen automatically. We check for the free cases below, and fall back to
+  // the context-free estimate otherwise.
+  case Instruction::Add:
+  case Instruction::Sub:
+
+    // We only contextualize sext and zext casts.
+    if (Ctx->Opcode != Instruction::SExt && Ctx->Opcode != Instruction::ZExt)
+      break;
+
+    // We have to look at the operands of widening and narrowing instructions.
+    // If we weren't given this information, there's nothing to do.
+    if (Ctx->User->Operands.size() != 2)
+      break;
+
+    // Type check the extend and it's user to ensure we have a valid context.
+    // Because we have an extend, it's source and destination types have to be
+    // integers or vectors of integers.
+    auto &Operands = Ctx->User->Operands;
+    assert(Src && Src->isIntOrIntVectorTy() && Dst &&
+           Dst->isIntOrIntVectorTy() && Ctx->User->InstrTy == Dst &&
+           Operands[0].InstrTy == Dst && Operands[1].InstrTy == Dst &&
+           "Invalid context");
+
+    // Check if this operation can be matched to a lengthening or widening
+    // vector operation. If not, give up.
+    auto IsWidening = Operands[1].Opcode == Ctx->Opcode;
+    auto IsLengthening = IsWidening && Operands[0].Opcode == Ctx->Opcode;
+    if (!IsWidening && !IsLengthening)
+      break;
+
+    // Now we legalize the types. If the resulting types are not vectors, this
+    // won't match the lengthening and widening operations.
+    auto SrcVT = TLI->getTypeLegalizationCost(DL, Src).second;
+    auto DstVT = TLI->getTypeLegalizationCost(DL, Dst).second;
+    if (!SrcVT.isVector() || !DstVT.isVector())
+      break;
+
+    // If we have legal vectors and the destination element size is twice the
+    // source element size, the extend can be performed automatically by the
+    // lengthening or widening instruction.
+    if (DstVT.getScalarSizeInBits() == 2 * SrcVT.getScalarSizeInBits())
+      return 0;
+  }
+
+  // If we can't do anything useful, fall back to the default implementation.
+  return getCastInstrCost(Opcode, Dst, Src, nullptr);
 }
 
 int AArch64TTIImpl::getVectorInstrCost(unsigned Opcode, Type *Val,
                                        unsigned Index) {
   assert(Val->isVectorTy() && "This must be a vector type");
 
   if (Index != -1U) {
     // Legalize the type.

 int AArch64TTIImpl::getArithmeticInstrCost(
     unsigned Opcode, Type *Ty, TTI::OperandValueKind Opd1Info,
     TTI::OperandValueKind Opd2Info, TTI::OperandValueProperties Opd1PropInfo,
     TTI::OperandValueProperties Opd2PropInfo) {
   // Legalize the type.
   std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Ty);
 
   int ISD = TLI->InstructionOpcodeToISD(Opcode);
 
   if (ISD == ISD::SDIV &&
       Opd2Info == TargetTransformInfo::OK_UniformConstantValue &&
       Opd2PropInfo == TargetTransformInfo::OP_PowerOf2) {
     // On AArch64, scalar signed division by constants power-of-two are
     // normally expanded to the sequence ADD + CMP + SELECT + SRA.
     // The OperandValue properties many not be same as that of previous
     // operation; conservatively assume OP_None.
     int Cost = getArithmeticInstrCost(Instruction::Add, Ty, Opd1Info, Opd2Info,

lib/Target/ARM/ARMTargetTransformInfo.h

unsigned getMaxInterleaveFactor(unsigned VF) {
     // These are out of order CPUs:
     if (ST->isCortexA15() || ST->isSwift())
       return 2;
     return 1;
   }
 
   int getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index, Type *SubTp);
 
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src);
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       TTI::InstrContext *Ctx);
 
   int getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy);
 
   int getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index);
 
   int getAddressComputationCost(Type *Val, bool IsComplex);
 
   int getFPOpCost(Type *Ty);

lib/Target/ARM/ARMTargetTransformInfo.cpp

int ARMTTIImpl::getIntImmCost(const APInt &Imm, Type *Ty) {
   if (SImmVal >= 0 && SImmVal < 256)
     return 1;
   if ((~ZImmVal < 256) || ARM_AM::isThumbImmShiftedVal(ZImmVal))
     return 2;
   // Load from constantpool.
   return 3;
 }
 
-int ARMTTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) {
+int ARMTTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                                 TTI::InstrContext *Ctx) {
   int ISD = TLI->InstructionOpcodeToISD(Opcode);
   assert(ISD && "Invalid opcode");
 
   // Single to/from double precision conversions.
   static const CostTblEntry NEONFltDblTbl[] = {
     // Vector fptrunc/fpext conversions.
     { ISD::FP_ROUND,   MVT::v2f64, 2 },
     { ISD::FP_EXTEND,  MVT::v2f32, 2 },
     { ISD::FP_EXTEND,  MVT::v4f32, 4 }
   };
 
   if (Src->isVectorTy() && ST->hasNEON() && (ISD == ISD::FP_ROUND ||
                                           ISD == ISD::FP_EXTEND)) {
     std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Src);
     if (const auto *Entry = CostTableLookup(NEONFltDblTbl, ISD, LT.second))
       return LT.first * Entry->Cost;
   }
 
   EVT SrcTy = TLI->getValueType(DL, Src);
   EVT DstTy = TLI->getValueType(DL, Dst);
 
   if (!SrcTy.isSimple() || !DstTy.isSimple())
-    return BaseT::getCastInstrCost(Opcode, Dst, Src);
+    return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 
   // Some arithmetic, load and store operations have specific instructions
   // to cast up/down their types automatically at no extra cost.
   // TODO: Get these tables to know at least what the related operations are.
   static const TypeConversionCostTblEntry NEONVectorConversionTbl[] = {
     { ISD::SIGN_EXTEND, MVT::v4i32, MVT::v4i16, 0 },
     { ISD::ZERO_EXTEND, MVT::v4i32, MVT::v4i16, 0 },
     { ISD::SIGN_EXTEND, MVT::v2i64, MVT::v2i32, 1 },
int ARMTTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
 
   if (SrcTy.isInteger()) {
     if (const auto *Entry = ConvertCostTableLookup(ARMIntegerConversionTbl, ISD,
                                                    DstTy.getSimpleVT(),
                                                    SrcTy.getSimpleVT()))
       return Entry->Cost;
   }
 
-  return BaseT::getCastInstrCost(Opcode, Dst, Src);
+  return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 }
 
 int ARMTTIImpl::getVectorInstrCost(unsigned Opcode, Type *ValTy,
                                    unsigned Index) {
   // Penalize inserting into an D-subregister. We end up with a three times
   // lower estimated throughput on swift.
   if (ST->isSwift() &&
       Opcode == Instruction::InsertElement &&

lib/Target/PowerPC/PPCTargetTransformInfo.h

public:
   unsigned getMaxInterleaveFactor(unsigned VF);
   int getArithmeticInstrCost(
       unsigned Opcode, Type *Ty,
       TTI::OperandValueKind Opd1Info = TTI::OK_AnyValue,
       TTI::OperandValueKind Opd2Info = TTI::OK_AnyValue,
       TTI::OperandValueProperties Opd1PropInfo = TTI::OP_None,
       TTI::OperandValueProperties Opd2PropInfo = TTI::OP_None);
   int getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index, Type *SubTp);
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src);
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       TTI::InstrContext *Ctx);
   int getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy);
   int getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index);
   int getMemoryOpCost(unsigned Opcode, Type *Src, unsigned Alignment,
                       unsigned AddressSpace);
   int getInterleavedMemoryOpCost(unsigned Opcode, Type *VecTy,
                                  unsigned Factor,
                                  ArrayRef<unsigned> Indices,
                                  unsigned Alignment,
                                  unsigned AddressSpace);
 
   /// @}
 };
 
 } // end namespace llvm
 
 #endif

lib/Target/PowerPC/PPCTargetTransformInfo.cpp

int PPCTTIImpl::getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index,
   // PPC, for both Altivec/VSX and QPX, support cheap arbitrary permutations
   // (at least in the sense that there need only be one non-loop-invariant
   // instruction). We need one such shuffle instruction for each actual
   // register (this is not true for arbitrary shuffles, but is true for the
   // structured types of shuffles covered by TTI::ShuffleKind).
   return LT.first;
 }
 
-int PPCTTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) {
+int PPCTTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                                 TTI::InstrContext *Ctx) {
   assert(TLI->InstructionOpcodeToISD(Opcode) && "Invalid opcode");
 
-  return BaseT::getCastInstrCost(Opcode, Dst, Src);
+  return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 }
 
 int PPCTTIImpl::getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy) {
   return BaseT::getCmpSelInstrCost(Opcode, ValTy, CondTy);
 }
 
 int PPCTTIImpl::getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index) {
   assert(Val->isVectorTy() && "This must be a vector type");

lib/Target/X86/X86TargetTransformInfo.h

public:
   unsigned getMaxInterleaveFactor(unsigned VF);
   int getArithmeticInstrCost(
       unsigned Opcode, Type *Ty,
       TTI::OperandValueKind Opd1Info = TTI::OK_AnyValue,
       TTI::OperandValueKind Opd2Info = TTI::OK_AnyValue,
       TTI::OperandValueProperties Opd1PropInfo = TTI::OP_None,
       TTI::OperandValueProperties Opd2PropInfo = TTI::OP_None);
   int getShuffleCost(TTI::ShuffleKind Kind, Type *Tp, int Index, Type *SubTp);
-  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src);
+  int getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                       TTI::InstrContext *Ctx);
   int getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy);
   int getVectorInstrCost(unsigned Opcode, Type *Val, unsigned Index);
   int getMemoryOpCost(unsigned Opcode, Type *Src, unsigned Alignment,
                       unsigned AddressSpace);
   int getMaskedMemoryOpCost(unsigned Opcode, Type *Src, unsigned Alignment,
                             unsigned AddressSpace);
   int getGatherScatterOpCost(unsigned Opcode, Type *DataTy, Value *Ptr,
                              bool VariableMask, unsigned Alignment);

lib/Target/X86/X86TargetTransformInfo.cpp

if (const auto *Entry = CostTableLookup(SSEAltShuffleTbl,
                                             ISD::VECTOR_SHUFFLE, LT.second))
       return LT.first * Entry->Cost;
     return BaseT::getShuffleCost(Kind, Tp, Index, SubTp);
   }
 
   return BaseT::getShuffleCost(Kind, Tp, Index, SubTp);
 }
 
-int X86TTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src) {
+int X86TTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
+                                 TTI::InstrContext *Ctx) {
   int ISD = TLI->InstructionOpcodeToISD(Opcode);
   assert(ISD && "Invalid opcode");
 
   // FIXME: Need a better design of the cost table to handle non-simple types of
   // potential massive combinations (elem_num x src_type x dst_type).
 
   static const TypeConversionCostTblEntry AVX512DQConversionTbl[] = {
     { ISD::UINT_TO_FP,  MVT::v2f64,  MVT::v2i64,  1 },
if (const auto *Entry = ConvertCostTableLookup(SSE2ConversionTbl, ISD,
       return LTSrc.first * Entry->Cost;
   }
 
   EVT SrcTy = TLI->getValueType(DL, Src);
   EVT DstTy = TLI->getValueType(DL, Dst);
 
   // The function getSimpleVT only handles simple value types.
   if (!SrcTy.isSimple() || !DstTy.isSimple())
-    return BaseT::getCastInstrCost(Opcode, Dst, Src);
+    return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 
   if (ST->hasDQI())
     if (const auto *Entry = ConvertCostTableLookup(AVX512DQConversionTbl, ISD,
                                                    DstTy.getSimpleVT(),
                                                    SrcTy.getSimpleVT()))
       return Entry->Cost;
 
   if (ST->hasAVX512())
int X86TTIImpl::getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
 
   if (ST->hasSSE2()) {
     if (const auto *Entry = ConvertCostTableLookup(SSE2ConversionTbl, ISD,
                                                    DstTy.getSimpleVT(),
                                                    SrcTy.getSimpleVT()))
       return Entry->Cost;
   }
 
-  return BaseT::getCastInstrCost(Opcode, Dst, Src);
+  return BaseT::getCastInstrCost(Opcode, Dst, Src, Ctx);
 }
 
 int X86TTIImpl::getCmpSelInstrCost(unsigned Opcode, Type *ValTy, Type *CondTy) {
   // Legalize the type.
   std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, ValTy);
 
   MVT MTy = LT.second;
 

lib/Transforms/Vectorize/SLPVectorizer.cpp

 /// Bottom Up SLP Vectorizer.
 class BoUpSLP {
 public:
   typedef SmallVector<Value *, 8> ValueList;
   typedef SmallVector<Instruction *, 16> InstrList;
   typedef SmallPtrSet<Value *, 16> ValueSet;
   typedef SmallVector<StoreInst *, 8> StoreList;
 
+  typedef TargetTransformInfo::InstrContext InstrCtx;
+
   BoUpSLP(Function *Func, ScalarEvolution *Se, TargetTransformInfo *Tti,
           TargetLibraryInfo *TLi, AliasAnalysis *Aa, LoopInfo *Li,
           DominatorTree *Dt, AssumptionCache *AC, DemandedBits *DB)
       : NumLoadsWantToKeepOrder(0), NumLoadsWantToChangeOrder(0), F(Func),
         SE(Se), TTI(Tti), TLI(TLi), AA(Aa), LI(Li), DT(Dt), AC(AC), DB(DB),
         Builder(Se->getContext()) {
     CodeMetrics::collectEphemeralValues(F, AC, EphValues);
   }
switch (Opcode) {
     case Instruction::FPExt:
     case Instruction::PtrToInt:
     case Instruction::IntToPtr:
     case Instruction::SIToFP:
     case Instruction::UIToFP:
     case Instruction::Trunc:
     case Instruction::FPTrunc:
     case Instruction::BitCast: {
-      Type *SrcTy = VL0->getOperand(0)->getType();
 
-      // Calculate the cost of this instruction.
-      int ScalarCost = VL.size() * TTI->getCastInstrCost(VL0->getOpcode(),
-                                                         VL0->getType(), SrcTy);
+      // Simplifying definitions to avoid confusion in the code below.
+      auto *SrcScalarTy = VL0->getOperand(0)->getType();
+      auto *SrcVectorTy = VectorType::get(SrcScalarTy, VL.size());
+      auto *DstScalarTy = VL0->getType();
+      auto *DstVectorTy = VecTy;
+      auto DstOpcode = VL0->getOpcode();
+
+      // The scalar and vector cost estimates. We will define these below.
+      int ScalarCost;
+      int VectorCost;
+
+      // If there is only one user of the scalar value, we can attempt to use
+      // the more precise, contextual cost estimates. If this is the case, we
+      // build up the context. Otherwise, we use the context-free estimates.
+      auto *TheUser = dyn_cast_or_null<Instruction>(*VL0->user_begin());
+      if (!TheUser || !TheUser->getType()->isIntegerTy() || !VL0->hasOneUse()) {
+        ScalarCost = TTI->getCastInstrCost(DstOpcode, DstScalarTy, SrcScalarTy);
+        VectorCost = TTI->getCastInstrCost(DstOpcode, DstVectorTy, SrcVectorTy);
+        ScalarCost *= VL.size();
+        return VectorCost - ScalarCost;
+      }
+
+      // We're going to use the contextual estimates. First define the contexts
+      // for both teh vector and scalar estimates.
+      InstrCtx ScalarCast, ScalarOper(SrcScalarTy), ScalarUser;
+      InstrCtx VectorCast, VectorOper(SrcVectorTy), VectorUser;
+      SmallVector<InstrCtx, 4> ScalarOperands;
+      SmallVector<InstrCtx, 4> VectorOperands;
+
+      // Next, gather information about the user.
+      auto UsrOpcode = TheUser->getOpcode();
+      auto *UsrScalarTy = TheUser->getType();
+      auto *UsrVectorTy = MinBWs.count(VL[0])
+                              ? DstVectorTy
+                              : VectorType::get(UsrScalarTy, VL.size());
+
+      // Finally, build the contexts. We first visit each operand of the user,
+      // and then connect the user to the cast.
+      for (Use &U : TheUser->operands()) {
+        auto *ScalarTy = U.get()->getType();
+        auto *VectorTy = VectorType::get(ScalarTy, VL.size());
+        auto *I = dyn_cast<Instruction>(U.get());
+        auto Opc = I ? I->getOpcode() : 0;
+        ScalarOperands.push_back(InstrCtx(ScalarTy, Opc, &ScalarUser));
+        if (MinBWs.count(VL[0]))
+          VectorOperands.push_back(InstrCtx(DstVectorTy, Opc, &VectorUser));
+        else
+          VectorOperands.push_back(InstrCtx(VectorTy, Opc, &VectorUser));
+      }
+      ScalarUser = InstrCtx(UsrScalarTy, UsrOpcode, ScalarOperands);
+      VectorUser = InstrCtx(UsrVectorTy, UsrOpcode, VectorOperands);
+      ScalarCast = InstrCtx(DstScalarTy, DstOpcode, ScalarOper, &ScalarUser);
+      VectorCast = InstrCtx(DstVectorTy, DstOpcode, VectorOper, &VectorUser);
+
+      // The contextual cost estimateis.
+      ScalarCost = TTI->getCastInstrCost(DstOpcode, DstScalarTy, SrcScalarTy,
+                                         &ScalarCast);
+      VectorCost = TTI->getCastInstrCost(DstOpcode, DstVectorTy, SrcVectorTy,
+                                         &VectorCast);
 
-      VectorType *SrcVecTy = VectorType::get(SrcTy, VL.size());
-      int VecCost = TTI->getCastInstrCost(VL0->getOpcode(), VecTy, SrcVecTy);
-      return VecCost - ScalarCost;
+      ScalarCost *= VL.size();
+      return VectorCost - ScalarCost;
     }
     case Instruction::FCmp:
     case Instruction::ICmp:
     case Instruction::Select:
     case Instruction::Add:
     case Instruction::FAdd:
     case Instruction::Sub:
     case Instruction::FSub:
for (ExternalUser &EU : ExternalUses) {
     // Uses by ephemeral values are free (because the ephemeral value will be
     // removed prior to code generation, and so the extraction will be
     // removed as well).
     if (EphValues.count(EU.User))
       continue;
 
     // If we plan to rewrite the tree in a smaller type, we will need to sign
     // extend the extracted value back to the original type. Here, we account
-    // for the extract and the added cost of the sign extend if needed.
+    // for the extract and the added cost of the sign extend if needed. We
+    // build a context for the extend that includes the extract to improve the
+    // accuracy of the estimate.
     auto *VecTy = VectorType::get(EU.Scalar->getType(), BundleWidth);
     auto *ScalarRoot = VectorizableTree[0].Scalars[0];
     if (MinBWs.count(ScalarRoot)) {
       auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot]);
       VecTy = VectorType::get(MinTy, BundleWidth);
-      ExtractCost +=
-          TTI->getCastInstrCost(Instruction::SExt, EU.Scalar->getType(), MinTy);
+      InstrCtx Vector(VecTy), Extract, Extend;
+      Extract = InstrCtx(MinTy, Instruction::ExtractElement, Vector, &Extend);
+      Extend = InstrCtx(EU.Scalar->getType(), Instruction::SExt, Extract);
+      ExtractCost += TTI->getCastInstrCost(
+          Instruction::SExt, EU.Scalar->getType(), MinTy, &Extend);
     }
     ExtractCost +=
         TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, EU.Lane);
   }
 
   int SpillCost = getSpillCost();
   Cost += SpillCost + ExtractCost;
 

test/Transforms/SLPVectorizer/AArch64/gather-reduce.ll

-; RUN: opt -S -slp-vectorizer -dce -instcombine < %s | FileCheck %s --check-prefix=PROFITABLE
-; RUN: opt -S -slp-vectorizer -slp-threshold=-12 -dce -instcombine < %s | FileCheck %s --check-prefix=UNPROFITABLE
+; RUN: opt -S -slp-vectorizer -dce -instcombine < %s | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-gnu"
 
 ; These tests check that we vectorize the index calculations in the
 ; gather-reduce pattern shown below. We check cases having i32 and i64
 ; subtraction.
 ;
 ; int gather_reduce_8x16(short *a, short *b, short *g, int n) {
 ;   int sum = 0;
 ;   for (int i = 0; i < n ; ++i) {
 ;     sum += g[*a++ - b[0]]; sum += g[*a++ - b[1]];
 ;     sum += g[*a++ - b[2]]; sum += g[*a++ - b[3]];
 ;     sum += g[*a++ - b[4]]; sum += g[*a++ - b[5]];
 ;     sum += g[*a++ - b[6]]; sum += g[*a++ - b[7]];
 ;   }
 ;   return sum;
 ; }
 
-; PROFITABLE-LABEL: @gather_reduce_8x16_i32
+; CHECK-LABEL: @gather_reduce_8x16_i32
 ;
-; PROFITABLE: [[L:%[a-zA-Z0-9.]+]] = load <8 x i16>
-; PROFITABLE: zext <8 x i16> [[L]] to <8 x i32>
-; PROFITABLE: [[S:%[a-zA-Z0-9.]+]] = sub nsw <8 x i32>
-; PROFITABLE: [[X:%[a-zA-Z0-9.]+]] = extractelement <8 x i32> [[S]]
-; PROFITABLE: sext i32 [[X]] to i64
+; CHECK: [[L:%[a-zA-Z0-9.]+]] = load <8 x i16>
+; CHECK: zext <8 x i16> [[L]] to <8 x i32>
+; CHECK: [[S:%[a-zA-Z0-9.]+]] = sub nsw <8 x i32>
+; CHECK: [[X:%[a-zA-Z0-9.]+]] = extractelement <8 x i32> [[S]]
+; CHECK: sext i32 [[X]] to i64
 ;
 define i32 @gather_reduce_8x16_i32(i16* nocapture readonly %a, i16* nocapture readonly %b, i16* nocapture readonly %g, i32 %n) {
 entry:
   %cmp.99 = icmp sgt i32 %n, 0
   br i1 %cmp.99, label %for.body.preheader, label %for.cond.cleanup
 
 for.body.preheader:
   br label %for.body
for.body:
   %23 = load i16, i16* %arrayidx64, align 2
   %conv65 = zext i16 %23 to i32
   %add66 = add nsw i32 %add57, %conv65
   %inc = add nuw nsw i32 %i.0103, 1
   %exitcond = icmp eq i32 %inc, %n
   br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
 }
 
-; UNPROFITABLE-LABEL: @gather_reduce_8x16_i64
+; CHECK-LABEL: @gather_reduce_8x16_i64
 ;
-; UNPROFITABLE: [[L:%[a-zA-Z0-9.]+]] = load <8 x i16>
-; UNPROFITABLE: zext <8 x i16> [[L]] to <8 x i32>
-; UNPROFITABLE: [[S:%[a-zA-Z0-9.]+]] = sub nsw <8 x i32>
-; UNPROFITABLE: [[X:%[a-zA-Z0-9.]+]] = extractelement <8 x i32> [[S]]
-; UNPROFITABLE: sext i32 [[X]] to i64
-;
-; TODO: Although we can now vectorize this case while converting the i64
-;       subtractions to i32, the cost model currently finds vectorization to be
-;       unprofitable. The cost model is penalizing the sign and zero
-;       extensions in the vectorized version, but they are actually free.
+; CHECK: [[L:%[a-zA-Z0-9.]+]] = load <8 x i16>
+; CHECK: zext <8 x i16> [[L]] to <8 x i32>
+; CHECK: [[S:%[a-zA-Z0-9.]+]] = sub nsw <8 x i32>
+; CHECK: [[X:%[a-zA-Z0-9.]+]] = extractelement <8 x i32> [[S]]
+; CHECK: sext i32 [[X]] to i64
 ;
 define i32 @gather_reduce_8x16_i64(i16* nocapture readonly %a, i16* nocapture readonly %b, i16* nocapture readonly %g, i32 %n) {
 entry:
   %cmp.99 = icmp sgt i32 %n, 0
   br i1 %cmp.99, label %for.body.preheader, label %for.cond.cleanup
 
 for.body.preheader:
   br label %for.body