[TTI] Change getOperandsScalarizationOverhead to take Type args

[dmgreen]

Summary

As a followup to D95291, getOperandsScalarizationOverhead was still using a VF as a vector factor if the arguments were scalar, and would assert on certain matrix intrinsics with differently sized vector arguments. This patch removes the VF arg, instead passing the Types through directly. This should allow it to more accurately compute the cost without having to guess at which operands will be vectorized, something difficult with more complex intrinsics.

[sdesmalen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

630–631

nit: if(auto *VecTy = dyn_cast<VectorType>(Ty))

647

Is this not doing something different than it did before, i.e. it now always uses InTy for each of the arguments to get the cost, where it previously used the types from each of the arguments?
I expected something similar to what you did in LoopVectorize with MaybeVectorizeType.

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

751

nit: clang-format? (over 80char I think)

[dmgreen]

Thanks for taking a look.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

647

Yeah. This version of getScalarizationOverhead is only called from getArithmeticInstrCost (or ARMTTIImpl::getArithmeticInstrCost etc). So the arguments/return value are all likely to all be the same value, and this is making that assumption. It is probably at least good enough as a heuristic, when compared to the else below. I've added a comment too.

[sdesmalen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

647

If this code assumes all arguments are of the same type, can you add an assert which checks this?

[dmgreen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

647

Unfortunately that isn't something that is always true. The type of Args may be incorrect, even the scalar type. The vectorizer can do bitwidth narrowing at the same time as vectorization, in which case the Ty will be the most accurate type available. If the original Arg type was an i32, but the vectorizer is vectorizing with it as an i8 to maximize the vector factor, a vXi8 is the correct type to use. That will be Ty for all the uses of this function.

I could try to move this array creation to the callers of the function. That would involve duplicating this bit of code in a few places, but after e771614bae0a05 removed the only other use of it, it should ensure these types are correct from the outside.

[dmgreen]

Pass Tys array to getScalarizationOverhead

[sdesmalen]

Thanks for the changes! I think moving the creation of the Type array to the callers is an improvement, it makes it more clear what types are used to calculate the cost.
Just a few more nits, and then I'm happy.

llvm/include/llvm/Analysis/TargetTransformInfo.h

726–727

Comment needs updating.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

638

nit: can you add a comment to describe the meaning of the arguments?

Just my mental understanding:

• InTy is the type of the vector(ized) node that needs scalarizing, hence it's always a VectorType.
• Args are the arguments to the node. These arguments are passed to see e.g. if any of them is a constant. If the question originates from the loopvectorizer, then the node and arguments are likely all vectorized with type InTy. Otherwise, they may have other types (such as the matrix intrinsic), but that information cannot be deduced from the arguments themselves yet, because the vectorization may not have happened yet.
• Tys are the actual types of each of the arguments after vectorizing, and are used to calculate the scalarization cost.

647

Thanks for the explanation and changes!

nit: This comment ("// Assume all arguments are of type Ty, ..") can now be removed.

[dmgreen]

Update comments.

[sdesmalen]

LGTM, thanks @dmgreen!

llvm/include/llvm/Analysis/TargetTransformInfo.h

728

nit: \p Tys

llvm/include/llvm/CodeGen/BasicTTIImpl.h

643

nit: \p RetTy, \p Args and \p Tys

[dmgreen]

Thanks

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
 
   /// Estimate the overhead of scalarizing an instruction. Insert and Extract
   /// are set if the demanded result elements need to be inserted and/or
   /// extracted from vectors.
   unsigned getScalarizationOverhead(VectorType *Ty, const APInt &DemandedElts,
                                     bool Insert, bool Extract) const;
 
   /// Estimate the overhead of scalarizing an instructions unique
   /// non-constant operands. The types of the arguments are ordinarily
   /// scalar, in which case the costs are multiplied with VF.
   unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                            unsigned VF) const;
+                                            ArrayRef<Type *> Tys) const;
 
   /// If target has efficient vector element load/store instructions, it can
   /// return true here so that insertion/extraction costs are not added to
   /// the scalarization cost of a load/store.
   bool supportsEfficientVectorElementLoadStore() const;
 
   /// Don't restrict interleaved unrolling to small loops.
   bool enableAggressiveInterleaving(bool LoopHasReductions) const;
public:
   virtual bool shouldBuildLookupTables() = 0;
   virtual bool shouldBuildLookupTablesForConstant(Constant *C) = 0;
   virtual bool useColdCCForColdCall(Function &F) = 0;
   virtual unsigned getScalarizationOverhead(VectorType *Ty,
                                             const APInt &DemandedElts,
                                             bool Insert, bool Extract) = 0;
   virtual unsigned
   getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                   unsigned VF) = 0;
+                                   ArrayRef<Type *> Tys) = 0;
   virtual bool supportsEfficientVectorElementLoadStore() = 0;
   virtual bool enableAggressiveInterleaving(bool LoopHasReductions) = 0;
   virtual MemCmpExpansionOptions
   enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const = 0;
   virtual bool enableInterleavedAccessVectorization() = 0;
   virtual bool enableMaskedInterleavedAccessVectorization() = 0;
   virtual bool isFPVectorizationPotentiallyUnsafe() = 0;
   virtual bool allowsMisalignedMemoryAccesses(LLVMContext &Context,
bool useColdCCForColdCall(Function &F) override {
     return Impl.useColdCCForColdCall(F);
   }
 
   unsigned getScalarizationOverhead(VectorType *Ty, const APInt &DemandedElts,
                                     bool Insert, bool Extract) override {
     return Impl.getScalarizationOverhead(Ty, DemandedElts, Insert, Extract);
   }
   unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                            unsigned VF) override {
-    return Impl.getOperandsScalarizationOverhead(Args, VF);
+                                            ArrayRef<Type *> Tys) override {
+    return Impl.getOperandsScalarizationOverhead(Args, Tys);
   }
 
   bool supportsEfficientVectorElementLoadStore() override {
     return Impl.supportsEfficientVectorElementLoadStore();
   }
 
   bool enableAggressiveInterleaving(bool LoopHasReductions) override {
     return Impl.enableAggressiveInterleaving(LoopHasReductions);

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool useColdCCForColdCall(Function &F) const { return false; }
 
   unsigned getScalarizationOverhead(VectorType *Ty, const APInt &DemandedElts,
                                     bool Insert, bool Extract) const {
     return 0;
   }
 
   unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                            unsigned VF) const {
+                                            ArrayRef<Type *> Tys) const {
     return 0;
   }
 
   bool supportsEfficientVectorElementLoadStore() const { return false; }
 
   bool enableAggressiveInterleaving(bool LoopHasReductions) const {
     return false;
   }

llvm/include/llvm/CodeGen/BasicTTIImpl.h

unsigned getScalarizationOverhead(VectorType *InTy, bool Insert,
     APInt DemandedElts = APInt::getAllOnesValue(Ty->getNumElements());
     return thisT()->getScalarizationOverhead(Ty, DemandedElts, Insert, Extract);
   }
 
   /// Estimate the overhead of scalarizing an instruction's unique
   /// non-constant operands. The types of the arguments are ordinarily
   /// scalar, in which case the costs are multiplied with VF.
   unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                            unsigned VF) {
+                                            ArrayRef<Type *> Tys) {
+    assert(Args.size() == Tys.size() && "Expected matching Args and Tys");
+
     unsigned Cost = 0;
     SmallPtrSet<const Value*, 4> UniqueOperands;
-    for (const Value *A : Args) {
+    for (int I = 0, E = Args.size(); I != E; I++) {
       // Disregard things like metadata arguments.
-      Type *Ty = A->getType();
+      const Value *A = Args[I];
+      Type *Ty = Tys[I];
       if (!Ty->isIntOrIntVectorTy() && !Ty->isFPOrFPVectorTy() &&
           !Ty->isPtrOrPtrVectorTy())
         continue;
 
       if (!isa<Constant>(A) && UniqueOperands.insert(A).second) {
         auto *VecTy = dyn_cast<VectorType>(Ty);
-        if (VecTy) {
-          // If A is a vector operand, VF should be 1 or correspond to A.
-          assert((VF == 1 ||
-                  VF == cast<FixedVectorType>(VecTy)->getNumElements()) &&
-                 "Vector argument does not match VF");
-        }
-        else
-          VecTy = FixedVectorType::get(Ty, VF);
-
+        if (VecTy)
           Cost += getScalarizationOverhead(VecTy, false, true);
       }
     }
 
     return Cost;
   }
 
   unsigned getScalarizationOverhead(VectorType *InTy,
                                     ArrayRef<const Value *> Args) {
     auto *Ty = cast<FixedVectorType>(InTy);
 
     unsigned Cost = 0;
 
     Cost += getScalarizationOverhead(Ty, true, false);
-    if (!Args.empty())
-      Cost += getOperandsScalarizationOverhead(Args, Ty->getNumElements());
-    else
+    if (!Args.empty()) {
+      SmallVector<Type *> Tys(Args.size(), Ty);
+      Cost += getOperandsScalarizationOverhead(Args, Tys);
+    } else
       // When no information on arguments is provided, we add the cost
       // associated with one argument as a heuristic.
       Cost += getScalarizationOverhead(Ty, false, true);
 
     return Cost;
   }
 
   unsigned getMaxInterleaveFactor(unsigned VF) { return 1; }
unsigned getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
     // intrinsic.
     unsigned ScalarizationCost = std::numeric_limits<unsigned>::max();
     if (RetVF.isVector()) {
       ScalarizationCost = 0;
       if (!RetTy->isVoidTy())
         ScalarizationCost +=
             getScalarizationOverhead(cast<VectorType>(RetTy), true, false);
       ScalarizationCost +=
-          getOperandsScalarizationOverhead(Args, RetVF.getKnownMinValue());
+          getOperandsScalarizationOverhead(Args, ICA.getArgTypes());
     }
 
     IntrinsicCostAttributes Attrs(IID, RetTy, ICA.getArgTypes(), FMF, I,
                                   ScalarizationCost);
     return thisT()->getTypeBasedIntrinsicInstrCost(Attrs, CostKind);
   }
 
   /// Get intrinsic cost based on argument types.

llvm/lib/Analysis/TargetTransformInfo.cpp

 unsigned
 TargetTransformInfo::getScalarizationOverhead(VectorType *Ty,
                                               const APInt &DemandedElts,
                                               bool Insert, bool Extract) const {
   return TTIImpl->getScalarizationOverhead(Ty, DemandedElts, Insert, Extract);
 }
 
 unsigned TargetTransformInfo::getOperandsScalarizationOverhead(
-    ArrayRef<const Value *> Args, unsigned VF) const {
-  return TTIImpl->getOperandsScalarizationOverhead(Args, VF);
+    ArrayRef<const Value *> Args, ArrayRef<Type *> Tys) const {
+  return TTIImpl->getOperandsScalarizationOverhead(Args, Tys);
 }
 
 bool TargetTransformInfo::supportsEfficientVectorElementLoadStore() const {
   return TTIImpl->supportsEfficientVectorElementLoadStore();
 }
 
 bool TargetTransformInfo::enableAggressiveInterleaving(
     bool LoopHasReductions) const {

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

if (!OrigTy.isSimple()) {
     // intrinsic. A vectorizer will pass a scalar RetTy and VF > 1, while
     // CostModel will pass a vector RetTy and VF is 1.
     unsigned ScalarizationCost = std::numeric_limits<unsigned>::max();
     if (RetVF > 1) {
       ScalarizationCost = 0;
       if (!RetTy->isVoidTy())
         ScalarizationCost +=
             getScalarizationOverhead(cast<VectorType>(RetTy), true, false);
-      ScalarizationCost += getOperandsScalarizationOverhead(Args, RetVF);
+      ScalarizationCost += getOperandsScalarizationOverhead(Args, ICA.getArgTypes());
     }
 
     IntrinsicCostAttributes Attrs(ICA.getID(), RetTy, ICA.getArgTypes(), FMF, I,
                                   ScalarizationCost);
     return getIntrinsicInstrCost(Attrs, CostKind);
   }
 
   // Legalize the type.

llvm/lib/Target/Hexagon/HexagonTargetTransformInfo.h

public:
   }
   bool enableInterleavedAccessVectorization() {
     return true;
   }
 
   unsigned getScalarizationOverhead(VectorType *Ty, const APInt &DemandedElts,
                                     bool Insert, bool Extract);
   unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
-                                            unsigned VF);
+                                            ArrayRef<Type *> Tys);
   unsigned getCallInstrCost(Function *F, Type *RetTy, ArrayRef<Type*> Tys,
                             TTI::TargetCostKind CostKind);
   unsigned getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
                                  TTI::TargetCostKind CostKind);
   unsigned getAddressComputationCost(Type *Tp, ScalarEvolution *SE,
             const SCEV *S);
   unsigned getMemoryOpCost(unsigned Opcode, Type *Src, MaybeAlign Alignment,
                            unsigned AddressSpace,

llvm/lib/Target/Hexagon/HexagonTargetTransformInfo.cpp

 
 unsigned HexagonTTIImpl::getScalarizationOverhead(VectorType *Ty,
                                                   const APInt &DemandedElts,
                                                   bool Insert, bool Extract) {
   return BaseT::getScalarizationOverhead(Ty, DemandedElts, Insert, Extract);
 }
 
 unsigned HexagonTTIImpl::getOperandsScalarizationOverhead(
-      ArrayRef<const Value*> Args, unsigned VF) {
-  return BaseT::getOperandsScalarizationOverhead(Args, VF);
+      ArrayRef<const Value*> Args, ArrayRef<Type *> Tys) {
+  return BaseT::getOperandsScalarizationOverhead(Args, Tys);
 }
 
 unsigned HexagonTTIImpl::getCallInstrCost(Function *F, Type *RetTy,
       ArrayRef<Type*> Tys, TTI::TargetCostKind CostKind) {
   return BaseT::getCallInstrCost(F, RetTy, Tys, CostKind);
 }
 
 unsigned

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

InstructionCost VectorCallCost =
       TTI.getCallInstrCost(nullptr, RetTy, Tys, TTI::TCK_RecipThroughput);
   if (VectorCallCost < Cost) {
     NeedToScalarize = false;
     Cost = VectorCallCost;
   }
   return Cost;
 }
 
-InstructionCost
-LoopVectorizationCostModel::getVectorIntrinsicCost(CallInst *CI,
-                                                   ElementCount VF) {
-  auto MaybeVectorizeType = [](Type *Elt, ElementCount VF) -> Type * {
-    if (VF.isScalar() || (!Elt->isIntOrPtrTy() && !Elt->isFloatingPointTy()))
-      return Elt;
-    return VectorType::get(Elt, VF);
-  };
+static Type *MaybeVectorizeType(Type *Elt, ElementCount VF) {
+}
 
+InstructionCost
+LoopVectorizationCostModel::getVectorIntrinsicCost(CallInst *CI,
+                                                   ElementCount VF) {
   Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
   assert(ID && "Expected intrinsic call!");
   Type *RetTy = MaybeVectorizeType(CI->getType(), VF);
   FastMathFlags FMF;
   if (auto *FPMO = dyn_cast<FPMathOperator>(CI))
     FMF = FPMO->getFastMathFlags();
 
   SmallVector<const Value *> Arguments(CI->arg_begin(), CI->arg_end());
   FunctionType *FTy = CI->getCalledFunction()->getFunctionType();
   SmallVector<Type *> ParamTys;
-  std::transform(FTy->param_begin(), FTy->param_end(), ParamTys.begin(),
+  std::transform(FTy->param_begin(), FTy->param_end(),
+                 std::back_inserter(ParamTys),
                  [&](Type *Ty) { return MaybeVectorizeType(Ty, VF); });
 
   IntrinsicCostAttributes CostAttrs(ID, RetTy, Arguments, ParamTys, FMF,
                                     dyn_cast<IntrinsicInst>(CI));
   return TTI.getIntrinsicInstrCost(CostAttrs,
                                    TargetTransformInfo::TCK_RecipThroughput);
 }
 
if (isa<StoreInst>(I) && TTI.supportsEfficientVectorElementLoadStore())
     return Cost;
 
   // Collect operands to consider.
   CallInst *CI = dyn_cast<CallInst>(I);
   Instruction::op_range Ops = CI ? CI->arg_operands() : I->operands();
 
   // Skip operands that do not require extraction/scalarization and do not incur
   // any overhead.
+  SmallVector<Type*> Tys;
+  for (auto *V : filterExtractingOperands(Ops, VF))
+    Tys.push_back(MaybeVectorizeType(V->getType(), VF));
   return Cost + TTI.getOperandsScalarizationOverhead(
-                    filterExtractingOperands(Ops, VF), VF.getKnownMinValue());
+                    filterExtractingOperands(Ops, VF), Tys);
 }
 
 void LoopVectorizationCostModel::setCostBasedWideningDecision(ElementCount VF) {
   if (VF.isScalar())
     return;
   NumPredStores = 0;
   for (BasicBlock *BB : TheLoop->blocks()) {
     // For each instruction in the old loop.

llvm/test/Analysis/CostModel/PowerPC/matrix.ll

+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
+; RUN: opt < %s -cost-model -analyze -mtriple=powerpc64-unknown-linux-gnu -mcpu=pwr7 -mattr=+vsx | FileCheck %s
+target datalayout = "E-m:e-i64:64-n32:64"
+target triple = "powerpc64-unknown-linux-gnu"
+
+; This test checks we don't crash on certain matrix operations, more than
+; checks the cost of the intrinsics per-se.
+
+define void @matrix() {
+; CHECK-LABEL: 'matrix'
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %matrix1 = call <1 x i32> @llvm.matrix.column.major.load.v1i32(i32* nonnull align 4 undef, i64 1, i1 false, i32 1, i32 1)
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 452 for instruction: %0 = call <10 x i32> @llvm.matrix.multiply.v10i32.v10i32.v1i32(<10 x i32> undef, <1 x i32> %matrix1, i32 10, i32 1, i32 1)
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: ret void
+;
+entry:
+  %matrix1 = call <1 x i32> @llvm.matrix.column.major.load.v1i32(i32* nonnull align 4 undef, i64 1, i1 false, i32 1, i32 1)
+  %0 = call <10 x i32> @llvm.matrix.multiply.v10i32.v10i32.v1i32(<10 x i32> undef, <1 x i32> %matrix1, i32 10, i32 1, i32 1)
+  ret void
+}
+
+declare <1 x i32> @llvm.matrix.column.major.load.v1i32(i32* nocapture, i64, i1 immarg, i32 immarg, i32 immarg) #2
+declare <10 x i32> @llvm.matrix.multiply.v10i32.v10i32.v1i32(<10 x i32>, <1 x i32>, i32 immarg, i32 immarg, i32 immarg) #3