[llvm][Codegen] Make `getVectorTypeBreakdownMVT` work with scalable types.

[efriedma]

llvm/lib/CodeGen/TargetLoweringBase.cpp

954

I'd prefer just making people write out isPowerOf2_32(EC.Min). This sort of operation should be rare, and I'd prefer to spell out what it actually means.

974

I'm not sure how you expect this assertion to work out. For the current set of defined MVTs, it might work for SVE, but this will explode as soon as someone adds MVTs for a new target, or tries to add support for a target that doesn't have f16 vectors, or we query this on a target that doesn't support scalable vectors.

There's also the larger problem of how legalization should work if we run into a scalable vector with an illegal element type. Not sure what the right answer is there.

986

At this point, the element count is known to be a power of two; maybe it would be simpler to work with the element types here? Those are never scalable, so it's more obvious that NextPowerOf2 does something sane.

[fpetrogalli]

Thank you for the review @efriedma.

I have updated the patch accordingly.

Francesco

[efriedma]

LGTM

llvm/include/llvm/Support/TypeSize.h

public:
 
   ElementCount(unsigned Min, bool Scalable)
   : Min(Min), Scalable(Scalable) {}
 
   ElementCount operator*(unsigned RHS) {
     return { Min * RHS, Scalable };
   }
   ElementCount operator/(unsigned RHS) {
+    assert(Min % RHS == 0 && "Min is not a multiple of RHS.");
     return { Min / RHS, Scalable };
   }
+  ElementCount operator>>=(unsigned RHS) { return {Min >> RHS, Scalable}; }
 
   bool operator==(const ElementCount& RHS) const {
     return Min == RHS.Min && Scalable == RHS.Scalable;
   }
   bool operator!=(const ElementCount& RHS) const {
     return !(*this == RHS);
   }
+  /// Checks whether Min is a power of 2. Note that for a ElementCount
+  /// instance where Scalable is set to true, this is not generally
+  /// true from a mathematical point of view, as the Scalable property
+  /// in this case is saying that the ElementCount is holding a value
+  /// that is and unknown multiple of MinSize. For example,
+  /// ElementCount(4, true).isPowerOf2() == true, but its symbolic
+  /// value is 4 x UNKNOWN.
+  bool isPowerOf2() { return isPowerOf2_32(Min); }
 };
 
 // This class is used to represent the size of types. If the type is of fixed
 // size, it will represent the exact size. If the type is a scalable vector,
 // it will represent the known minimum size.
 class TypeSize {
   uint64_t MinSize;   // The known minimum size.
   bool IsScalable;    // If true, then the runtime size is an integer multiple
public:
   friend TypeSize operator*(const unsigned LHS, const TypeSize &RHS) {
     return { LHS * RHS.MinSize, RHS.IsScalable };
   }
 
   TypeSize operator/(unsigned RHS) const {
     return { MinSize / RHS, IsScalable };
   }
 
+  friend TypeSize operator/(const TypeSize &LHS, const TypeSize &RHS) {
+    assert(LHS.IsScalable == RHS.IsScalable &&
+           "Division of scalable and fixed types");
+
+    return {LHS.MinSize / RHS.MinSize, RHS.IsScalable};
+  }
   // Return the minimum size with the assumption that the size is exact.
   // Use in places where a scalable size doesn't make sense (e.g. non-vector
   // types, or vectors in backends which don't support scalable vectors).
   uint64_t getFixedSize() const {
     assert(!IsScalable && "Request for a fixed size on a scalable object");
     return MinSize;
   }
 
public:
   }
 
   // Returns true if the type size is non-zero.
   bool isNonZero() const { return MinSize != 0; }
 
   // Returns true if the type size is zero.
   bool isZero() const { return MinSize == 0; }
 
+  /// Checks whether MinSize is a power of 2. Note that for a TypeSize
+  /// instance where IsScalable is set to true, this is not generally
+  /// true from a mathematical point of view, as the IsScalable
+  /// property in this case is saying that the TypeSize is holding a
+  /// value that is and unknown multiple of MinSize. For example,
+  /// TypeSize(4, true).isPowerOf2() == true, but its symbolic value
+  /// is 4 x UNKNOWN.
+  bool isPowerOf2() const { return isPowerOf2_32(MinSize); }
+
+  /// Returns a TypeSize where MinSize is set to be the smallest power
+  /// of 2 greater than the MinSize of the current instance. The
+  /// IsScalable fiels is preserved.
+  static TypeSize getNextPowerOf2(TypeSize &TS) {
+    return {NextPowerOf2(TS.MinSize), TS.IsScalable};
+  }
+
   // Casts to a uint64_t if this is a fixed-width size.
   //
   // This interface is deprecated and will be removed in a future version
   // of LLVM in favour of upgrading uses that rely on this implicit conversion
   // to uint64_t. Calls to functions that return a TypeSize should use the
   // proper interfaces to TypeSize.
   // In practice this is mostly calls to MVT/EVT::getSizeInBits().
   //

llvm/lib/CodeGen/TargetLoweringBase.cpp

TargetLoweringBase::getTypeConversion(LLVMContext &Context, EVT VT) const {
   return LegalizeKind(TypeSplitVector, NVT);
 }
 
 static unsigned getVectorTypeBreakdownMVT(MVT VT, MVT &IntermediateVT,
                                           unsigned &NumIntermediates,
                                           MVT &RegisterVT,
                                           TargetLoweringBase *TLI) {
   // Figure out the right, legal destination reg to copy into.
-  unsigned NumElts = VT.getVectorNumElements();
+  ElementCount EC = VT.getVectorElementCount();
   MVT EltTy = VT.getVectorElementType();
 
   unsigned NumVectorRegs = 1;
 
-  // FIXME: We don't support non-power-of-2-sized vectors for now.  Ideally we
-  // could break down into LHS/RHS like LegalizeDAG does.
-  if (!isPowerOf2_32(NumElts)) {
-    NumVectorRegs = NumElts;
-    NumElts = 1;
-  }
-
-  // Divide the input until we get to a supported size.  This will always
-  // end with a scalar if the target doesn't support vectors.
-  while (NumElts > 1 && !TLI->isTypeLegal(MVT::getVectorVT(EltTy, NumElts))) {
-    NumElts >>= 1;
+  // FIXME: We don't support non-power-of-2-sized vectors for now.
+  // Ideally we could break down into LHS/RHS like LegalizeDAG does.
+  if (!EC.isPowerOf2()) {
+    // Split EC to unit size (scalable property is preserved).
+    NumVectorRegs = EC.Min;
+    EC = EC / NumVectorRegs;
+  }
+
+  // Divide the input until we get to a supported size. For example,
+  // on a target where MVT::nxv4i32 is a legal type, MVT::nxv32i32 is
+  // split in 8 vectors of type MVT::nxv4i32. On a target where
+  // MVT::v4i16 is a legal type, MVT::v64i16 is split in 16 vectors of
+  // type MVT::v4i16. This loop will always end up with an EC that
+  // represent a scalar or a scalable scalar. For scalable EC, to the
+  // best of our knowledge, we do not expect to ever get to EC = {1,
+  // true} for any of the LLVM Targets. The assertion following the
+  // loop capture such situation.
+  while (EC.Min > 1 && !TLI->isTypeLegal(MVT::getVectorVT(EltTy, EC))) {
+    EC >>= 1;
     NumVectorRegs <<= 1;
   }
+  assert(!(EC.Scalable && EC.Min == 1) && "The split has resulted in a"
+                                          "scalable vector with one lane.");
 
   NumIntermediates = NumVectorRegs;
 
-  MVT NewVT = MVT::getVectorVT(EltTy, NumElts);
+  MVT NewVT = MVT::getVectorVT(EltTy, EC);
   if (!TLI->isTypeLegal(NewVT))
     NewVT = EltTy;
   IntermediateVT = NewVT;
 
-  unsigned NewVTSize = NewVT.getSizeInBits();
+  TypeSize NewVTSize = NewVT.getSizeInBits();
 
   // Convert sizes such as i33 to i64.
-  if (!isPowerOf2_32(NewVTSize))
-    NewVTSize = NextPowerOf2(NewVTSize);
+  if (!NewVTSize.isPowerOf2())
+    NewVTSize = TypeSize::getNextPowerOf2(NewVTSize);
 
   MVT DestVT = TLI->getRegisterType(NewVT);
   RegisterVT = DestVT;
   if (EVT(DestVT).bitsLT(NewVT))    // Value is expanded, e.g. i64 -> i16.
     return NumVectorRegs*(NewVTSize/DestVT.getSizeInBits());
 
   // Otherwise, promotion or legal types use the same number of registers as
   // the vector decimated to the appropriate level.