[LoopVectorizer] NFCI: Calculate register usage based on TLI.getTypeLegalizationCost.

[sdesmalen]

Summary

This is more accurate than dividing the bitwidth based on the element count by the
maximum register size, as it can just reuse whatever has been calculated for
legalization of these types.

This change is also necessary when calculating register usage for scalable vectors, where
the legalization of these types cannot be done based on the widest register size, because
that does not take the 'vscale' component into account.

[SjoerdMeijer]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

301

If it's only about this, perhaps better not to create yet another TTI hook? Perhaps just query getTypeLegalizationCost directly in the LV?

[sdesmalen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

301

I wasn't sure if this was by design, but TLI is currently not a member of the LoopVectorizer. I'm happy to add it, but I thought it had to go through TTI so that target lowering is abstracted from the more high-level loop-vectorization pass.

[SjoerdMeijer]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

301

Ah sorry, looking at it closer, I think you're right, calling this from the LV would be a layering violation....

[SjoerdMeijer]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

262

Perhaps 0 is not a very good default?

[sdesmalen]

Changed default for getRegUsageForType to return 1 instead of 0.

[sdesmalen]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

262

Good point!

[SjoerdMeijer]

Looks like a good change to me.

[sdesmalen]

Thanks @SjoerdMeijer!

[uabelho]

Hi,

I hit an assertion with this patch:
opt: ../lib/IR/Type.cpp:620: static llvm::FixedVectorType *llvm::FixedVectorType::get(llvm::Type *, unsigned int): Assertion `isValidElementType(ElementType) && "Element type of a VectorType must " "be an integer, floating point, or " "pointer type."' failed.

Reproduce with
opt -loop-vectorize -mtriple x86_64 -o /dev/null extractvalue-crash.ll

with the input

; RUN: opt -loop-vectorize -mtriple x86_64 -o /dev/null %s
%rec6 = type { i16 }

; Don't crash on trying to vectorize an extractvalue.

@b = global %rec6 zeroinitializer

define void @f1() {
  %_tmp2 = icmp ne i16 0, 0
  br i1 %_tmp2, label %bb1.preheader, label %bb3

bb1.preheader:                                    ; preds = %0
  %_tmp4 = load %rec6, %rec6* @b
  br label %bb1

bb1:                                              ; preds = %bb1.preheader, %bb1
  %_tmp81 = phi i16 [ 0, %bb1.preheader ], [ %_tmp7, %bb1 ]
  %_tmp4.fca.0.extract = extractvalue %rec6 %_tmp4, 0
  %_tmp7 = sub i16 %_tmp81, 1
  %_tmp9 = icmp ne i16 %_tmp7, 0
  br i1 %_tmp9, label %bb1, label %bb3.loopexit

bb3.loopexit:                                     ; preds = %bb1
  %_tmp7.lcssa = phi i16 [ %_tmp7, %bb1 ]
  br label %bb3

bb3:                                              ; preds = %bb3.loopexit, %0
  ret void
}

[sdesmalen]

In D91059#2388573, @uabelho wrote:

Hi,

I hit an assertion with this patch:
opt: ../lib/IR/Type.cpp:620: static llvm::FixedVectorType *llvm::FixedVectorType::get(llvm::Type *, unsigned int): Assertion `isValidElementType(ElementType) && "Element type of a VectorType must " "be an integer, floating point, or " "pointer type."' failed.

Reproduce with
opt -loop-vectorize -mtriple x86_64 -o /dev/null extractvalue-crash.ll

with the input

; RUN: opt -loop-vectorize -mtriple x86_64 -o /dev/null %s
%rec6 = type { i16 }

; Don't crash on trying to vectorize an extractvalue.

@b = global %rec6 zeroinitializer

define void @f1() {
  %_tmp2 = icmp ne i16 0, 0
  br i1 %_tmp2, label %bb1.preheader, label %bb3

bb1.preheader:                                    ; preds = %0
  %_tmp4 = load %rec6, %rec6* @b
  br label %bb1

bb1:                                              ; preds = %bb1.preheader, %bb1
  %_tmp81 = phi i16 [ 0, %bb1.preheader ], [ %_tmp7, %bb1 ]
  %_tmp4.fca.0.extract = extractvalue %rec6 %_tmp4, 0
  %_tmp7 = sub i16 %_tmp81, 1
  %_tmp9 = icmp ne i16 %_tmp7, 0
  br i1 %_tmp9, label %bb1, label %bb3.loopexit

bb3.loopexit:                                     ; preds = %bb1
  %_tmp7.lcssa = phi i16 [ %_tmp7, %bb1 ]
  br label %bb3

bb3:                                              ; preds = %bb3.loopexit, %0
  ret void
}

Thanks for reporting @uabelho, I'll revert the patch for now and fix the issue!

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5805

The example provided by @uabelho fails for the load instruction (with a struct as result type) which is outside the loop body and isn't actually vectorized, so counting it as a vector register usage seems overly conservative. The extractvalue instruction in the body itself is vectorized, which is fine, since the extracted value from the aggregate is a valid element type for a vector.

I suggest addressing the issue with:

-  if (Ty->isTokenTy())
+  if (Ty->isTokenTy() || !VectorType::isValidElementType(Ty))
      return 0U;

@uabelho @SjoerdMeijer are you happy with me re-landing the patch with the above change?

[uabelho]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5805

I've verified that the above change fixes the problem and in the limited testing I've done of it I haven't seen any new problems. Thanks!

@SjoerdMeijer ?

[SjoerdMeijer]

yeah, LGTM too

[RKSimon]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5795

@sdesmalen I'm seeing unused variable (written but not used) warnings for MaxSafeDepDist - can this be removed?

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5795

Thanks for letting me know! I've addressed it in rG41c9f4c1cea7314c0a6a32f2e59dab60cd575c1f

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// then/else to before if.
   bool isProfitableToHoist(Instruction *I) const;
 
   bool useAA() const;
 
   /// Return true if this type is legal.
   bool isTypeLegal(Type *Ty) const;
 
+  /// Returns the estimated number of registers required to represent \p Ty.
+  unsigned getRegUsageForType(Type *Ty) const;
+
   /// Return true if switches should be turned into lookup tables for the
   /// target.
   bool shouldBuildLookupTables() const;
 
   /// Return true if switches should be turned into lookup tables
   /// containing this constant value for the target.
   bool shouldBuildLookupTablesForConstant(Constant *C) const;
 
public:
   virtual int getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                    int64_t BaseOffset, bool HasBaseReg,
                                    int64_t Scale, unsigned AddrSpace) = 0;
   virtual bool LSRWithInstrQueries() = 0;
   virtual bool isTruncateFree(Type *Ty1, Type *Ty2) = 0;
   virtual bool isProfitableToHoist(Instruction *I) = 0;
   virtual bool useAA() = 0;
   virtual bool isTypeLegal(Type *Ty) = 0;
+  virtual unsigned getRegUsageForType(Type *Ty) = 0;
   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,
public:
   bool isTruncateFree(Type *Ty1, Type *Ty2) override {
     return Impl.isTruncateFree(Ty1, Ty2);
   }
   bool isProfitableToHoist(Instruction *I) override {
     return Impl.isProfitableToHoist(I);
   }
   bool useAA() override { return Impl.useAA(); }
   bool isTypeLegal(Type *Ty) override { return Impl.isTypeLegal(Ty); }
+  unsigned getRegUsageForType(Type *Ty) override {
+    return Impl.getRegUsageForType(Ty);
+  }
   bool shouldBuildLookupTables() override {
     return Impl.shouldBuildLookupTables();
   }
   bool shouldBuildLookupTablesForConstant(Constant *C) override {
     return Impl.shouldBuildLookupTablesForConstant(C);
   }
   bool useColdCCForColdCall(Function &F) override {
     return Impl.useColdCCForColdCall(F);

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool isTruncateFree(Type *Ty1, Type *Ty2) { return false; }
 
   bool isProfitableToHoist(Instruction *I) { return true; }
 
   bool useAA() { return false; }
 
   bool isTypeLegal(Type *Ty) { return false; }
 
+  unsigned getRegUsageForType(Type *Ty) { return 1; }
+
   bool shouldBuildLookupTables() { return true; }
   bool shouldBuildLookupTablesForConstant(Constant *C) { return true; }
 
   bool useColdCCForColdCall(Function &F) { return false; }
 
   unsigned getScalarizationOverhead(VectorType *Ty, const APInt &DemandedElts,
                                     bool Insert, bool Extract) {
     return 0;

llvm/include/llvm/CodeGen/BasicTTIImpl.h

public:
 
   bool useAA() const { return getST()->useAA(); }
 
   bool isTypeLegal(Type *Ty) {
     EVT VT = getTLI()->getValueType(DL, Ty);
     return getTLI()->isTypeLegal(VT);
   }
 
+  unsigned getRegUsageForType(Type *Ty) {
+    return getTLI()->getTypeLegalizationCost(DL, Ty).first;
+  }
+
   int getGEPCost(Type *PointeeType, const Value *Ptr,
                  ArrayRef<const Value *> Operands) {
     return BaseT::getGEPCost(PointeeType, Ptr, Operands);
   }
 
   unsigned getEstimatedNumberOfCaseClusters(const SwitchInst &SI,
                                             unsigned &JumpTableSize,
                                             ProfileSummaryInfo *PSI,

llvm/lib/Analysis/TargetTransformInfo.cpp

 }
 
 bool TargetTransformInfo::useAA() const { return TTIImpl->useAA(); }
 
 bool TargetTransformInfo::isTypeLegal(Type *Ty) const {
   return TTIImpl->isTypeLegal(Ty);
 }
 
+unsigned TargetTransformInfo::getRegUsageForType(Type *Ty) const {
+  return TTIImpl->getRegUsageForType(Ty);
+}
+
 bool TargetTransformInfo::shouldBuildLookupTables() const {
   return TTIImpl->shouldBuildLookupTables();
 }
 bool TargetTransformInfo::shouldBuildLookupTablesForConstant(
     Constant *C) const {
   return TTIImpl->shouldBuildLookupTablesForConstant(C);
 }
 

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

LoopVectorizationCostModel::calculateRegisterUsage(ArrayRef<ElementCount> VFs) {
   for (auto &Interval : EndPoint)
     TransposeEnds[Interval.second].push_back(Interval.first);
 
   SmallPtrSet<Instruction *, 8> OpenIntervals;
 
   // Get the size of the widest register.
   unsigned MaxSafeDepDist = -1U;
   if (Legal->getMaxSafeDepDistBytes() != -1U)
     MaxSafeDepDist = Legal->getMaxSafeDepDistBytes() * 8;
-  unsigned WidestRegister =
-      std::min(TTI.getRegisterBitWidth(true), MaxSafeDepDist);
   const DataLayout &DL = TheFunction->getParent()->getDataLayout();
 
   SmallVector<RegisterUsage, 8> RUs(VFs.size());
   SmallVector<SmallMapVector<unsigned, unsigned, 4>, 8> MaxUsages(VFs.size());
 
   LLVM_DEBUG(dbgs() << "LV(REG): Calculating max register usage:\n");
 
   // A lambda that gets the register usage for the given type and VF.
-  auto GetRegUsage = [&DL, WidestRegister](Type *Ty, ElementCount VF) {
+  auto GetRegUsage = [&DL, &TTI=TTI](Type *Ty, ElementCount VF) {
     if (Ty->isTokenTy())
       return 0U;
-    unsigned TypeSize = DL.getTypeSizeInBits(Ty->getScalarType());
-    assert(!VF.isScalable() && "scalable vectors not yet supported.");
-    return std::max<unsigned>(1, VF.getKnownMinValue() * TypeSize /
-                                     WidestRegister);
+    return TTI.getRegUsageForType(VectorType::get(Ty, VF));
   };
 
   for (unsigned int i = 0, s = IdxToInstr.size(); i < s; ++i) {
     Instruction *I = IdxToInstr[i];
 
     // Remove all of the instructions that end at this location.
     InstrList &List = TransposeEnds[i];
     for (Instruction *ToRemove : List)